<commit_before>package goxmpp

import (
	"encoding/xml"
	"io"
)

type Stream struct {
	XMLName    xml.Name `xml:"http://etherx.jabber.org/streams stream"`
	ID         string   `xml:"id,attr,omitempty"`
	From       string   `xml:"from,attr,omitempty"`
	To         string   `xml:"to,attr,omitempty"`
	Version    string   `xml:"version,attr,omitempty"`
}

type Mechanisms struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl mechanisms"`
	Names   []string `xml:"mechanism"`
}

type DigestMD5Auth struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl auth"`
	ID      string   `xml:"id,attr"`
}

type StartTLS struct {
	XMLName  xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-tls starttls"`
	Required bool     `xml:"required,omitempty"`
}

type Features struct {
	// http://tools.ietf.org/html/rfc3920#section-4.6
	XMLName          xml.Name    `xml:"stream:features"`
	Mechanisms       *Mechanisms
	StartTLS         *StartTLS
	Bind             bool        `xml:"urn:ietf:params:xml:ns:xmpp-bind bind,omitempty"`
	Session          bool        `xml:"urn:ietf:params:xml:ns:xmpp-session session,omitempty"`
}

type Error struct {
	Type   string `xml:"type,attr,omitempty"`
}

type Stanza struct {
	From   string `xml:"from,attr,omitempty"`
	To     string `xml:"to,attr,omitempty"`
	Type   string `xml:"type,attr,omitempty"`
	ID     string `xml:"id,attr,omitempty"`
	Error  Error
}

type Message struct {
	XMLName xml.Name `xml:"message"`
	Stanza
	Body    string   `xml:"body,omitempty"`
}

type VersionQuery struct {
	// http://xmpp.org/extensions/xep-0092.html
	XMLName xml.Name `xml:"jabber:iq:version query"`
	Name    string   `xml:"name,attr,omitempty"`
	Version string   `xml:"version,attr,omitempty"`
	OS      string   `xml:"os,attr,omitempty"`
}

type TimeQuery struct {
	// http://xmpp.org/extensions/xep-0202.html
	XMLName xml.Name `xml:"urn:xmpp:time time"`
	TZO     string   `xml:"tzo,omitempty"`
	UTC     string   `xml:"utc,omitempty"`
}

type DiscoInfoQuery struct {
	// http://xmpp.org/extensions/xep-0030.html
	XMLName xml.Name `xml:"http://jabber.org/protocol/disco#info query"`
}

type DiscoItemsQuery struct {
	// http://xmpp.org/extensions/xep-0030.html
	XMLName xml.Name `xml:"http://jabber.org/protocol/disco#items query"`
}

type PingQuery struct {
	// http://xmpp.org/extensions/xep-0199.html
	XMLName xml.Name `xml:"urn:xmpp:ping ping"`
}

type StatsQuery struct {
	// http://xmpp.org/extensions/xep-0039.html
	XMLName xml.Name `xml:"http://jabber.org/protocol/stats query"`
}

type LastQuery struct {
	// http://xmpp.org/extensions/xep-0012.html
	XMLName xml.Name `xml:"jabber:iq:last query"`
	Seconds int      `xml:"seconds,attr,omitempty"`
}

type PrivacyQuery struct {
	// http://xmpp.org/rfcs/rfc3921.html
	XMLName xml.Name `xml:"jabber:iq:privacy query"`
}

type IQ struct {
	Stanza
	XMLName         xml.Name `xml:"iq"`
	VersionQuery    VersionQuery
	TimeQuery       TimeQuery
	DiscoInfoQuery  DiscoInfoQuery
	DiscoItemsQuery DiscoItemsQuery
	PingQuery       PingQuery
	StatsQuery      StatsQuery
	LastQuery       LastQuery
	PrivacyQuery    PrivacyQuery
}

type Presence struct {
	Stanza
	XMLName  xml.Name `xml:"presence"`
	Show     string   `xml:"show,omitempty"`
	Status   string   `xml:"status,omitempty"`
	Priority int      `xml:"priority,omitempty"`
}

type StreamWrapper struct {
	RW io.ReadWriter
	Encoder *xml.Encoder
	Decoder *xml.Decoder
}

func NewStreamWrapper(rw io.ReadWriter) *StreamWrapper {
	return &StreamWrapper{
		RW: rw,
		Encoder: xml.NewEncoder(rw),
		Decoder: xml.NewDecoder(rw) }
}

func (sw *StreamWrapper) ReadStreamOpen() (*Stream, error) {
	stream := Stream{}
	for {
		t, err := sw.Decoder.Token()
		if err != nil { return nil, err }
		switch t := t.(type) {
		case xml.ProcInst:
			// Good.
		case xml.StartElement:
			if t.Name.Local == "stream" {
				for _, attr := range t.Attr {
					switch attr.Name.Local {
					case "to": stream.To = attr.Value
					case "from": stream.From = attr.Value
					case "version": stream.Version = attr.Value
					}
				}

				return &stream, nil
			}
		}
	}
}

// TODO(artem): refactor
func (sw *StreamWrapper) WriteStreamOpen(stream *Stream, default_namespace string) (err error) {
	data := xml.Header

	data += "<stream:stream xmlns='" + default_namespace + "' xmlns:stream='" + stream.XMLName.Space + "'"
	if stream.ID != "" {
		data += " id='" + stream.ID + "'"
	}
	if stream.From != "" {
		data += " from='" + stream.From + "'"
	}
	if stream.To != "" {
		data += " to='" + stream.To + "'"
	}
	if stream.Version != "" {
		data += " version='" + stream.Version + "'"
	}
	data += ">"

	_, err = io.WriteString(sw.RW, data)
	return
}

func (sw *StreamWrapper) ReadXMLChunk(types map[[2]string](func(xml.StartElement) interface{})) (interface{}, error) {
	for {
		t, err := sw.Decoder.Token()
		if err != nil { return nil, err }
		if element, ok := t.(xml.StartElement); ok {
			// TODO(artem): handle </stream:stream> etc
			if tp, ok := types[[2]string{element.Name.Local, element.Name.Space}]; ok {
				value := tp(element)
				err = sw.Decoder.DecodeElement(value, &element)
				return value, err
			}
		}
	}
}

/*
func ReadStanza(d *xml.Decoder) (interface{}, error) {
	var element xml.StartElement
	for {
		t, err := d.Token()
		if err != nil { return nil, err }
		var ok bool
		element, ok = t.(xml.StartElement)
		var stanza interface{}
		if ok {
			switch element.Name.Local {
			case "message": stanza = new(Message)
			case "iq": stanza = new(IQ)
			case "presence": stanza = new(Presence)
			}
		}
		err = d.DecodeElement(stanza, &element)
		return stanza, err
	}
}
*/
<commit_msg>PlainAuth struct for auth mech PLAIN<commit_after>package goxmpp

import (
	"encoding/xml"
	"io"
)

type Stream struct {
	XMLName    xml.Name `xml:"http://etherx.jabber.org/streams stream"`
	ID         string   `xml:"id,attr,omitempty"`
	From       string   `xml:"from,attr,omitempty"`
	To         string   `xml:"to,attr,omitempty"`
	Version    string   `xml:"version,attr,omitempty"`
}

type Mechanisms struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl mechanisms"`
	Names   []string `xml:"mechanism"`
}

type DigestMD5Auth struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl auth"`
	ID      string   `xml:"id,attr"`
}

type PlainAuth struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl auth"`
	Nonce   string   `xml:",chardata"`
}

type StartTLS struct {
	XMLName  xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-tls starttls"`
	Required bool     `xml:"required,omitempty"`
}

type Features struct {
	// http://tools.ietf.org/html/rfc3920#section-4.6
	XMLName          xml.Name    `xml:"stream:features"`
	Mechanisms       *Mechanisms
	StartTLS         *StartTLS
	Bind             bool        `xml:"urn:ietf:params:xml:ns:xmpp-bind bind,omitempty"`
	Session          bool        `xml:"urn:ietf:params:xml:ns:xmpp-session session,omitempty"`
}

type Error struct {
	Type   string `xml:"type,attr,omitempty"`
}

type Stanza struct {
	From   string `xml:"from,attr,omitempty"`
	To     string `xml:"to,attr,omitempty"`
	Type   string `xml:"type,attr,omitempty"`
	ID     string `xml:"id,attr,omitempty"`
	Error  Error
}

type Message struct {
	XMLName xml.Name `xml:"message"`
	Stanza
	Body    string   `xml:"body,omitempty"`
}

type VersionQuery struct {
	// http://xmpp.org/extensions/xep-0092.html
	XMLName xml.Name `xml:"jabber:iq:version query"`
	Name    string   `xml:"name,attr,omitempty"`
	Version string   `xml:"version,attr,omitempty"`
	OS      string   `xml:"os,attr,omitempty"`
}

type TimeQuery struct {
	// http://xmpp.org/extensions/xep-0202.html
	XMLName xml.Name `xml:"urn:xmpp:time time"`
	TZO     string   `xml:"tzo,omitempty"`
	UTC     string   `xml:"utc,omitempty"`
}

type DiscoInfoQuery struct {
	// http://xmpp.org/extensions/xep-0030.html
	XMLName xml.Name `xml:"http://jabber.org/protocol/disco#info query"`
}

type DiscoItemsQuery struct {
	// http://xmpp.org/extensions/xep-0030.html
	XMLName xml.Name `xml:"http://jabber.org/protocol/disco#items query"`
}

type PingQuery struct {
	// http://xmpp.org/extensions/xep-0199.html
	XMLName xml.Name `xml:"urn:xmpp:ping ping"`
}

type StatsQuery struct {
	// http://xmpp.org/extensions/xep-0039.html
	XMLName xml.Name `xml:"http://jabber.org/protocol/stats query"`
}

type LastQuery struct {
	// http://xmpp.org/extensions/xep-0012.html
	XMLName xml.Name `xml:"jabber:iq:last query"`
	Seconds int      `xml:"seconds,attr,omitempty"`
}

type PrivacyQuery struct {
	// http://xmpp.org/rfcs/rfc3921.html
	XMLName xml.Name `xml:"jabber:iq:privacy query"`
}

type IQ struct {
	Stanza
	XMLName         xml.Name `xml:"iq"`
	VersionQuery    VersionQuery
	TimeQuery       TimeQuery
	DiscoInfoQuery  DiscoInfoQuery
	DiscoItemsQuery DiscoItemsQuery
	PingQuery       PingQuery
	StatsQuery      StatsQuery
	LastQuery       LastQuery
	PrivacyQuery    PrivacyQuery
}

type Presence struct {
	Stanza
	XMLName  xml.Name `xml:"presence"`
	Show     string   `xml:"show,omitempty"`
	Status   string   `xml:"status,omitempty"`
	Priority int      `xml:"priority,omitempty"`
}

type StreamWrapper struct {
	RW io.ReadWriter
	Encoder *xml.Encoder
	Decoder *xml.Decoder
}

func NewStreamWrapper(rw io.ReadWriter) *StreamWrapper {
	return &StreamWrapper{
		RW: rw,
		Encoder: xml.NewEncoder(rw),
		Decoder: xml.NewDecoder(rw) }
}

func (sw *StreamWrapper) ReadStreamOpen() (*Stream, error) {
	stream := Stream{}
	for {
		t, err := sw.Decoder.Token()
		if err != nil { return nil, err }
		switch t := t.(type) {
		case xml.ProcInst:
			// Good.
		case xml.StartElement:
			if t.Name.Local == "stream" {
				for _, attr := range t.Attr {
					switch attr.Name.Local {
					case "to": stream.To = attr.Value
					case "from": stream.From = attr.Value
					case "version": stream.Version = attr.Value
					}
				}

				return &stream, nil
			}
		}
	}
}

// TODO(artem): refactor
func (sw *StreamWrapper) WriteStreamOpen(stream *Stream, default_namespace string) (err error) {
	data := xml.Header

	data += "<stream:stream xmlns='" + default_namespace + "' xmlns:stream='" + stream.XMLName.Space + "'"
	if stream.ID != "" {
		data += " id='" + stream.ID + "'"
	}
	if stream.From != "" {
		data += " from='" + stream.From + "'"
	}
	if stream.To != "" {
		data += " to='" + stream.To + "'"
	}
	if stream.Version != "" {
		data += " version='" + stream.Version + "'"
	}
	data += ">"

	_, err = io.WriteString(sw.RW, data)
	return
}

func (sw *StreamWrapper) ReadXMLChunk(types map[[2]string](func(xml.StartElement) interface{})) (interface{}, error) {
	for {
		t, err := sw.Decoder.Token()
		if err != nil { return nil, err }
		if element, ok := t.(xml.StartElement); ok {
			// TODO(artem): handle </stream:stream> etc
			if tp, ok := types[[2]string{element.Name.Local, element.Name.Space}]; ok {
				value := tp(element)
				err = sw.Decoder.DecodeElement(value, &element)
				return value, err
			}
		}
	}
}

/*
func ReadStanza(d *xml.Decoder) (interface{}, error) {
	var element xml.StartElement
	for {
		t, err := d.Token()
		if err != nil { return nil, err }
		var ok bool
		element, ok = t.(xml.StartElement)
		var stanza interface{}
		if ok {
			switch element.Name.Local {
			case "message": stanza = new(Message)
			case "iq": stanza = new(IQ)
			case "presence": stanza = new(Presence)
			}
		}
		err = d.DecodeElement(stanza, &element)
		return stanza, err
	}
}
*/
<|endoftext|>


<commit_before>package cluster

import (
	"crypto/tls"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"net"
	"time"

	"github.com/pkg/errors"

	lxd "github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/version"
)

// UserAgentNotifier used to distinguish between a regular client request and an internal cluster request when
// notifying other nodes of a cluster change.
const UserAgentNotifier = "lxd-cluster-notifier"

// UserAgentJoiner used to distinguish between a regular client request and an internal cluster request when
// joining a node to a cluster.
const UserAgentJoiner = "lxd-cluster-joiner"

// Connect is a convenience around lxd.ConnectLXD that configures the client
// with the correct parameters for node-to-node communication.
//
// If 'notify' switch is true, then the user agent will be set to the special
// to the UserAgentNotifier value, which can be used in some cases to distinguish
// between a regular client request and an internal cluster request.
func Connect(address string, cert *shared.CertInfo, notify bool) (lxd.InstanceServer, error) {
	// Wait for a connection to the events API first for non-notify connections.
	if !notify {
		connected := false
		for i := 0; i < 20; i++ {
			listenersLock.Lock()
			_, ok := listeners[address]
			listenersLock.Unlock()

			if ok {
				connected = true
				break
			}

			time.Sleep(500 * time.Millisecond)
		}

		if !connected {
			return nil, fmt.Errorf("Missing event connection with target cluster member")
		}
	}

	args := &lxd.ConnectionArgs{
		TLSServerCert: string(cert.PublicKey()),
		TLSClientCert: string(cert.PublicKey()),
		TLSClientKey:  string(cert.PrivateKey()),
		SkipGetServer: true,
		UserAgent:     version.UserAgent,
	}
	if notify {
		args.UserAgent = UserAgentNotifier
	}

	url := fmt.Sprintf("https://%s", address)
	return lxd.ConnectLXD(url, args)
}

// ConnectIfInstanceIsRemote figures out the address of the node which is
// running the container with the given name. If it's not the local node will
// connect to it and return the connected client, otherwise it will just return
// nil.
func ConnectIfInstanceIsRemote(cluster *db.Cluster, project, name string, cert *shared.CertInfo, instanceType instancetype.Type) (lxd.InstanceServer, error) {
	var address string // Node address
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		address, err = tx.GetNodeAddressOfInstance(project, name, instanceType)
		return err
	})
	if err != nil {
		return nil, err
	}
	if address == "" {
		// The container is running right on this node, no need to connect.
		return nil, nil
	}
	return Connect(address, cert, false)
}

// ConnectIfVolumeIsRemote figures out the address of the node on which the
// volume with the given name is defined. If it's not the local node will
// connect to it and return the connected client, otherwise it will just return
// nil.
//
// If there is more than one node with a matching volume name, an error is
// returned.
func ConnectIfVolumeIsRemote(cluster *db.Cluster, poolID int64, volumeName string, volumeType int, cert *shared.CertInfo) (lxd.InstanceServer, error) {
	var addresses []string // Node addresses
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		addresses, err = tx.GetStorageVolumeNodeAddresses(poolID, "default", volumeName, volumeType)
		return err
	})
	if err != nil {
		return nil, err
	}

	if len(addresses) > 1 {
		var driver string
		err := cluster.Transaction(func(tx *db.ClusterTx) error {
			var err error
			driver, err = tx.GetStoragePoolDriver(poolID)
			return err
		})
		if err != nil {
			return nil, err
		}

		if driver == "ceph" || driver == "cephfs" {
			return nil, nil
		}

		return nil, fmt.Errorf("more than one node has a volume named %s", volumeName)
	}

	address := addresses[0]
	if address == "" {
		return nil, nil
	}

	return Connect(address, cert, false)
}

// SetupTrust is a convenience around InstanceServer.CreateCertificate that
// adds the given client certificate to the trusted pool of the cluster at the
// given address, using the given password.
func SetupTrust(cert, targetAddress, targetCert, targetPassword string) error {
	// Connect to the target cluster node.
	args := &lxd.ConnectionArgs{
		TLSServerCert: targetCert,
		UserAgent:     version.UserAgent,
	}

	target, err := lxd.ConnectLXD(fmt.Sprintf("https://%s", targetAddress), args)
	if err != nil {
		return errors.Wrap(err, "failed to connect to target cluster node")
	}

	block, _ := pem.Decode([]byte(cert))
	if block == nil {
		return fmt.Errorf("failed to decode certificate")
	}

	certificate := base64.StdEncoding.EncodeToString(block.Bytes)
	post := api.CertificatesPost{
		Password:    targetPassword,
		Certificate: certificate,
	}

	fingerprint, err := shared.CertFingerprintStr(cert)
	if err != nil {
		return errors.Wrap(err, "failed to calculate fingerprint")
	}

	post.Name = fmt.Sprintf("lxd.cluster.%s", fingerprint)
	post.Type = "client"

	err = target.CreateCertificate(post)
	if err != nil && err.Error() != "Certificate already in trust store" {
		return errors.Wrap(err, "Failed to add client cert to cluster")
	}

	return nil
}

// HasConnectivity probes the member with the given address for connectivity.
func HasConnectivity(cert *shared.CertInfo, address string) bool {
	config, err := tlsClientConfig(cert)
	if err != nil {
		return false
	}

	var conn net.Conn
	dialer := &net.Dialer{Timeout: time.Second}
	conn, err = tls.DialWithDialer(dialer, "tcp", address, config)
	if err == nil {
		conn.Close()
		return true
	}
	return false
}
<commit_msg>lxd/cluster/connect: Adds ClientType type and UserAgentClientType function<commit_after>package cluster

import (
	"crypto/tls"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"net"
	"time"

	"github.com/pkg/errors"

	lxd "github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/version"
)

// UserAgentNotifier used to distinguish between a regular client request and an internal cluster request when
// notifying other nodes of a cluster change.
const UserAgentNotifier = "lxd-cluster-notifier"

// UserAgentJoiner used to distinguish between a regular client request and an internal cluster request when
// joining a node to a cluster.
const UserAgentJoiner = "lxd-cluster-joiner"

// ClientType indicates which sort of client type is being used.
type ClientType string

// ClientTypeNotifier cluster notification client.
const ClientTypeNotifier ClientType = "notifier"

// ClientTypeJoiner cluster joiner client.
const ClientTypeJoiner ClientType = "joiner"

// ClientTypeNormal normal client.
const ClientTypeNormal ClientType = "normal"

// UserAgentClientType converts user agent to client type.
func UserAgentClientType(userAgent string) ClientType {
	switch userAgent {
	case UserAgentNotifier:
		return ClientTypeNotifier
	case UserAgentJoiner:
		return ClientTypeJoiner
	}

	return ClientTypeNormal
}

// Connect is a convenience around lxd.ConnectLXD that configures the client
// with the correct parameters for node-to-node communication.
//
// If 'notify' switch is true, then the user agent will be set to the special
// to the UserAgentNotifier value, which can be used in some cases to distinguish
// between a regular client request and an internal cluster request.
func Connect(address string, cert *shared.CertInfo, notify bool) (lxd.InstanceServer, error) {
	// Wait for a connection to the events API first for non-notify connections.
	if !notify {
		connected := false
		for i := 0; i < 20; i++ {
			listenersLock.Lock()
			_, ok := listeners[address]
			listenersLock.Unlock()

			if ok {
				connected = true
				break
			}

			time.Sleep(500 * time.Millisecond)
		}

		if !connected {
			return nil, fmt.Errorf("Missing event connection with target cluster member")
		}
	}

	args := &lxd.ConnectionArgs{
		TLSServerCert: string(cert.PublicKey()),
		TLSClientCert: string(cert.PublicKey()),
		TLSClientKey:  string(cert.PrivateKey()),
		SkipGetServer: true,
		UserAgent:     version.UserAgent,
	}
	if notify {
		args.UserAgent = UserAgentNotifier
	}

	url := fmt.Sprintf("https://%s", address)
	return lxd.ConnectLXD(url, args)
}

// ConnectIfInstanceIsRemote figures out the address of the node which is
// running the container with the given name. If it's not the local node will
// connect to it and return the connected client, otherwise it will just return
// nil.
func ConnectIfInstanceIsRemote(cluster *db.Cluster, project, name string, cert *shared.CertInfo, instanceType instancetype.Type) (lxd.InstanceServer, error) {
	var address string // Node address
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		address, err = tx.GetNodeAddressOfInstance(project, name, instanceType)
		return err
	})
	if err != nil {
		return nil, err
	}
	if address == "" {
		// The container is running right on this node, no need to connect.
		return nil, nil
	}
	return Connect(address, cert, false)
}

// ConnectIfVolumeIsRemote figures out the address of the node on which the
// volume with the given name is defined. If it's not the local node will
// connect to it and return the connected client, otherwise it will just return
// nil.
//
// If there is more than one node with a matching volume name, an error is
// returned.
func ConnectIfVolumeIsRemote(cluster *db.Cluster, poolID int64, volumeName string, volumeType int, cert *shared.CertInfo) (lxd.InstanceServer, error) {
	var addresses []string // Node addresses
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		addresses, err = tx.GetStorageVolumeNodeAddresses(poolID, "default", volumeName, volumeType)
		return err
	})
	if err != nil {
		return nil, err
	}

	if len(addresses) > 1 {
		var driver string
		err := cluster.Transaction(func(tx *db.ClusterTx) error {
			var err error
			driver, err = tx.GetStoragePoolDriver(poolID)
			return err
		})
		if err != nil {
			return nil, err
		}

		if driver == "ceph" || driver == "cephfs" {
			return nil, nil
		}

		return nil, fmt.Errorf("more than one node has a volume named %s", volumeName)
	}

	address := addresses[0]
	if address == "" {
		return nil, nil
	}

	return Connect(address, cert, false)
}

// SetupTrust is a convenience around InstanceServer.CreateCertificate that
// adds the given client certificate to the trusted pool of the cluster at the
// given address, using the given password.
func SetupTrust(cert, targetAddress, targetCert, targetPassword string) error {
	// Connect to the target cluster node.
	args := &lxd.ConnectionArgs{
		TLSServerCert: targetCert,
		UserAgent:     version.UserAgent,
	}

	target, err := lxd.ConnectLXD(fmt.Sprintf("https://%s", targetAddress), args)
	if err != nil {
		return errors.Wrap(err, "failed to connect to target cluster node")
	}

	block, _ := pem.Decode([]byte(cert))
	if block == nil {
		return fmt.Errorf("failed to decode certificate")
	}

	certificate := base64.StdEncoding.EncodeToString(block.Bytes)
	post := api.CertificatesPost{
		Password:    targetPassword,
		Certificate: certificate,
	}

	fingerprint, err := shared.CertFingerprintStr(cert)
	if err != nil {
		return errors.Wrap(err, "failed to calculate fingerprint")
	}

	post.Name = fmt.Sprintf("lxd.cluster.%s", fingerprint)
	post.Type = "client"

	err = target.CreateCertificate(post)
	if err != nil && err.Error() != "Certificate already in trust store" {
		return errors.Wrap(err, "Failed to add client cert to cluster")
	}

	return nil
}

// HasConnectivity probes the member with the given address for connectivity.
func HasConnectivity(cert *shared.CertInfo, address string) bool {
	config, err := tlsClientConfig(cert)
	if err != nil {
		return false
	}

	var conn net.Conn
	dialer := &net.Dialer{Timeout: time.Second}
	conn, err = tls.DialWithDialer(dialer, "tcp", address, config)
	if err == nil {
		conn.Close()
		return true
	}
	return false
}
<|endoftext|>


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

package api

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/control-center/serviced/auth"
	"github.com/control-center/serviced/config"
	"github.com/control-center/serviced/dao"
	"github.com/control-center/serviced/dao/client"
	"github.com/control-center/serviced/rpc/agent"
	"github.com/control-center/serviced/rpc/master"
	"github.com/control-center/serviced/utils"
	"github.com/control-center/serviced/validation"

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

type api struct {
	master master.ClientInterface
	agent  *agent.Client
	docker *dockerclient.Client
	dao    dao.ControlPlane // Deprecated
}

var hostAuthenticated bool

func New() API {
	// let lazy init populate each interface as necessary
	return NewAPI(nil, nil, nil, nil)
}

// New creates a new API type
func NewAPI(master master.ClientInterface, agent *agent.Client, docker *dockerclient.Client, dao dao.ControlPlane) API {
	return &api{master: master, agent: agent, docker: docker, dao: dao}
}

// Starts the agent or master services on this host
func (a *api) StartServer() error {
	options := config.GetOptions()
	configureLoggingForLogstash(options.LogstashURL)
	log := log.WithFields(logrus.Fields{
		"staticips": options.StaticIPs,
	})
	log.Debug("Starting server")

	if err := utils.SetCiphers("http", options.TLSCiphers); err != nil {
		return fmt.Errorf("unable to set HTTP TLS Ciphers %v", err)
	}
	log.WithFields(logrus.Fields{
		"ciphers": strings.Join(options.TLSCiphers, ","),
	}).Debug("Set supported TLS ciphers for HTTP")

	if err := utils.SetMinTLS("http", options.TLSMinVersion); err != nil {
		return fmt.Errorf("unable to set minimum HTTP TLS version %v", err)
	}
	log.WithFields(logrus.Fields{
		"minversion": options.TLSMinVersion,
	}).Debug("Set minimum TLS version for HTTP")

	if err := utils.SetCiphers("mux", options.MUXTLSCiphers); err != nil {
		return fmt.Errorf("unable to set MUX TLS Ciphers %v", err)
	}
	log.WithFields(logrus.Fields{
		"ciphers": strings.Join(options.MUXTLSCiphers, ","),
	}).Debug("Set supported TLS ciphers for the mux")

	if err := utils.SetMinTLS("mux", options.MUXTLSMinVersion); err != nil {
		return fmt.Errorf("unable to set minimum MUX TLS version %v", err)
	}
	log.WithFields(logrus.Fields{
		"minversion": options.MUXTLSMinVersion,
	}).Debug("Set minimum TLS version for the mux")

	if err := utils.SetCiphers("rpc", options.RPCTLSCiphers); err != nil {
		return fmt.Errorf("unable to set RPC TLS Ciphers %v", err)
	}
	log.WithFields(logrus.Fields{
		"ciphers": strings.Join(options.RPCTLSCiphers, ","),
	}).Debug("Set supported TLS ciphers for RPC")

	if err := utils.SetMinTLS("rpc", options.RPCTLSMinVersion); err != nil {
		return fmt.Errorf("unable to set minimum RPC TLS version %v", err)
	}
	log.WithFields(logrus.Fields{
		"minversion": options.RPCTLSMinVersion,
	}).Debug("Set minimum TLS version for RPC")

	if len(options.CPUProfile) > 0 {
		f, err := os.Create(options.CPUProfile)
		if err != nil {
			log.WithFields(logrus.Fields{
				"file": options.CPUProfile,
			}).WithError(err).Fatal("Unable to create CPU profile file")
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	d, err := newDaemon(options.Endpoint, options.StaticIPs, options.MasterPoolID, time.Duration(options.TokenExpiration)*time.Second)
	if err != nil {
		return err
	}
	return d.run()
}

func configureLoggingForLogstash(logstashURL string) {
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "unknown"
	}
	glog.SetLogstashType("serviced-" + hostname)
	glog.SetLogstashURL(logstashURL)
}

// Opens a connection to the master if not already connected
func (a *api) connectMaster() (master.ClientInterface, error) {
	if a.master == nil {
		var err error
		a.master, err = master.NewClient(config.GetOptions().Endpoint)
		if err != nil {
			return nil, fmt.Errorf("could not create a client to the master: %s", err)
		}
		if err := a.authenticateHost(); err != nil {
			return nil, err
		}
	}
	return a.master, nil
}

// Opens a connection to the agent if not already connected
func (a *api) connectAgent(address string) (*agent.Client, error) {
	if a.agent == nil {
		var err error
		a.agent, err = agent.NewClient(address)
		if err != nil {
			return nil, fmt.Errorf("could not create a client to the agent: %s", err)
		}
		if err := a.authenticateHost(); err != nil {
			return nil, err
		}
	}
	return a.agent, nil
}

// Opens a connection to docker if not already connected
func (a *api) connectDocker() (*dockerclient.Client, error) {
	if a.docker == nil {
		const DockerEndpoint string = "unix:///var/run/docker.sock"
		var err error
		if a.docker, err = dockerclient.NewClient(DockerEndpoint); err != nil {
			return nil, fmt.Errorf("could not create a client to docker: %s", err)
		}
	}
	return a.docker, nil
}

// DEPRECATED: Opens a connection to the DAO if not already connected
func (a *api) connectDAO() (dao.ControlPlane, error) {
	if a.dao == nil {
		var err error
		a.dao, err = client.NewControlClient(config.GetOptions().Endpoint)
		if err != nil {
			return nil, fmt.Errorf("could not create a client to the agent: %s", err)
		}
		if err := a.authenticateHost(); err != nil {
			return nil, err
		}
	}
	return a.dao, nil
}

// This will authenticate the host once to get a valid token for any CLI commands
//  that require it.
func (a *api) authenticateHost() error {
	if hostAuthenticated {
		return nil
	}

	options := config.GetOptions()

	// Try to load the master keys, fail silently if they don't exist
	masterKeyFile := filepath.Join(options.IsvcsPath, auth.MasterKeyFileName)
	if err := auth.LoadMasterKeyFile(masterKeyFile); err != nil {
		log.WithError(err).Debug("Unable to load master keys")
	}

	// Load the delegate keys
	delegateKeyFile := filepath.Join(options.EtcPath, auth.DelegateKeyFileName)
	if err := auth.LoadDelegateKeysFromFile(delegateKeyFile); err != nil {
		return err
	}

	// Get our host ID
	myHostID, err := utils.HostID()
	if err != nil {
		return err
	} else if err := validation.ValidHostID(myHostID); err != nil {
		return err
	}

	// Load an auth token once
	tokenFile := filepath.Join(options.EtcPath, auth.TokenFileName)
	getToken := func() (string, int64, error) {
		return a.AuthenticateHost(myHostID)
	}

	if _, err := auth.RefreshToken(getToken, tokenFile); err != nil {
		return err
	}

	hostAuthenticated = true
	return nil
}
<commit_msg>fix smoke test failure<commit_after>// Copyright 2014 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//	   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package api

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/control-center/serviced/auth"
	"github.com/control-center/serviced/config"
	"github.com/control-center/serviced/dao"
	"github.com/control-center/serviced/dao/client"
	"github.com/control-center/serviced/rpc/agent"
	"github.com/control-center/serviced/rpc/master"
	"github.com/control-center/serviced/utils"
	"github.com/control-center/serviced/validation"

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

type api struct {
	master master.ClientInterface
	agent  *agent.Client
	docker *dockerclient.Client
	dao    dao.ControlPlane // Deprecated
}

var hostAuthenticated bool

func New() API {
	// let lazy init populate each interface as necessary
	return NewAPI(nil, nil, nil, nil)
}

// New creates a new API type
func NewAPI(master master.ClientInterface, agent *agent.Client, docker *dockerclient.Client, dao dao.ControlPlane) API {
	return &api{master: master, agent: agent, docker: docker, dao: dao}
}

// Starts the agent or master services on this host
func (a *api) StartServer() error {
	options := config.GetOptions()
	configureLoggingForLogstash(options.LogstashURL)
	log := log.WithFields(logrus.Fields{
		"staticips": options.StaticIPs,
	})
	log.Debug("Starting server")

	if err := utils.SetCiphers("http", options.TLSCiphers); err != nil {
		return fmt.Errorf("unable to set HTTP TLS Ciphers %v", err)
	}
	log.WithFields(logrus.Fields{
		"ciphers": strings.Join(options.TLSCiphers, ","),
	}).Debug("Set supported TLS ciphers for HTTP")

	if err := utils.SetMinTLS("http", options.TLSMinVersion); err != nil {
		return fmt.Errorf("unable to set minimum HTTP TLS version %v", err)
	}
	log.WithFields(logrus.Fields{
		"minversion": options.TLSMinVersion,
	}).Debug("Set minimum TLS version for HTTP")

	if err := utils.SetCiphers("mux", options.MUXTLSCiphers); err != nil {
		return fmt.Errorf("unable to set MUX TLS Ciphers %v", err)
	}
	log.WithFields(logrus.Fields{
		"ciphers": strings.Join(options.MUXTLSCiphers, ","),
	}).Debug("Set supported TLS ciphers for the mux")

	if err := utils.SetMinTLS("mux", options.MUXTLSMinVersion); err != nil {
		return fmt.Errorf("unable to set minimum MUX TLS version %v", err)
	}
	log.WithFields(logrus.Fields{
		"minversion": options.MUXTLSMinVersion,
	}).Debug("Set minimum TLS version for the mux")

	if err := utils.SetCiphers("rpc", options.RPCTLSCiphers); err != nil {
		return fmt.Errorf("unable to set RPC TLS Ciphers %v", err)
	}
	log.WithFields(logrus.Fields{
		"ciphers": strings.Join(options.RPCTLSCiphers, ","),
	}).Debug("Set supported TLS ciphers for RPC")

	if err := utils.SetMinTLS("rpc", options.RPCTLSMinVersion); err != nil {
		return fmt.Errorf("unable to set minimum RPC TLS version %v", err)
	}
	log.WithFields(logrus.Fields{
		"minversion": options.RPCTLSMinVersion,
	}).Debug("Set minimum TLS version for RPC")

	if len(options.CPUProfile) > 0 {
		f, err := os.Create(options.CPUProfile)
		if err != nil {
			log.WithFields(logrus.Fields{
				"file": options.CPUProfile,
			}).WithError(err).Fatal("Unable to create CPU profile file")
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	d, err := newDaemon(options.Endpoint, options.StaticIPs, options.MasterPoolID, time.Duration(options.TokenExpiration)*time.Second)
	if err != nil {
		return err
	}
	return d.run()
}

func configureLoggingForLogstash(logstashURL string) {
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "unknown"
	}
	glog.SetLogstashType("serviced-" + hostname)
	glog.SetLogstashURL(logstashURL)
}

// Opens a connection to the master if not already connected
func (a *api) connectMaster() (master.ClientInterface, error) {
	if a.master == nil {
		var err error
		a.master, err = master.NewClient(config.GetOptions().Endpoint)
		if err != nil {
			return nil, fmt.Errorf("could not create a client to the master: %s", err)
		}
		a.authenticateHost()
	}
	return a.master, nil
}

// Opens a connection to the agent if not already connected
func (a *api) connectAgent(address string) (*agent.Client, error) {
	if a.agent == nil {
		var err error
		a.agent, err = agent.NewClient(address)
		if err != nil {
			return nil, fmt.Errorf("could not create a client to the agent: %s", err)
		}
		a.authenticateHost()
	}
	return a.agent, nil
}

// Opens a connection to docker if not already connected
func (a *api) connectDocker() (*dockerclient.Client, error) {
	if a.docker == nil {
		const DockerEndpoint string = "unix:///var/run/docker.sock"
		var err error
		if a.docker, err = dockerclient.NewClient(DockerEndpoint); err != nil {
			return nil, fmt.Errorf("could not create a client to docker: %s", err)
		}
	}
	return a.docker, nil
}

// DEPRECATED: Opens a connection to the DAO if not already connected
func (a *api) connectDAO() (dao.ControlPlane, error) {
	if a.dao == nil {
		var err error
		a.dao, err = client.NewControlClient(config.GetOptions().Endpoint)
		if err != nil {
			return nil, fmt.Errorf("could not create a client to the agent: %s", err)
		}
		a.authenticateHost()
	}
	return a.dao, nil
}

// This will authenticate the host once to get a valid token for any CLI commands
//  that require it.
func (a *api) authenticateHost() error {
	if hostAuthenticated {
		return nil
	}

	options := config.GetOptions()

	// Try to load the master keys, fail silently if they don't exist
	masterKeyFile := filepath.Join(options.IsvcsPath, auth.MasterKeyFileName)
	if err := auth.LoadMasterKeyFile(masterKeyFile); err != nil {
		log.WithError(err).Debug("Unable to load master keys")
	}

	// Load the delegate keys
	delegateKeyFile := filepath.Join(options.EtcPath, auth.DelegateKeyFileName)
	if err := auth.LoadDelegateKeysFromFile(delegateKeyFile); err != nil {
		return err
	}

	// Get our host ID
	myHostID, err := utils.HostID()
	if err != nil {
		return err
	} else if err := validation.ValidHostID(myHostID); err != nil {
		return err
	}

	// Load an auth token once
	tokenFile := filepath.Join(options.EtcPath, auth.TokenFileName)
	getToken := func() (string, int64, error) {
		return a.AuthenticateHost(myHostID)
	}

	if _, err := auth.RefreshToken(getToken, tokenFile); err != nil {
		return err
	}

	hostAuthenticated = true
	return nil
}
<|endoftext|>


<commit_before>package dnsmasq

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/lxc/lxd/lxd/project"
	"github.com/lxc/lxd/lxd/storage/filesystem"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/subprocess"
	"github.com/lxc/lxd/shared/version"
)

const staticAllocationDeviceSeparator = "."

// DHCPAllocation represents an IP allocation from dnsmasq.
type DHCPAllocation struct {
	IP     net.IP
	Name   string
	MAC    net.HardwareAddr
	Static bool
}

// ConfigMutex used to coordinate access to the dnsmasq config files.
var ConfigMutex sync.Mutex

// UpdateStaticEntry writes a single dhcp-host line for a network/instance combination.
func UpdateStaticEntry(network string, projectName string, instanceName string, deviceName string, netConfig map[string]string, hwaddr string, ipv4Address string, ipv6Address string) error {
	hwaddr = strings.ToLower(hwaddr)
	line := hwaddr

	// Generate the dhcp-host line
	if ipv4Address != "" {
		line += fmt.Sprintf(",%s", ipv4Address)
	}

	if ipv6Address != "" {
		line += fmt.Sprintf(",[%s]", ipv6Address)
	}

	if netConfig["dns.mode"] == "" || netConfig["dns.mode"] == "managed" {
		line += fmt.Sprintf(",%s", project.DNS(projectName, instanceName))
	}

	if line == hwaddr {
		return nil
	}

	entryFile := dnsMasqEntryFileName(projectName, instanceName, deviceName)
	err := ioutil.WriteFile(shared.VarPath("networks", network, "dnsmasq.hosts", entryFile), []byte(line+"\n"), 0644)
	if err != nil {
		return err
	}

	return nil
}

// RemoveStaticEntry removes a single dhcp-host line for a network/instance combination.
func RemoveStaticEntry(network string, projectName string, instanceName string, deviceName string) error {
	entryFile := dnsMasqEntryFileName(projectName, instanceName, deviceName)
	err := os.Remove(shared.VarPath("networks", network, "dnsmasq.hosts", entryFile))
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	return nil
}

// Kill kills dnsmasq for a particular network (or optionally reloads it).
func Kill(name string, reload bool) error {
	pidPath := shared.VarPath("networks", name, "dnsmasq.pid")

	// If the pid file doesn't exist, there is no process to kill.
	if !shared.PathExists(pidPath) {
		return nil
	}

	// Import saved subprocess details
	p, err := subprocess.ImportProcess(pidPath)
	if err != nil {
		return fmt.Errorf("Could not read pid file: %s", err)
	}

	if reload {
		err = p.Reload()
		if err != nil && err != subprocess.ErrNotRunning {
			return fmt.Errorf("Could not reload dnsmasq: %s", err)
		}

		return nil
	}

	err = p.Stop()
	if err != nil && err != subprocess.ErrNotRunning {
		return fmt.Errorf("Unable to kill dnsmasq: %s", err)
	}

	time.Sleep(100 * time.Millisecond) // Give OS time to release sockets.

	return nil
}

// GetVersion returns the version of dnsmasq.
func GetVersion() (*version.DottedVersion, error) {
	output, err := shared.RunCommandCLocale("dnsmasq", "--version")
	if err != nil {
		return nil, fmt.Errorf("Failed to check dnsmasq version: %v", err)
	}

	lines := strings.Split(string(output), " ")
	return version.Parse(lines[2])
}

// DHCPStaticAllocation retrieves the dnsmasq statically allocated MAC and IPs for an instance.
// Returns MAC, IPv4 and IPv6 DHCPAllocation structs respectively.
// The projectName, instanceName, and deviceName are used to generate the entryFileName. If the entryFileName
// argument is provided, then projectName, instanceName, and deviceName are ignored.
func DHCPStaticAllocation(network, projectName, instanceName, deviceName, entryFileName string) (net.HardwareAddr, DHCPAllocation, DHCPAllocation, error) {
	var IPv4, IPv6 DHCPAllocation
	var mac net.HardwareAddr

	if entryFileName == "" {
		entryFileName = dnsMasqEntryFileName(projectName, instanceName, deviceName)
	}

	file, err := os.Open(shared.VarPath("networks", network, "dnsmasq.hosts", entryFileName))
	if err != nil {
		return nil, IPv4, IPv6, err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		fields := strings.SplitN(scanner.Text(), ",", -1)
		for _, field := range fields {
			// Check if field is IPv4 or IPv6 address.
			if strings.Count(field, ".") == 3 {
				IP := net.ParseIP(field)
				if IP.To4() == nil {
					return nil, IPv4, IPv6, fmt.Errorf("Error parsing IP address %q", field)
				}
				IPv4 = DHCPAllocation{Name: instanceName, Static: true, IP: IP.To4(), MAC: mac}

			} else if strings.HasPrefix(field, "[") && strings.HasSuffix(field, "]") {
				IP := net.ParseIP(field[1 : len(field)-1])
				if IP == nil {
					return nil, IPv4, IPv6, fmt.Errorf("Error parsing IP address %q", field)
				}
				IPv6 = DHCPAllocation{Name: instanceName, Static: true, IP: IP, MAC: mac}
			} else if strings.Count(field, ":") == 5 {
				// This field is expected to come first, so that mac variable can be used with
				// populating the DHCPAllocation structs too.
				mac, err = net.ParseMAC(field)
				if err != nil {
					return nil, IPv4, IPv6, fmt.Errorf("Error parsing MAC address %q", field)
				}
			}
		}
	}
	if err := scanner.Err(); err != nil {
		return nil, IPv4, IPv6, err
	}

	return mac, IPv4, IPv6, nil
}

// DHCPAllAllocations returns a map of IPs currently allocated (statically and dynamically)
// in dnsmasq for a specific network. The returned map is keyed by a 16 byte array representing
// the net.IP format. The value of each map item is a DHCPAllocation struct containing at least
// whether the allocation was static or dynamic and optionally instance name or MAC address.
// MAC addresses are only included for dynamic IPv4 allocations (where name is not reliable).
// Static allocations are not overridden by dynamic allocations, allowing for instance name to be
// included for static IPv6 allocations. IPv6 addresses that are dynamically assigned cannot be
// reliably linked to instances using either name or MAC because dnsmasq does not record the MAC
// address for these records, and the recorded host name can be set by the instance if the dns.mode
// for the network is set to "dynamic" and so cannot be trusted, so in this case we do not return
// any identifying info.
func DHCPAllAllocations(network string) (map[[4]byte]DHCPAllocation, map[[16]byte]DHCPAllocation, error) {
	IPv4s := make(map[[4]byte]DHCPAllocation)
	IPv6s := make(map[[16]byte]DHCPAllocation)

	// First read all statically allocated IPs.
	files, err := ioutil.ReadDir(shared.VarPath("networks", network, "dnsmasq.hosts"))
	if err != nil && os.IsNotExist(err) {
		return nil, nil, err
	}

	for _, entry := range files {
		_, IPv4, IPv6, err := DHCPStaticAllocation(network, "", "", "", entry.Name())
		if err != nil {
			return nil, nil, err
		}

		if IPv4.IP != nil {
			var IPKey [4]byte
			copy(IPKey[:], IPv4.IP.To4())
			IPv4s[IPKey] = IPv4
		}

		if IPv6.IP != nil {
			var IPKey [16]byte
			copy(IPKey[:], IPv6.IP.To16())
			IPv6s[IPKey] = IPv6
		}
	}

	// Next read all dynamic allocated IPs.
	file, err := os.Open(shared.VarPath("networks", network, "dnsmasq.leases"))
	if err != nil {
		return nil, nil, err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())
		if len(fields) == 5 {
			IP := net.ParseIP(fields[2])
			if IP == nil {
				return nil, nil, fmt.Errorf("Error parsing IP address: %v", fields[2])
			}

			// Handle IPv6 addresses.
			if IP.To4() == nil {
				var IPKey [16]byte
				copy(IPKey[:], IP.To16())

				// Don't replace IPs from static config as more reliable.
				if IPv6s[IPKey].Name != "" {
					continue
				}

				IPv6s[IPKey] = DHCPAllocation{
					Static: false,
					IP:     IP.To16(),
				}
			} else {
				// MAC only available in IPv4 leases.
				MAC, err := net.ParseMAC(fields[1])
				if err != nil {
					return nil, nil, err
				}

				var IPKey [4]byte
				copy(IPKey[:], IP.To4())

				// Don't replace IPs from static config as more reliable.
				if IPv4s[IPKey].Name != "" {
					continue
				}

				IPv4s[IPKey] = DHCPAllocation{
					MAC:    MAC,
					Static: false,
					IP:     IP.To4(),
				}
			}
		}
	}
	if err := scanner.Err(); err != nil {
		return nil, nil, err
	}

	return IPv4s, IPv6s, nil
}

func dnsMasqEntryFileName(projectName string, instanceName string, deviceName string) string {
	escapedDeviceName := filesystem.PathNameEncode(deviceName)

	return strings.Join([]string{project.Instance(projectName, instanceName), escapedDeviceName}, ".")
}
<commit_msg>lxd/dnsmasq: Renames dnsMasqEntryFileName to StaticAllocationFileName<commit_after>package dnsmasq

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/lxc/lxd/lxd/project"
	"github.com/lxc/lxd/lxd/storage/filesystem"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/subprocess"
	"github.com/lxc/lxd/shared/version"
)

const staticAllocationDeviceSeparator = "."

// DHCPAllocation represents an IP allocation from dnsmasq.
type DHCPAllocation struct {
	IP     net.IP
	Name   string
	MAC    net.HardwareAddr
	Static bool
}

// ConfigMutex used to coordinate access to the dnsmasq config files.
var ConfigMutex sync.Mutex

// UpdateStaticEntry writes a single dhcp-host line for a network/instance combination.
func UpdateStaticEntry(network string, projectName string, instanceName string, deviceName string, netConfig map[string]string, hwaddr string, ipv4Address string, ipv6Address string) error {
	hwaddr = strings.ToLower(hwaddr)
	line := hwaddr

	// Generate the dhcp-host line
	if ipv4Address != "" {
		line += fmt.Sprintf(",%s", ipv4Address)
	}

	if ipv6Address != "" {
		line += fmt.Sprintf(",[%s]", ipv6Address)
	}

	if netConfig["dns.mode"] == "" || netConfig["dns.mode"] == "managed" {
		line += fmt.Sprintf(",%s", project.DNS(projectName, instanceName))
	}

	if line == hwaddr {
		return nil
	}

	entryFile := dnsMasqEntryFileName(projectName, instanceName, deviceName)
	err := ioutil.WriteFile(shared.VarPath("networks", network, "dnsmasq.hosts", entryFile), []byte(line+"\n"), 0644)
	if err != nil {
		return err
	}

	return nil
}

// RemoveStaticEntry removes a single dhcp-host line for a network/instance combination.
func RemoveStaticEntry(network string, projectName string, instanceName string, deviceName string) error {
	entryFile := dnsMasqEntryFileName(projectName, instanceName, deviceName)
	err := os.Remove(shared.VarPath("networks", network, "dnsmasq.hosts", entryFile))
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	return nil
}

// Kill kills dnsmasq for a particular network (or optionally reloads it).
func Kill(name string, reload bool) error {
	pidPath := shared.VarPath("networks", name, "dnsmasq.pid")

	// If the pid file doesn't exist, there is no process to kill.
	if !shared.PathExists(pidPath) {
		return nil
	}

	// Import saved subprocess details
	p, err := subprocess.ImportProcess(pidPath)
	if err != nil {
		return fmt.Errorf("Could not read pid file: %s", err)
	}

	if reload {
		err = p.Reload()
		if err != nil && err != subprocess.ErrNotRunning {
			return fmt.Errorf("Could not reload dnsmasq: %s", err)
		}

		return nil
	}

	err = p.Stop()
	if err != nil && err != subprocess.ErrNotRunning {
		return fmt.Errorf("Unable to kill dnsmasq: %s", err)
	}

	time.Sleep(100 * time.Millisecond) // Give OS time to release sockets.

	return nil
}

// GetVersion returns the version of dnsmasq.
func GetVersion() (*version.DottedVersion, error) {
	output, err := shared.RunCommandCLocale("dnsmasq", "--version")
	if err != nil {
		return nil, fmt.Errorf("Failed to check dnsmasq version: %v", err)
	}

	lines := strings.Split(string(output), " ")
	return version.Parse(lines[2])
}

// DHCPStaticAllocation retrieves the dnsmasq statically allocated MAC and IPs for an instance.
// Returns MAC, IPv4 and IPv6 DHCPAllocation structs respectively.
// The projectName, instanceName, and deviceName are used to generate the entryFileName. If the entryFileName
// argument is provided, then projectName, instanceName, and deviceName are ignored.
func DHCPStaticAllocation(network, projectName, instanceName, deviceName, entryFileName string) (net.HardwareAddr, DHCPAllocation, DHCPAllocation, error) {
	var IPv4, IPv6 DHCPAllocation
	var mac net.HardwareAddr

	if entryFileName == "" {
		entryFileName = dnsMasqEntryFileName(projectName, instanceName, deviceName)
	}

	file, err := os.Open(shared.VarPath("networks", network, "dnsmasq.hosts", entryFileName))
	if err != nil {
		return nil, IPv4, IPv6, err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		fields := strings.SplitN(scanner.Text(), ",", -1)
		for _, field := range fields {
			// Check if field is IPv4 or IPv6 address.
			if strings.Count(field, ".") == 3 {
				IP := net.ParseIP(field)
				if IP.To4() == nil {
					return nil, IPv4, IPv6, fmt.Errorf("Error parsing IP address %q", field)
				}
				IPv4 = DHCPAllocation{Name: instanceName, Static: true, IP: IP.To4(), MAC: mac}

			} else if strings.HasPrefix(field, "[") && strings.HasSuffix(field, "]") {
				IP := net.ParseIP(field[1 : len(field)-1])
				if IP == nil {
					return nil, IPv4, IPv6, fmt.Errorf("Error parsing IP address %q", field)
				}
				IPv6 = DHCPAllocation{Name: instanceName, Static: true, IP: IP, MAC: mac}
			} else if strings.Count(field, ":") == 5 {
				// This field is expected to come first, so that mac variable can be used with
				// populating the DHCPAllocation structs too.
				mac, err = net.ParseMAC(field)
				if err != nil {
					return nil, IPv4, IPv6, fmt.Errorf("Error parsing MAC address %q", field)
				}
			}
		}
	}
	if err := scanner.Err(); err != nil {
		return nil, IPv4, IPv6, err
	}

	return mac, IPv4, IPv6, nil
}

// DHCPAllAllocations returns a map of IPs currently allocated (statically and dynamically)
// in dnsmasq for a specific network. The returned map is keyed by a 16 byte array representing
// the net.IP format. The value of each map item is a DHCPAllocation struct containing at least
// whether the allocation was static or dynamic and optionally instance name or MAC address.
// MAC addresses are only included for dynamic IPv4 allocations (where name is not reliable).
// Static allocations are not overridden by dynamic allocations, allowing for instance name to be
// included for static IPv6 allocations. IPv6 addresses that are dynamically assigned cannot be
// reliably linked to instances using either name or MAC because dnsmasq does not record the MAC
// address for these records, and the recorded host name can be set by the instance if the dns.mode
// for the network is set to "dynamic" and so cannot be trusted, so in this case we do not return
// any identifying info.
func DHCPAllAllocations(network string) (map[[4]byte]DHCPAllocation, map[[16]byte]DHCPAllocation, error) {
	IPv4s := make(map[[4]byte]DHCPAllocation)
	IPv6s := make(map[[16]byte]DHCPAllocation)

	// First read all statically allocated IPs.
	files, err := ioutil.ReadDir(shared.VarPath("networks", network, "dnsmasq.hosts"))
	if err != nil && os.IsNotExist(err) {
		return nil, nil, err
	}

	for _, entry := range files {
		_, IPv4, IPv6, err := DHCPStaticAllocation(network, "", "", "", entry.Name())
		if err != nil {
			return nil, nil, err
		}

		if IPv4.IP != nil {
			var IPKey [4]byte
			copy(IPKey[:], IPv4.IP.To4())
			IPv4s[IPKey] = IPv4
		}

		if IPv6.IP != nil {
			var IPKey [16]byte
			copy(IPKey[:], IPv6.IP.To16())
			IPv6s[IPKey] = IPv6
		}
	}

	// Next read all dynamic allocated IPs.
	file, err := os.Open(shared.VarPath("networks", network, "dnsmasq.leases"))
	if err != nil {
		return nil, nil, err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())
		if len(fields) == 5 {
			IP := net.ParseIP(fields[2])
			if IP == nil {
				return nil, nil, fmt.Errorf("Error parsing IP address: %v", fields[2])
			}

			// Handle IPv6 addresses.
			if IP.To4() == nil {
				var IPKey [16]byte
				copy(IPKey[:], IP.To16())

				// Don't replace IPs from static config as more reliable.
				if IPv6s[IPKey].Name != "" {
					continue
				}

				IPv6s[IPKey] = DHCPAllocation{
					Static: false,
					IP:     IP.To16(),
				}
			} else {
				// MAC only available in IPv4 leases.
				MAC, err := net.ParseMAC(fields[1])
				if err != nil {
					return nil, nil, err
				}

				var IPKey [4]byte
				copy(IPKey[:], IP.To4())

				// Don't replace IPs from static config as more reliable.
				if IPv4s[IPKey].Name != "" {
					continue
				}

				IPv4s[IPKey] = DHCPAllocation{
					MAC:    MAC,
					Static: false,
					IP:     IP.To4(),
				}
			}
		}
	}
	if err := scanner.Err(); err != nil {
		return nil, nil, err
	}

	return IPv4s, IPv6s, nil
}

// StaticAllocationFileName returns the file name to use for a dnsmasq instance device static allocation.
func StaticAllocationFileName(projectName string, instanceName string, deviceName string) string {
	escapedDeviceName := filesystem.PathNameEncode(deviceName)

	return strings.Join([]string{project.Instance(projectName, instanceName), escapedDeviceName}, staticAllocationDeviceSeparator)
}
<|endoftext|>


<commit_before>package maas

import (
	"fmt"
	"net/url"
	"strings"

	"github.com/juju/gomaasapi"

	"github.com/lxc/lxd/lxd/project"
)

// Instance is a MAAS specific instance interface.
// This is used rather than instance.Instance to avoid import loops.
type Instance interface {
	Name() string
	Project() string
}

// Controller represents a MAAS server's machine functions
type Controller struct {
	url string

	srv     gomaasapi.Controller
	srvRaw  gomaasapi.Client
	machine gomaasapi.Machine
}

// ContainerInterface represents a MAAS connected network interface on the container
type ContainerInterface struct {
	Name       string
	MACAddress string
	Subnets    []ContainerInterfaceSubnet
}

// ContainerInterfaceSubnet represents an interface's subscription to a MAAS subnet
type ContainerInterfaceSubnet struct {
	Name    string
	Address string
}

func parseInterfaces(interfaces []ContainerInterface) (map[string]ContainerInterface, error) {
	// Quick checks.
	if len(interfaces) == 0 {
		return nil, fmt.Errorf("At least one interface must be provided")
	}

	// Parse the MACs and interfaces
	macInterfaces := map[string]ContainerInterface{}
	for _, iface := range interfaces {
		_, ok := macInterfaces[iface.MACAddress]
		if ok {
			return nil, fmt.Errorf("MAAS doesn't allow duplicate MAC addresses")
		}

		if iface.MACAddress == "" {
			return nil, fmt.Errorf("Interfaces must have a MAC address")
		}

		if len(iface.Subnets) == 0 {
			return nil, fmt.Errorf("Interfaces must have at least one subnet")
		}

		macInterfaces[iface.MACAddress] = iface
	}

	return macInterfaces, nil
}

// NewController returns a new Controller using the specific MAAS server and machine
func NewController(url string, key string, machine string) (*Controller, error) {
	baseURL := fmt.Sprintf("%s/api/2.0/", url)

	// Connect to MAAS
	srv, err := gomaasapi.NewController(gomaasapi.ControllerArgs{
		BaseURL: baseURL,
		APIKey:  key,
	})
	if err != nil {
		// Juju errors aren't user-friendly, try to extract what actually happened
		if !strings.Contains(err.Error(), "unsupported version") {
			return nil, err
		}

		return nil, fmt.Errorf("Unable to connect MAAS at '%s': %v", baseURL,
			strings.Split(strings.Split(err.Error(), "unsupported version: ")[1], " (")[0])
	}

	srvRaw, err := gomaasapi.NewAuthenticatedClient(baseURL, key)
	if err != nil {
		return nil, err
	}

	// Find the right machine
	machines, err := srv.Machines(gomaasapi.MachinesArgs{Hostnames: []string{machine}})
	if err != nil {
		return nil, err
	}

	if len(machines) != 1 {
		return nil, fmt.Errorf("Couldn't find the specified machine: %s", machine)
	}

	// Setup the struct
	c := Controller{}
	c.srv = srv
	c.srvRaw = *srvRaw
	c.machine = machines[0]
	c.url = baseURL

	return &c, err
}

func (c *Controller) getDomain(inst Instance) string {
	fields := strings.Split(c.machine.FQDN(), ".")
	domain := strings.Join(fields[1:], ".")

	if inst.Project() == project.Default {
		return domain
	}

	return fmt.Sprintf("%s.%s", inst.Project(), domain)
}

func (c *Controller) getDevice(name string, domain string) (gomaasapi.Device, error) {
	devs, err := c.machine.Devices(gomaasapi.DevicesArgs{
		Hostname: []string{name},
		Domain:   domain,
	})
	if err != nil {
		return nil, err
	}

	if len(devs) != 1 {
		return nil, fmt.Errorf("Couldn't find the specified instance: %s", name)
	}

	return devs[0], nil
}

func (c *Controller) getSubnets() (map[string]gomaasapi.Subnet, error) {
	// Get all the spaces
	spaces, err := c.srv.Spaces()
	if err != nil {
		return nil, err
	}

	// Get all the subnets
	subnets := map[string]gomaasapi.Subnet{}
	for _, space := range spaces {
		for _, subnet := range space.Subnets() {
			subnets[subnet.Name()] = subnet
		}
	}

	return subnets, nil
}

// CreateContainer defines a new MAAS device for the controller
func (c *Controller) CreateContainer(inst Instance, interfaces []ContainerInterface) error {
	// Parse the provided interfaces
	macInterfaces, err := parseInterfaces(interfaces)
	if err != nil {
		return err
	}

	// Get all the subnets
	subnets, err := c.getSubnets()
	if err != nil {
		return err
	}

	// Validation
	if len(interfaces) < 1 {
		return fmt.Errorf("Empty list of MAAS interface provided")
	}

	for _, iface := range interfaces {
		if len(iface.Subnets) < 1 {
			return fmt.Errorf("Bad subnet provided for interface '%s'", iface.Name)
		}

		for _, subnet := range iface.Subnets {
			_, ok := subnets[subnet.Name]
			if !ok {
				return fmt.Errorf("Subnet '%s' doesn't exist in MAAS", interfaces[0].Subnets[0].Name)
			}
		}
	}

	// Create the device and first interface
	device, err := c.machine.CreateDevice(gomaasapi.CreateMachineDeviceArgs{
		Hostname:      inst.Name(),
		Domain:        c.getDomain(inst),
		InterfaceName: interfaces[0].Name,
		MACAddress:    interfaces[0].MACAddress,
		VLAN:          subnets[interfaces[0].Subnets[0].Name].VLAN(),
	})
	if err != nil {
		return err
	}

	// Wipe the container entry if anything fails
	success := false
	defer func() {
		if success == true {
			return
		}

		_ = c.DeleteContainer(inst)
	}()

	// Create the rest of the interfaces
	for _, iface := range interfaces[1:] {
		_, err := device.CreateInterface(gomaasapi.CreateInterfaceArgs{
			Name:       iface.Name,
			MACAddress: iface.MACAddress,
			VLAN:       subnets[iface.Subnets[0].Name].VLAN(),
		})
		if err != nil {
			return err
		}
	}

	// Get a fresh copy of the device
	device, err = c.getDevice(inst.Name(), c.getDomain(inst))
	if err != nil {
		return err
	}

	// Setup the interfaces
	for _, entry := range device.InterfaceSet() {
		// Get our record
		iface, ok := macInterfaces[entry.MACAddress()]
		if !ok {
			return fmt.Errorf("MAAS created an interface with a bad MAC: %s", entry.MACAddress())
		}

		// Add the subnets
		for _, subnet := range iface.Subnets {
			err := entry.LinkSubnet(gomaasapi.LinkSubnetArgs{
				Mode:      gomaasapi.LinkModeStatic,
				Subnet:    subnets[subnet.Name],
				IPAddress: subnet.Address,
			})
			if err != nil {
				return err
			}
		}
	}

	success = true
	return nil
}

// DefinedContainer returns true if the container is defined in MAAS
func (c *Controller) DefinedContainer(inst Instance) (bool, error) {
	devs, err := c.machine.Devices(gomaasapi.DevicesArgs{
		Hostname: []string{inst.Name()},
		Domain:   c.getDomain(inst),
	})
	if err != nil {
		return false, err
	}

	if len(devs) == 1 {
		return true, nil
	}

	return false, nil
}

// UpdateContainer updates the MAAS device's interfaces with the new provided state
func (c *Controller) UpdateContainer(inst Instance, interfaces []ContainerInterface) error {
	// Parse the provided interfaces
	macInterfaces, err := parseInterfaces(interfaces)
	if err != nil {
		return err
	}

	// Get all the subnets
	subnets, err := c.getSubnets()
	if err != nil {
		return err
	}

	device, err := c.getDevice(inst.Name(), c.getDomain(inst))
	if err != nil {
		return err
	}

	// Validation
	if len(interfaces) < 1 {
		return fmt.Errorf("Empty list of MAAS interface provided")
	}

	for _, iface := range interfaces {
		if len(iface.Subnets) < 1 {
			return fmt.Errorf("Bad subnet provided for interface '%s'", iface.Name)
		}

		for _, subnet := range iface.Subnets {
			_, ok := subnets[subnet.Name]
			if !ok {
				return fmt.Errorf("Subnet '%s' doesn't exist in MAAS", interfaces[0].Subnets[0].Name)
			}
		}
	}

	// Iterate over existing interfaces, drop all removed ones and update existing ones
	existingInterfaces := map[string]gomaasapi.Interface{}
	for _, entry := range device.InterfaceSet() {
		// Check if the interface has been removed from the container
		iface, ok := macInterfaces[entry.MACAddress()]
		if !ok {
			// Delete the interface in MAAS
			err = entry.Delete()
			if err != nil {
				return err
			}

			continue
		}

		// Update the subnets
		existingSubnets := map[string]gomaasapi.Subnet{}
		for _, link := range entry.Links() {
			// Check if the MAAS subnet matches any of the container's
			found := false
			for _, subnet := range iface.Subnets {
				if subnet.Name == link.Subnet().Name() {
					if subnet.Address == "" || subnet.Address == link.IPAddress() {
						found = true
					}
					break
				}
			}

			// If no exact match could be found, remove it from MAAS
			if !found {
				err = entry.UnlinkSubnet(link.Subnet())
				if err != nil {
					return err
				}

				continue
			}

			// Record the existing up to date subnet
			existingSubnets[link.Subnet().Name()] = link.Subnet()
		}

		// Add any missing (or updated) subnet to MAAS
		for _, subnet := range iface.Subnets {
			// Check that it's not configured yet
			_, ok := existingSubnets[subnet.Name]
			if ok {
				continue
			}

			// Add the link
			err := entry.LinkSubnet(gomaasapi.LinkSubnetArgs{
				Mode:      gomaasapi.LinkModeStatic,
				Subnet:    subnets[subnet.Name],
				IPAddress: subnet.Address,
			})
			if err != nil {
				return err
			}
		}

		// Record the interface has being configured
		existingInterfaces[entry.MACAddress()] = entry
	}

	// Iterate over expected interfaces, add any missing one
	for _, iface := range macInterfaces {
		_, ok := existingInterfaces[iface.MACAddress]
		if ok {
			// We already have it so just move on
			continue
		}

		// Create the new interface
		entry, err := device.CreateInterface(gomaasapi.CreateInterfaceArgs{
			Name:       iface.Name,
			MACAddress: iface.MACAddress,
			VLAN:       subnets[iface.Subnets[0].Name].VLAN(),
		})
		if err != nil {
			return err
		}

		// Add the subnets
		for _, subnet := range iface.Subnets {
			err := entry.LinkSubnet(gomaasapi.LinkSubnetArgs{
				Mode:      gomaasapi.LinkModeStatic,
				Subnet:    subnets[subnet.Name],
				IPAddress: subnet.Address,
			})
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// RenameContainer renames the MAAS device for the container without releasing any allocation
func (c *Controller) RenameContainer(inst Instance, newName string) error {
	device, err := c.getDevice(inst.Name(), c.getDomain(inst))
	if err != nil {
		return err
	}

	// FIXME: We should convince the Juju folks to implement an Update() method on Device
	uri, err := url.Parse(fmt.Sprintf("%s/devices/%s/", c.url, device.SystemID()))
	if err != nil {
		return err
	}

	values := url.Values{}
	values.Set("hostname", newName)

	_, err = c.srvRaw.Put(uri, values)
	if err != nil {
		return err
	}

	return nil
}

// DeleteContainer removes the MAAS device for the container
func (c *Controller) DeleteContainer(inst Instance) error {
	device, err := c.getDevice(inst.Name(), c.getDomain(inst))
	if err != nil {
		return err
	}

	err = device.Delete()
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>lxd/maas: Fixes gosimple lint errors.<commit_after>package maas

import (
	"fmt"
	"net/url"
	"strings"

	"github.com/juju/gomaasapi"

	"github.com/lxc/lxd/lxd/project"
)

// Instance is a MAAS specific instance interface.
// This is used rather than instance.Instance to avoid import loops.
type Instance interface {
	Name() string
	Project() string
}

// Controller represents a MAAS server's machine functions
type Controller struct {
	url string

	srv     gomaasapi.Controller
	srvRaw  gomaasapi.Client
	machine gomaasapi.Machine
}

// ContainerInterface represents a MAAS connected network interface on the container
type ContainerInterface struct {
	Name       string
	MACAddress string
	Subnets    []ContainerInterfaceSubnet
}

// ContainerInterfaceSubnet represents an interface's subscription to a MAAS subnet
type ContainerInterfaceSubnet struct {
	Name    string
	Address string
}

func parseInterfaces(interfaces []ContainerInterface) (map[string]ContainerInterface, error) {
	// Quick checks.
	if len(interfaces) == 0 {
		return nil, fmt.Errorf("At least one interface must be provided")
	}

	// Parse the MACs and interfaces
	macInterfaces := map[string]ContainerInterface{}
	for _, iface := range interfaces {
		_, ok := macInterfaces[iface.MACAddress]
		if ok {
			return nil, fmt.Errorf("MAAS doesn't allow duplicate MAC addresses")
		}

		if iface.MACAddress == "" {
			return nil, fmt.Errorf("Interfaces must have a MAC address")
		}

		if len(iface.Subnets) == 0 {
			return nil, fmt.Errorf("Interfaces must have at least one subnet")
		}

		macInterfaces[iface.MACAddress] = iface
	}

	return macInterfaces, nil
}

// NewController returns a new Controller using the specific MAAS server and machine
func NewController(url string, key string, machine string) (*Controller, error) {
	baseURL := fmt.Sprintf("%s/api/2.0/", url)

	// Connect to MAAS
	srv, err := gomaasapi.NewController(gomaasapi.ControllerArgs{
		BaseURL: baseURL,
		APIKey:  key,
	})
	if err != nil {
		// Juju errors aren't user-friendly, try to extract what actually happened
		if !strings.Contains(err.Error(), "unsupported version") {
			return nil, err
		}

		return nil, fmt.Errorf("Unable to connect MAAS at '%s': %v", baseURL,
			strings.Split(strings.Split(err.Error(), "unsupported version: ")[1], " (")[0])
	}

	srvRaw, err := gomaasapi.NewAuthenticatedClient(baseURL, key)
	if err != nil {
		return nil, err
	}

	// Find the right machine
	machines, err := srv.Machines(gomaasapi.MachinesArgs{Hostnames: []string{machine}})
	if err != nil {
		return nil, err
	}

	if len(machines) != 1 {
		return nil, fmt.Errorf("Couldn't find the specified machine: %s", machine)
	}

	// Setup the struct
	c := Controller{}
	c.srv = srv
	c.srvRaw = *srvRaw
	c.machine = machines[0]
	c.url = baseURL

	return &c, err
}

func (c *Controller) getDomain(inst Instance) string {
	fields := strings.Split(c.machine.FQDN(), ".")
	domain := strings.Join(fields[1:], ".")

	if inst.Project() == project.Default {
		return domain
	}

	return fmt.Sprintf("%s.%s", inst.Project(), domain)
}

func (c *Controller) getDevice(name string, domain string) (gomaasapi.Device, error) {
	devs, err := c.machine.Devices(gomaasapi.DevicesArgs{
		Hostname: []string{name},
		Domain:   domain,
	})
	if err != nil {
		return nil, err
	}

	if len(devs) != 1 {
		return nil, fmt.Errorf("Couldn't find the specified instance: %s", name)
	}

	return devs[0], nil
}

func (c *Controller) getSubnets() (map[string]gomaasapi.Subnet, error) {
	// Get all the spaces
	spaces, err := c.srv.Spaces()
	if err != nil {
		return nil, err
	}

	// Get all the subnets
	subnets := map[string]gomaasapi.Subnet{}
	for _, space := range spaces {
		for _, subnet := range space.Subnets() {
			subnets[subnet.Name()] = subnet
		}
	}

	return subnets, nil
}

// CreateContainer defines a new MAAS device for the controller
func (c *Controller) CreateContainer(inst Instance, interfaces []ContainerInterface) error {
	// Parse the provided interfaces
	macInterfaces, err := parseInterfaces(interfaces)
	if err != nil {
		return err
	}

	// Get all the subnets
	subnets, err := c.getSubnets()
	if err != nil {
		return err
	}

	// Validation
	if len(interfaces) < 1 {
		return fmt.Errorf("Empty list of MAAS interface provided")
	}

	for _, iface := range interfaces {
		if len(iface.Subnets) < 1 {
			return fmt.Errorf("Bad subnet provided for interface '%s'", iface.Name)
		}

		for _, subnet := range iface.Subnets {
			_, ok := subnets[subnet.Name]
			if !ok {
				return fmt.Errorf("Subnet '%s' doesn't exist in MAAS", interfaces[0].Subnets[0].Name)
			}
		}
	}

	// Create the device and first interface
	device, err := c.machine.CreateDevice(gomaasapi.CreateMachineDeviceArgs{
		Hostname:      inst.Name(),
		Domain:        c.getDomain(inst),
		InterfaceName: interfaces[0].Name,
		MACAddress:    interfaces[0].MACAddress,
		VLAN:          subnets[interfaces[0].Subnets[0].Name].VLAN(),
	})
	if err != nil {
		return err
	}

	// Wipe the container entry if anything fails
	success := false
	defer func() {
		if success {
			return
		}

		_ = c.DeleteContainer(inst)
	}()

	// Create the rest of the interfaces
	for _, iface := range interfaces[1:] {
		_, err := device.CreateInterface(gomaasapi.CreateInterfaceArgs{
			Name:       iface.Name,
			MACAddress: iface.MACAddress,
			VLAN:       subnets[iface.Subnets[0].Name].VLAN(),
		})
		if err != nil {
			return err
		}
	}

	// Get a fresh copy of the device
	device, err = c.getDevice(inst.Name(), c.getDomain(inst))
	if err != nil {
		return err
	}

	// Setup the interfaces
	for _, entry := range device.InterfaceSet() {
		// Get our record
		iface, ok := macInterfaces[entry.MACAddress()]
		if !ok {
			return fmt.Errorf("MAAS created an interface with a bad MAC: %s", entry.MACAddress())
		}

		// Add the subnets
		for _, subnet := range iface.Subnets {
			err := entry.LinkSubnet(gomaasapi.LinkSubnetArgs{
				Mode:      gomaasapi.LinkModeStatic,
				Subnet:    subnets[subnet.Name],
				IPAddress: subnet.Address,
			})
			if err != nil {
				return err
			}
		}
	}

	success = true
	return nil
}

// DefinedContainer returns true if the container is defined in MAAS
func (c *Controller) DefinedContainer(inst Instance) (bool, error) {
	devs, err := c.machine.Devices(gomaasapi.DevicesArgs{
		Hostname: []string{inst.Name()},
		Domain:   c.getDomain(inst),
	})
	if err != nil {
		return false, err
	}

	if len(devs) == 1 {
		return true, nil
	}

	return false, nil
}

// UpdateContainer updates the MAAS device's interfaces with the new provided state
func (c *Controller) UpdateContainer(inst Instance, interfaces []ContainerInterface) error {
	// Parse the provided interfaces
	macInterfaces, err := parseInterfaces(interfaces)
	if err != nil {
		return err
	}

	// Get all the subnets
	subnets, err := c.getSubnets()
	if err != nil {
		return err
	}

	device, err := c.getDevice(inst.Name(), c.getDomain(inst))
	if err != nil {
		return err
	}

	// Validation
	if len(interfaces) < 1 {
		return fmt.Errorf("Empty list of MAAS interface provided")
	}

	for _, iface := range interfaces {
		if len(iface.Subnets) < 1 {
			return fmt.Errorf("Bad subnet provided for interface '%s'", iface.Name)
		}

		for _, subnet := range iface.Subnets {
			_, ok := subnets[subnet.Name]
			if !ok {
				return fmt.Errorf("Subnet '%s' doesn't exist in MAAS", interfaces[0].Subnets[0].Name)
			}
		}
	}

	// Iterate over existing interfaces, drop all removed ones and update existing ones
	existingInterfaces := map[string]gomaasapi.Interface{}
	for _, entry := range device.InterfaceSet() {
		// Check if the interface has been removed from the container
		iface, ok := macInterfaces[entry.MACAddress()]
		if !ok {
			// Delete the interface in MAAS
			err = entry.Delete()
			if err != nil {
				return err
			}

			continue
		}

		// Update the subnets
		existingSubnets := map[string]gomaasapi.Subnet{}
		for _, link := range entry.Links() {
			// Check if the MAAS subnet matches any of the container's
			found := false
			for _, subnet := range iface.Subnets {
				if subnet.Name == link.Subnet().Name() {
					if subnet.Address == "" || subnet.Address == link.IPAddress() {
						found = true
					}
					break
				}
			}

			// If no exact match could be found, remove it from MAAS
			if !found {
				err = entry.UnlinkSubnet(link.Subnet())
				if err != nil {
					return err
				}

				continue
			}

			// Record the existing up to date subnet
			existingSubnets[link.Subnet().Name()] = link.Subnet()
		}

		// Add any missing (or updated) subnet to MAAS
		for _, subnet := range iface.Subnets {
			// Check that it's not configured yet
			_, ok := existingSubnets[subnet.Name]
			if ok {
				continue
			}

			// Add the link
			err := entry.LinkSubnet(gomaasapi.LinkSubnetArgs{
				Mode:      gomaasapi.LinkModeStatic,
				Subnet:    subnets[subnet.Name],
				IPAddress: subnet.Address,
			})
			if err != nil {
				return err
			}
		}

		// Record the interface has being configured
		existingInterfaces[entry.MACAddress()] = entry
	}

	// Iterate over expected interfaces, add any missing one
	for _, iface := range macInterfaces {
		_, ok := existingInterfaces[iface.MACAddress]
		if ok {
			// We already have it so just move on
			continue
		}

		// Create the new interface
		entry, err := device.CreateInterface(gomaasapi.CreateInterfaceArgs{
			Name:       iface.Name,
			MACAddress: iface.MACAddress,
			VLAN:       subnets[iface.Subnets[0].Name].VLAN(),
		})
		if err != nil {
			return err
		}

		// Add the subnets
		for _, subnet := range iface.Subnets {
			err := entry.LinkSubnet(gomaasapi.LinkSubnetArgs{
				Mode:      gomaasapi.LinkModeStatic,
				Subnet:    subnets[subnet.Name],
				IPAddress: subnet.Address,
			})
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// RenameContainer renames the MAAS device for the container without releasing any allocation
func (c *Controller) RenameContainer(inst Instance, newName string) error {
	device, err := c.getDevice(inst.Name(), c.getDomain(inst))
	if err != nil {
		return err
	}

	// FIXME: We should convince the Juju folks to implement an Update() method on Device
	uri, err := url.Parse(fmt.Sprintf("%s/devices/%s/", c.url, device.SystemID()))
	if err != nil {
		return err
	}

	values := url.Values{}
	values.Set("hostname", newName)

	_, err = c.srvRaw.Put(uri, values)
	if err != nil {
		return err
	}

	return nil
}

// DeleteContainer removes the MAAS device for the container
func (c *Controller) DeleteContainer(inst Instance) error {
	device, err := c.getDevice(inst.Name(), c.getDomain(inst))
	if err != nil {
		return err
	}

	err = device.Delete()
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package migration

// GetRsyncFeaturesSlice returns a slice of strings representing the supported RSYNC features
func (m *MigrationHeader) GetRsyncFeaturesSlice() []string {
	features := []string{}
	if m == nil {
		return features
	}
	if m.RsyncFeatures != nil {
		if m.RsyncFeatures.Xattrs != nil && *m.RsyncFeatures.Xattrs == true {
			features = append(features, "xattrs")
		}

		if m.RsyncFeatures.Delete != nil && *m.RsyncFeatures.Delete == true {
			features = append(features, "delete")
		}

		if m.RsyncFeatures.Compress != nil && *m.RsyncFeatures.Compress == true {
			features = append(features, "compress")
		}

		if m.RsyncFeatures.Bidirectional != nil && *m.RsyncFeatures.Bidirectional == true {
			features = append(features, "bidirectional")
		}
	}

	return features
}

// GetZfsFeaturesSlice returns a slice of strings representing the supported ZFS features
func (m *MigrationHeader) GetZfsFeaturesSlice() []string {
	features := []string{}
	if m == nil {
		return features
	}

	if m.ZfsFeatures != nil {
		if m.ZfsFeatures.Compress != nil && *m.ZfsFeatures.Compress == true {
			features = append(features, "compress")
		}
	}

	return features
}
<commit_msg>lxd/migration/utils: Adds GetBtrfsFeaturesSlice function<commit_after>package migration

// BTRFSFeatureMigrationHeader indicates a migration header will be sent/recv in data channel first.
const BTRFSFeatureMigrationHeader = "migration_header"

// BTRFSFeatureSubvolumes indicates migration can send/recv subvolumes.
const BTRFSFeatureSubvolumes = "header_subvolumes"

// GetRsyncFeaturesSlice returns a slice of strings representing the supported RSYNC features
func (m *MigrationHeader) GetRsyncFeaturesSlice() []string {
	features := []string{}
	if m == nil {
		return features
	}
	if m.RsyncFeatures != nil {
		if m.RsyncFeatures.Xattrs != nil && *m.RsyncFeatures.Xattrs == true {
			features = append(features, "xattrs")
		}

		if m.RsyncFeatures.Delete != nil && *m.RsyncFeatures.Delete == true {
			features = append(features, "delete")
		}

		if m.RsyncFeatures.Compress != nil && *m.RsyncFeatures.Compress == true {
			features = append(features, "compress")
		}

		if m.RsyncFeatures.Bidirectional != nil && *m.RsyncFeatures.Bidirectional == true {
			features = append(features, "bidirectional")
		}
	}

	return features
}

// GetZfsFeaturesSlice returns a slice of strings representing the supported ZFS features
func (m *MigrationHeader) GetZfsFeaturesSlice() []string {
	features := []string{}
	if m == nil {
		return features
	}

	if m.ZfsFeatures != nil {
		if m.ZfsFeatures.Compress != nil && *m.ZfsFeatures.Compress == true {
			features = append(features, "compress")
		}
	}

	return features
}

// GetBtrfsFeaturesSlice returns a slice of strings representing the supported BTRFS features.
func (m *MigrationHeader) GetBtrfsFeaturesSlice() []string {
	features := []string{}
	if m == nil {
		return features
	}

	if m.BtrfsFeatures != nil {
		if m.BtrfsFeatures.MigrationHeader != nil && *m.BtrfsFeatures.MigrationHeader == true {
			features = append(features, BTRFSFeatureMigrationHeader)
		}

		if m.BtrfsFeatures.HeaderSubvolumes != nil && *m.BtrfsFeatures.HeaderSubvolumes == true {
			features = append(features, BTRFSFeatureSubvolumes)
		}
	}

	return features
}
<|endoftext|>


<commit_before>package migration

// BTRFSFeatureMigrationHeader indicates a migration header will be sent/recv in data channel after index header.
const BTRFSFeatureMigrationHeader = "migration_header"

// BTRFSFeatureSubvolumes indicates migration can send/recv subvolumes.
const BTRFSFeatureSubvolumes = "header_subvolumes"

// BTRFSFeatureSubvolumeUUIDs indicates that the header will include subvolume UUIDs.
const BTRFSFeatureSubvolumeUUIDs = "header_subvolume_uuids"

// ZFSFeatureMigrationHeader indicates a migration header will be sent/recv in data channel after index header.
const ZFSFeatureMigrationHeader = "migration_header"

// GetRsyncFeaturesSlice returns a slice of strings representing the supported RSYNC features
func (m *MigrationHeader) GetRsyncFeaturesSlice() []string {
	features := []string{}
	if m == nil {
		return features
	}
	if m.RsyncFeatures != nil {
		if m.RsyncFeatures.Xattrs != nil && *m.RsyncFeatures.Xattrs == true {
			features = append(features, "xattrs")
		}

		if m.RsyncFeatures.Delete != nil && *m.RsyncFeatures.Delete == true {
			features = append(features, "delete")
		}

		if m.RsyncFeatures.Compress != nil && *m.RsyncFeatures.Compress == true {
			features = append(features, "compress")
		}

		if m.RsyncFeatures.Bidirectional != nil && *m.RsyncFeatures.Bidirectional == true {
			features = append(features, "bidirectional")
		}
	}

	return features
}

// GetZfsFeaturesSlice returns a slice of strings representing the supported ZFS features
func (m *MigrationHeader) GetZfsFeaturesSlice() []string {
	features := []string{}
	if m == nil {
		return features
	}

	if m.ZfsFeatures != nil {
		if m.ZfsFeatures.Compress != nil && *m.ZfsFeatures.Compress == true {
			features = append(features, "compress")
		}

		if m.ZfsFeatures.MigrationHeader != nil && *m.ZfsFeatures.MigrationHeader == true {
			features = append(features, ZFSFeatureMigrationHeader)
		}
	}

	return features
}

// GetBtrfsFeaturesSlice returns a slice of strings representing the supported BTRFS features.
func (m *MigrationHeader) GetBtrfsFeaturesSlice() []string {
	features := []string{}
	if m == nil {
		return features
	}

	if m.BtrfsFeatures != nil {
		if m.BtrfsFeatures.MigrationHeader != nil && *m.BtrfsFeatures.MigrationHeader == true {
			features = append(features, BTRFSFeatureMigrationHeader)
		}

		if m.BtrfsFeatures.HeaderSubvolumes != nil && *m.BtrfsFeatures.HeaderSubvolumes == true {
			features = append(features, BTRFSFeatureSubvolumes)
		}

		if m.BtrfsFeatures.HeaderSubvolumeUuids != nil && *m.BtrfsFeatures.HeaderSubvolumeUuids == true {
			features = append(features, BTRFSFeatureSubvolumeUUIDs)
		}
	}

	return features
}
<commit_msg>lxd/migration/utils: Adds IndexHeaderVersion constant<commit_after>package migration

// IndexHeaderVersion version of the index header to be sent/recv.
const IndexHeaderVersion uint32 = 1

// BTRFSFeatureMigrationHeader indicates a migration header will be sent/recv in data channel after index header.
const BTRFSFeatureMigrationHeader = "migration_header"

// BTRFSFeatureSubvolumes indicates migration can send/recv subvolumes.
const BTRFSFeatureSubvolumes = "header_subvolumes"

// BTRFSFeatureSubvolumeUUIDs indicates that the header will include subvolume UUIDs.
const BTRFSFeatureSubvolumeUUIDs = "header_subvolume_uuids"

// ZFSFeatureMigrationHeader indicates a migration header will be sent/recv in data channel after index header.
const ZFSFeatureMigrationHeader = "migration_header"

// GetRsyncFeaturesSlice returns a slice of strings representing the supported RSYNC features
func (m *MigrationHeader) GetRsyncFeaturesSlice() []string {
	features := []string{}
	if m == nil {
		return features
	}
	if m.RsyncFeatures != nil {
		if m.RsyncFeatures.Xattrs != nil && *m.RsyncFeatures.Xattrs == true {
			features = append(features, "xattrs")
		}

		if m.RsyncFeatures.Delete != nil && *m.RsyncFeatures.Delete == true {
			features = append(features, "delete")
		}

		if m.RsyncFeatures.Compress != nil && *m.RsyncFeatures.Compress == true {
			features = append(features, "compress")
		}

		if m.RsyncFeatures.Bidirectional != nil && *m.RsyncFeatures.Bidirectional == true {
			features = append(features, "bidirectional")
		}
	}

	return features
}

// GetZfsFeaturesSlice returns a slice of strings representing the supported ZFS features
func (m *MigrationHeader) GetZfsFeaturesSlice() []string {
	features := []string{}
	if m == nil {
		return features
	}

	if m.ZfsFeatures != nil {
		if m.ZfsFeatures.Compress != nil && *m.ZfsFeatures.Compress == true {
			features = append(features, "compress")
		}

		if m.ZfsFeatures.MigrationHeader != nil && *m.ZfsFeatures.MigrationHeader == true {
			features = append(features, ZFSFeatureMigrationHeader)
		}
	}

	return features
}

// GetBtrfsFeaturesSlice returns a slice of strings representing the supported BTRFS features.
func (m *MigrationHeader) GetBtrfsFeaturesSlice() []string {
	features := []string{}
	if m == nil {
		return features
	}

	if m.BtrfsFeatures != nil {
		if m.BtrfsFeatures.MigrationHeader != nil && *m.BtrfsFeatures.MigrationHeader == true {
			features = append(features, BTRFSFeatureMigrationHeader)
		}

		if m.BtrfsFeatures.HeaderSubvolumes != nil && *m.BtrfsFeatures.HeaderSubvolumes == true {
			features = append(features, BTRFSFeatureSubvolumes)
		}

		if m.BtrfsFeatures.HeaderSubvolumeUuids != nil && *m.BtrfsFeatures.HeaderSubvolumeUuids == true {
			features = append(features, BTRFSFeatureSubvolumeUUIDs)
		}
	}

	return features
}
<|endoftext|>


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

package ebiten

import (
	"github.com/hajimehoshi/ebiten/internal/graphics/opengl"
)

// Filter represents the type of filter to be used when an image is maginified or minified.
type Filter int

// Filters
const (
	FilterNearest Filter = iota
	FilterLinear
)

func glFilter(c *opengl.Context, filter Filter) opengl.Filter {
	switch filter {
	case FilterNearest:
		return c.Nearest
	case FilterLinear:
		return c.Linear
	}
	panic("not reach")
}
<commit_msg>graphics: Add comment<commit_after>// Copyright 2014 Hajime Hoshi
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ebiten

import (
	"github.com/hajimehoshi/ebiten/internal/graphics/opengl"
)

// Filter represents the type of filter to be used when an image is maginified or minified.
type Filter int

const (
	FilterNearest Filter = iota // nearest (crisp-edged) filter
	FilterLinear                // linear filter
)

func glFilter(c *opengl.Context, filter Filter) opengl.Filter {
	switch filter {
	case FilterNearest:
		return c.Nearest
	case FilterLinear:
		return c.Linear
	}
	panic("not reach")
}
<|endoftext|>


<commit_before>package urxvtermbox

import (
	"fmt"
	"os"
	"os/exec"
	"time"

	"github.com/kr/pty"
	"github.com/mpasternacki/termbox-go"
)

func UrxvtPty(args ...string) (*os.File, func(), error) {
	master, slave, err := pty.Open()
	if err != nil {
		return nil, nil, err
	}

	cmd := exec.Command("urxvt", append([]string{"-pty-fd", "3"}, args...)...)
	fini := func() {
		slave.Close()
		if cmd.Process != nil {
			cmd.Wait()
		}
	}

	cmd.ExtraFiles = []*os.File{master}
	cmd.Start()

	// poll until urxvt is done starting
	for {
		rows, _, err := pty.Getsize(slave)
		if err != nil {
			fini()
			return nil, nil, err
		}
		if rows > 0 {
			break
		}
		time.Sleep(5 * time.Millisecond)
	}

	return slave, fini, nil
}

func TermboxUrxvt(width, height int, args ...string) (func(), error) {
	if width == 0 {
		width = 80
	}
	if height == 0 {
		height = 25
	}

	pty, fini, err := UrxvtPty(append(
		[]string{"-geometry", fmt.Sprintf("%dx%d", width, height)},
		args...)...)

	if err != nil {
		return nil, err
	}

	origTerminalDevice := termbox.TerminalDevice
	termbox.TerminalDevice = pty.Name()
	return func() {
		termbox.TerminalDevice = origTerminalDevice
		fini()
	}, nil
}
<commit_msg>urxvtermbox: disable scroll bar<commit_after>package urxvtermbox

import (
	"fmt"
	"os"
	"os/exec"
	"time"

	"github.com/kr/pty"
	"github.com/mpasternacki/termbox-go"
)

func UrxvtPty(args ...string) (*os.File, func(), error) {
	master, slave, err := pty.Open()
	if err != nil {
		return nil, nil, err
	}

	cmd := exec.Command("urxvt", append([]string{"-pty-fd", "3"}, args...)...)
	fini := func() {
		slave.Close()
		if cmd.Process != nil {
			cmd.Wait()
		}
	}

	cmd.ExtraFiles = []*os.File{master}
	cmd.Start()

	// poll until urxvt is done starting
	for {
		rows, _, err := pty.Getsize(slave)
		if err != nil {
			fini()
			return nil, nil, err
		}
		if rows > 0 {
			break
		}
		time.Sleep(5 * time.Millisecond)
	}

	return slave, fini, nil
}

func TermboxUrxvt(width, height int, args ...string) (func(), error) {
	if width == 0 {
		width = 80
	}
	if height == 0 {
		height = 25
	}

	pty, fini, err := UrxvtPty(append(
		[]string{"+sb", "-geometry", fmt.Sprintf("%dx%d", width, height)},
		args...)...)

	if err != nil {
		return nil, err
	}

	origTerminalDevice := termbox.TerminalDevice
	termbox.TerminalDevice = pty.Name()
	return func() {
		termbox.TerminalDevice = origTerminalDevice
		fini()
	}, nil
}
<|endoftext|>


<commit_before>package graphite

import (
	"bytes"
	"fmt"
	"log"
	"net"
	"time"
)

// Graphite is a struct that defines the relevant properties of a graphite
// connection
type Graphite struct {
	Host     string
	Port     int
	Protocol string
	Timeout  time.Duration
	Prefix   string
	conn     net.Conn
	nop      bool
}

// defaultTimeout is the default number of seconds that we're willing to wait
// before forcing the connection establishment to fail
const defaultTimeout = 5

// IsNop is a getter for *graphite.Graphite.nop
func (graphite *Graphite) IsNop() bool {
	if graphite.nop {
		return true
	}
	return false
}

// Given a Graphite struct, Connect populates the Graphite.conn field with an
// appropriate TCP connection
func (graphite *Graphite) Connect() error {
	if !graphite.IsNop() {
		if graphite.conn != nil {
			graphite.conn.Close()
		}

		address := fmt.Sprintf("%s:%d", graphite.Host, graphite.Port)

		if graphite.Timeout == 0 {
			graphite.Timeout = defaultTimeout * time.Second
		}

		var err error
		var conn net.Conn

		if graphite.Protocol == "udp" {
			udpAddr, err := net.ResolveUDPAddr("udp", address)
			if err != nil {
				return err
			}
			conn, err = net.DialUDP(graphite.Protocol, nil, udpAddr)
		} else {
			conn, err = net.DialTimeout(graphite.Protocol, address, graphite.Timeout)
		}

		if err != nil {
			return err
		}

		graphite.conn = conn
	}

	return nil
}

// Given a Graphite struct, Disconnect closes the Graphite.conn field
func (graphite *Graphite) Disconnect() error {
	err := graphite.conn.Close()
	graphite.conn = nil
	return err
}

// Given a Metric struct, the SendMetric method sends the supplied metric to the
// Graphite connection that the method is called upon
func (graphite *Graphite) SendMetric(metric Metric) error {
	metrics := make([]Metric, 1)
	metrics[0] = metric

	return graphite.sendMetrics(metrics)
}

// Given a slice of Metrics, the SendMetrics method sends the metrics, as a
// batch, to the Graphite connection that the method is called upon
func (graphite *Graphite) SendMetrics(metrics []Metric) error {
	return graphite.sendMetrics(metrics)
}

// sendMetrics is an internal function that is used to write to the TCP
// connection in order to communicate metrics to the remote Graphite host
func (graphite *Graphite) sendMetrics(metrics []Metric) error {
	if graphite.IsNop() {
		for _, metric := range metrics {
			log.Printf("Graphite: %s\n", metric)
		}
		return nil
	}
	zeroed_metric := Metric{} // ignore unintialized metrics
	buf := bytes.NewBufferString("")
	for _, metric := range metrics {
		if metric == zeroed_metric {
			continue // ignore unintialized metrics
		}
		if metric.Timestamp == 0 {
			metric.Timestamp = time.Now().Unix()
		}
		metric_name := ""
		if graphite.Prefix != "" {
			metric_name = fmt.Sprintf("%s.%s", graphite.Prefix, metric.Name)
		} else {
			metric_name = metric.Name
		}
		if graphite.Protocol == "udp" {
			fmt.Fprintf(graphite.conn, "%s %s %d\n", metric_name, metric.Value, metric.Timestamp)
			continue
		}
		buf.WriteString(fmt.Sprintf("%s %s %d\n", metric_name, metric.Value, metric.Timestamp))
	}
	if graphite.Protocol == "tcp" {
		_, err := graphite.conn.Write(buf.Bytes())
		//fmt.Print("Sent msg:", buf.String(), "'")
		if err != nil {
			return err
		}
	}
	return nil
}

// The SimpleSend method can be used to just pass a metric name and value and
// have it be sent to the Graphite host with the current timestamp
func (graphite *Graphite) SimpleSend(stat string, value string) error {
	metrics := make([]Metric, 1)
	metrics[0] = NewMetric(stat, value, time.Now().Unix())
	err := graphite.sendMetrics(metrics)
	if err != nil {
		return err
	}
	return nil
}

// NewGraphite is a factory method that's used to create a new Graphite
func NewGraphite(host string, port int) (*Graphite, error) {
	return GraphiteFactory("tcp", host, port, "")
}

// NewGraphiteWithMetricPrefix is a factory method that's used to create a new Graphite with a metric prefix
func NewGraphiteWithMetricPrefix(host string, port int, prefix string) (*Graphite, error) {
	return GraphiteFactory("tcp", host, port, prefix)
}

// When a UDP connection to Graphite is required
func NewGraphiteUDP(host string, port int) (*Graphite, error) {
	return GraphiteFactory("udp", host, port, "")
}

// NewGraphiteNop is a factory method that returns a Graphite struct but will
// not actually try to send any packets to a remote host and, instead, will just
// log. This is useful if you want to use Graphite in a project but don't want
// to make Graphite a requirement for the project.
func NewGraphiteNop(host string, port int) *Graphite {
	graphiteNop, _ := GraphiteFactory("nop", host, port, "")
	return graphiteNop
}

func GraphiteFactory(protocol string, host string, port int, prefix string) (*Graphite, error) {
	var graphite *Graphite

	switch protocol {
	case "tcp":
		graphite = &Graphite{Host: host, Port: port, Protocol: "tcp", Prefix: prefix}
	case "udp":
		graphite = &Graphite{Host: host, Port: port, Protocol: "udp", Prefix: prefix}
	case "nop":
		graphite = &Graphite{Host: host, Port: port, nop: true}
	}

	err := graphite.Connect()
	if err != nil {
		return nil, err
	}

	return graphite, nil
}
<commit_msg>#1: Add Graphite.DisableLog<commit_after>package graphite

import (
	"bytes"
	"fmt"
	"log"
	"net"
	"time"
)

// Graphite is a struct that defines the relevant properties of a graphite
// connection
type Graphite struct {
	Host     string
	Port     int
	Protocol string
	Timeout  time.Duration
	Prefix   string
	conn     net.Conn
	nop      bool
	DisableLog bool
}

// defaultTimeout is the default number of seconds that we're willing to wait
// before forcing the connection establishment to fail
const defaultTimeout = 5

// IsNop is a getter for *graphite.Graphite.nop
func (graphite *Graphite) IsNop() bool {
	if graphite.nop {
		return true
	}
	return false
}

// Given a Graphite struct, Connect populates the Graphite.conn field with an
// appropriate TCP connection
func (graphite *Graphite) Connect() error {
	if !graphite.IsNop() {
		if graphite.conn != nil {
			graphite.conn.Close()
		}

		address := fmt.Sprintf("%s:%d", graphite.Host, graphite.Port)

		if graphite.Timeout == 0 {
			graphite.Timeout = defaultTimeout * time.Second
		}

		var err error
		var conn net.Conn

		if graphite.Protocol == "udp" {
			udpAddr, err := net.ResolveUDPAddr("udp", address)
			if err != nil {
				return err
			}
			conn, err = net.DialUDP(graphite.Protocol, nil, udpAddr)
		} else {
			conn, err = net.DialTimeout(graphite.Protocol, address, graphite.Timeout)
		}

		if err != nil {
			return err
		}

		graphite.conn = conn
	}

	return nil
}

// Given a Graphite struct, Disconnect closes the Graphite.conn field
func (graphite *Graphite) Disconnect() error {
	err := graphite.conn.Close()
	graphite.conn = nil
	return err
}

// Given a Metric struct, the SendMetric method sends the supplied metric to the
// Graphite connection that the method is called upon
func (graphite *Graphite) SendMetric(metric Metric) error {
	metrics := make([]Metric, 1)
	metrics[0] = metric

	return graphite.sendMetrics(metrics)
}

// Given a slice of Metrics, the SendMetrics method sends the metrics, as a
// batch, to the Graphite connection that the method is called upon
func (graphite *Graphite) SendMetrics(metrics []Metric) error {
	return graphite.sendMetrics(metrics)
}

// sendMetrics is an internal function that is used to write to the TCP
// connection in order to communicate metrics to the remote Graphite host
func (graphite *Graphite) sendMetrics(metrics []Metric) error {
	if graphite.IsNop() {
		if !graphite.DisableLog {
			for _, metric := range metrics {
				log.Printf("Graphite: %s\n", metric)
			}
		}
		return nil
	}
	zeroed_metric := Metric{} // ignore unintialized metrics
	buf := bytes.NewBufferString("")
	for _, metric := range metrics {
		if metric == zeroed_metric {
			continue // ignore unintialized metrics
		}
		if metric.Timestamp == 0 {
			metric.Timestamp = time.Now().Unix()
		}
		metric_name := ""
		if graphite.Prefix != "" {
			metric_name = fmt.Sprintf("%s.%s", graphite.Prefix, metric.Name)
		} else {
			metric_name = metric.Name
		}
		if graphite.Protocol == "udp" {
			fmt.Fprintf(graphite.conn, "%s %s %d\n", metric_name, metric.Value, metric.Timestamp)
			continue
		}
		buf.WriteString(fmt.Sprintf("%s %s %d\n", metric_name, metric.Value, metric.Timestamp))
	}
	if graphite.Protocol == "tcp" {
		_, err := graphite.conn.Write(buf.Bytes())
		//fmt.Print("Sent msg:", buf.String(), "'")
		if err != nil {
			return err
		}
	}
	return nil
}

// The SimpleSend method can be used to just pass a metric name and value and
// have it be sent to the Graphite host with the current timestamp
func (graphite *Graphite) SimpleSend(stat string, value string) error {
	metrics := make([]Metric, 1)
	metrics[0] = NewMetric(stat, value, time.Now().Unix())
	err := graphite.sendMetrics(metrics)
	if err != nil {
		return err
	}
	return nil
}

// NewGraphite is a factory method that's used to create a new Graphite
func NewGraphite(host string, port int) (*Graphite, error) {
	return GraphiteFactory("tcp", host, port, "")
}

// NewGraphiteWithMetricPrefix is a factory method that's used to create a new Graphite with a metric prefix
func NewGraphiteWithMetricPrefix(host string, port int, prefix string) (*Graphite, error) {
	return GraphiteFactory("tcp", host, port, prefix)
}

// When a UDP connection to Graphite is required
func NewGraphiteUDP(host string, port int) (*Graphite, error) {
	return GraphiteFactory("udp", host, port, "")
}

// NewGraphiteNop is a factory method that returns a Graphite struct but will
// not actually try to send any packets to a remote host and, instead, will just
// log. This is useful if you want to use Graphite in a project but don't want
// to make Graphite a requirement for the project.
func NewGraphiteNop(host string, port int) *Graphite {
	graphiteNop, _ := GraphiteFactory("nop", host, port, "")
	return graphiteNop
}

func GraphiteFactory(protocol string, host string, port int, prefix string) (*Graphite, error) {
	var graphite *Graphite

	switch protocol {
	case "tcp":
		graphite = &Graphite{Host: host, Port: port, Protocol: "tcp", Prefix: prefix}
	case "udp":
		graphite = &Graphite{Host: host, Port: port, Protocol: "udp", Prefix: prefix}
	case "nop":
		graphite = &Graphite{Host: host, Port: port, nop: true}
	}

	err := graphite.Connect()
	if err != nil {
		return nil, err
	}

	return graphite, nil
}
<|endoftext|>