<commit_before>// Copyright 2012 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

package layers

import (
	"fmt"
	"strconv"

	"github.com/google/gopacket"
)

// TCPPort is a port in a TCP layer.
type TCPPort uint16

// UDPPort is a port in a UDP layer.
type UDPPort uint16

// RUDPPort is a port in a RUDP layer.
type RUDPPort uint8

// SCTPPort is a port in a SCTP layer.
type SCTPPort uint16

// UDPLitePort is a port in a UDPLite layer.
type UDPLitePort uint16

// RUDPPortNames contains the string names for all RUDP ports.
var RUDPPortNames = map[RUDPPort]string{}

// UDPLitePortNames contains the string names for all UDPLite ports.
var UDPLitePortNames = map[UDPLitePort]string{}

// {TCP,UDP,SCTP}PortNames can be found in iana_ports.go

// String returns the port as "number(name)" if there's a well-known port name,
// or just "number" if there isn't.  Well-known names are stored in
// TCPPortNames.
func (a TCPPort) String() string {
	if name, ok := TCPPortNames[a]; ok {
		return fmt.Sprintf("%d(%s)", a, name)
	}
	return strconv.Itoa(int(a))
}

// String returns the port as "number(name)" if there's a well-known port name,
// or just "number" if there isn't.  Well-known names are stored in
// UDPPortNames.
func (a UDPPort) String() string {
	if name, ok := UDPPortNames[a]; ok {
		return fmt.Sprintf("%d(%s)", a, name)
	}
	return strconv.Itoa(int(a))
}

// LayerType returns a LayerType that would be able to decode the
// application payload. It use some well-known port such as 53 for DNS.
//
// Returns gopacket.LayerTypePayload for unknown/unsupported port numbers.
func (a UDPPort) LayerType() gopacket.LayerType {
	lt := udpPortLayerType[uint16(a)]
	if lt != 0 {
		return lt
	}
	return gopacket.LayerTypePayload
}

var udpPortLayerType [65536]gopacket.LayerType

// RegisterUDPPortLayerType create a new mapping between an UDPPort
// and an underlaying LayerType.
func RegisterUDPPortLayerType(portNum uint16, layerType gopacket.LayerType) uint16 {
	udpPortLayerType[portNum] = layerType
	return portNum
}

// String returns the port as "number(name)" if there's a well-known port name,
// or just "number" if there isn't.  Well-known names are stored in
// RUDPPortNames.
func (a RUDPPort) String() string {
	if name, ok := RUDPPortNames[a]; ok {
		return fmt.Sprintf("%d(%s)", a, name)
	}
	return strconv.Itoa(int(a))
}

// String returns the port as "number(name)" if there's a well-known port name,
// or just "number" if there isn't.  Well-known names are stored in
// SCTPPortNames.
func (a SCTPPort) String() string {
	if name, ok := SCTPPortNames[a]; ok {
		return fmt.Sprintf("%d(%s)", a, name)
	}
	return strconv.Itoa(int(a))
}

// String returns the port as "number(name)" if there's a well-known port name,
// or just "number" if there isn't.  Well-known names are stored in
// UDPLitePortNames.
func (a UDPLitePort) String() string {
	if name, ok := UDPLitePortNames[a]; ok {
		return fmt.Sprintf("%d(%s)", a, name)
	}
	return strconv.Itoa(int(a))
}

func init() {
	RegisterUDPPortLayerType(53, LayerTypeDNS)
	RegisterUDPPortLayerType(123, LayerTypeNTP)
	RegisterUDPPortLayerType(4789, LayerTypeVXLAN)
	RegisterUDPPortLayerType(67, LayerTypeDHCPv4)
	RegisterUDPPortLayerType(68, LayerTypeDHCPv4)
	RegisterUDPPortLayerType(6343, LayerTypeSFlow)
}
<commit_msg>Improve UDPPort LayerType mapping code<commit_after>// Copyright 2012 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

package layers

import (
	"fmt"
	"strconv"

	"github.com/google/gopacket"
)

// TCPPort is a port in a TCP layer.
type TCPPort uint16

// UDPPort is a port in a UDP layer.
type UDPPort uint16

// RUDPPort is a port in a RUDP layer.
type RUDPPort uint8

// SCTPPort is a port in a SCTP layer.
type SCTPPort uint16

// UDPLitePort is a port in a UDPLite layer.
type UDPLitePort uint16

// RUDPPortNames contains the string names for all RUDP ports.
var RUDPPortNames = map[RUDPPort]string{}

// UDPLitePortNames contains the string names for all UDPLite ports.
var UDPLitePortNames = map[UDPLitePort]string{}

// {TCP,UDP,SCTP}PortNames can be found in iana_ports.go

// String returns the port as "number(name)" if there's a well-known port name,
// or just "number" if there isn't.  Well-known names are stored in
// TCPPortNames.
func (a TCPPort) String() string {
	if name, ok := TCPPortNames[a]; ok {
		return fmt.Sprintf("%d(%s)", a, name)
	}
	return strconv.Itoa(int(a))
}

// String returns the port as "number(name)" if there's a well-known port name,
// or just "number" if there isn't.  Well-known names are stored in
// UDPPortNames.
func (a UDPPort) String() string {
	if name, ok := UDPPortNames[a]; ok {
		return fmt.Sprintf("%d(%s)", a, name)
	}
	return strconv.Itoa(int(a))
}

// LayerType returns a LayerType that would be able to decode the
// application payload. It use some well-known port such as 53 for DNS.
//
// Returns gopacket.LayerTypePayload for unknown/unsupported port numbers.
func (a UDPPort) LayerType() gopacket.LayerType {
	lt := udpPortLayerType[uint16(a)]
	if lt != 0 {
		return lt
	}
	return gopacket.LayerTypePayload
}

var udpPortLayerType = [65536]gopacket.LayerType{
	53:   LayerTypeDNS,
	123:  LayerTypeNTP,
	4789: LayerTypeVXLAN,
	67:   LayerTypeDHCPv4,
	68:   LayerTypeDHCPv4,
	6343: LayerTypeSFlow,
}

// RegisterUDPPortLayerType create a new mapping between an UDPPort
// and an underlaying LayerType.
func RegisterUDPPortLayerType(port UDPPort, layerType gopacket.LayerType) {
	udpPortLayerType[port] = layerType
}

// String returns the port as "number(name)" if there's a well-known port name,
// or just "number" if there isn't.  Well-known names are stored in
// RUDPPortNames.
func (a RUDPPort) String() string {
	if name, ok := RUDPPortNames[a]; ok {
		return fmt.Sprintf("%d(%s)", a, name)
	}
	return strconv.Itoa(int(a))
}

// String returns the port as "number(name)" if there's a well-known port name,
// or just "number" if there isn't.  Well-known names are stored in
// SCTPPortNames.
func (a SCTPPort) String() string {
	if name, ok := SCTPPortNames[a]; ok {
		return fmt.Sprintf("%d(%s)", a, name)
	}
	return strconv.Itoa(int(a))
}

// String returns the port as "number(name)" if there's a well-known port name,
// or just "number" if there isn't.  Well-known names are stored in
// UDPLitePortNames.
func (a UDPLitePort) String() string {
	if name, ok := UDPLitePortNames[a]; ok {
		return fmt.Sprintf("%d(%s)", a, name)
	}
	return strconv.Itoa(int(a))
}
<|endoftext|>


<commit_before>package google

import (
	"fmt"
	"regexp"
	"strings"

	"google.golang.org/api/googleapi"
)

const (
	resolveImageFamilyRegex = "[-_a-zA-Z0-9]*"
	resolveImageImageRegex  = "[-_a-zA-Z0-9]*"
)

var (
	resolveImageProjectImage           = regexp.MustCompile(fmt.Sprintf("projects/(%s)/global/images/(%s)$", ProjectRegex, resolveImageImageRegex))
	resolveImageProjectFamily          = regexp.MustCompile(fmt.Sprintf("projects/(%s)/global/images/family/(%s)$", ProjectRegex, resolveImageFamilyRegex))
	resolveImageGlobalImage            = regexp.MustCompile(fmt.Sprintf("^global/images/(%s)$", resolveImageImageRegex))
	resolveImageGlobalFamily           = regexp.MustCompile(fmt.Sprintf("^global/images/family/(%s)$", resolveImageFamilyRegex))
	resolveImageFamilyFamily           = regexp.MustCompile(fmt.Sprintf("^family/(%s)$", resolveImageFamilyRegex))
	resolveImageProjectImageShorthand  = regexp.MustCompile(fmt.Sprintf("^(%s)/(%s)$", ProjectRegex, resolveImageImageRegex))
	resolveImageProjectFamilyShorthand = regexp.MustCompile(fmt.Sprintf("^(%s)/(%s)$", ProjectRegex, resolveImageFamilyRegex))
	resolveImageFamily                 = regexp.MustCompile(fmt.Sprintf("^(%s)$", resolveImageFamilyRegex))
	resolveImageImage                  = regexp.MustCompile(fmt.Sprintf("^(%s)$", resolveImageImageRegex))
	resolveImageLink                   = regexp.MustCompile(fmt.Sprintf("^https://www.googleapis.com/compute/[a-z0-9]+/projects/(%s)/global/images/(%s)", ProjectRegex, resolveImageImageRegex))

	windowsSqlImage         = regexp.MustCompile("^sql-(?:server-)?([0-9]{4})-([a-z]+)-windows-(?:server-)?([0-9]{4})(?:-r([0-9]+))?-dc-v[0-9]+$")
	canonicalUbuntuLtsImage = regexp.MustCompile("^ubuntu-(minimal-)?([0-9]+)-")
	cosLtsImage             = regexp.MustCompile("^cos-([0-9]+)-")
)

// built-in projects to look for images/families containing the string
// on the left in
var imageMap = map[string]string{
	"centos":      "centos-cloud",
	"coreos":      "coreos-cloud",
	"debian":      "debian-cloud",
	"opensuse":    "opensuse-cloud",
	"rhel":        "rhel-cloud",
	"sles":        "suse-cloud",
	"ubuntu":      "ubuntu-os-cloud",
	"windows":     "windows-cloud",
	"windows-sql": "windows-sql-cloud",
}

func resolveImageImageExists(c *Config, project, name, userAgent string) (bool, error) {
	if _, err := c.NewComputeClient(userAgent).Images.Get(project, name).Do(); err == nil {
		return true, nil
	} else if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
		return false, nil
	} else {
		return false, fmt.Errorf("Error checking if image %s exists: %s", name, err)
	}
}

func resolveImageFamilyExists(c *Config, project, name, userAgent string) (bool, error) {
	if _, err := c.NewComputeClient(userAgent).Images.GetFromFamily(project, name).Do(); err == nil {
		return true, nil
	} else if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
		return false, nil
	} else {
		return false, fmt.Errorf("Error checking if family %s exists: %s", name, err)
	}
}

func sanityTestRegexMatches(expected int, got []string, regexType, name string) error {
	if len(got)-1 != expected { // subtract one, index zero is the entire matched expression
		return fmt.Errorf("Expected %d %s regex matches, got %d for %s", expected, regexType, len(got)-1, name)
	}
	return nil
}

// If the given name is a URL, return it.
// If it's in the form projects/{project}/global/images/{image}, return it
// If it's in the form projects/{project}/global/images/family/{family}, return it
// If it's in the form global/images/{image}, return it
// If it's in the form global/images/family/{family}, return it
// If it's in the form family/{family}, check if it's a family in the current project. If it is, return it as global/images/family/{family}.
//    If not, check if it could be a GCP-provided family, and if it exists. If it does, return it as projects/{project}/global/images/family/{family}.
// If it's in the form {project}/{family-or-image}, check if it's an image in the named project. If it is, return it as projects/{project}/global/images/{image}.
//    If not, check if it's a family in the named project. If it is, return it as projects/{project}/global/images/family/{family}.
// If it's in the form {family-or-image}, check if it's an image in the current project. If it is, return it as global/images/{image}.
//    If not, check if it could be a GCP-provided image, and if it exists. If it does, return it as projects/{project}/global/images/{image}.
//    If not, check if it's a family in the current project. If it is, return it as global/images/family/{family}.
//    If not, check if it could be a GCP-provided family, and if it exists. If it does, return it as projects/{project}/global/images/family/{family}
func resolveImage(c *Config, project, name, userAgent string) (string, error) {
	var builtInProject string
	for k, v := range imageMap {
		if strings.Contains(name, k) {
			builtInProject = v
			break
		}
	}
	switch {
	case resolveImageLink.MatchString(name): // https://www.googleapis.com/compute/v1/projects/xyz/global/images/xyz
		return name, nil
	case resolveImageProjectImage.MatchString(name): // projects/xyz/global/images/xyz
		res := resolveImageProjectImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/%s", res[1], res[2]), nil
	case resolveImageProjectFamily.MatchString(name): // projects/xyz/global/images/family/xyz
		res := resolveImageProjectFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
	case resolveImageGlobalImage.MatchString(name): // global/images/xyz
		res := resolveImageGlobalImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("global/images/%s", res[1]), nil
	case resolveImageGlobalFamily.MatchString(name): // global/images/family/xyz
		res := resolveImageGlobalFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("global/images/family/%s", res[1]), nil
	case resolveImageFamilyFamily.MatchString(name): // family/xyz
		res := resolveImageFamilyFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "family family", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, project, res[1], userAgent); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/family/%s", res[1]), nil
		}
		if builtInProject != "" {
			if ok, err := resolveImageFamilyExists(c, builtInProject, res[1], userAgent); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/family/%s", builtInProject, res[1]), nil
			}
		}
	case resolveImageProjectImageShorthand.MatchString(name): // xyz/xyz
		res := resolveImageProjectImageShorthand.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project image shorthand", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageImageExists(c, res[1], res[2], userAgent); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("projects/%s/global/images/%s", res[1], res[2]), nil
		}
		fallthrough // check if it's a family
	case resolveImageProjectFamilyShorthand.MatchString(name): // xyz/xyz
		res := resolveImageProjectFamilyShorthand.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family shorthand", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, res[1], res[2], userAgent); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
		}
	case resolveImageImage.MatchString(name): // xyz
		res := resolveImageImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "image", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageImageExists(c, project, res[1], userAgent); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/%s", res[1]), nil
		}
		if builtInProject != "" {
			// check the images GCP provides
			if ok, err := resolveImageImageExists(c, builtInProject, res[1], userAgent); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/%s", builtInProject, res[1]), nil
			}
		}
		fallthrough // check if the name is a family, instead of an image
	case resolveImageFamily.MatchString(name): // xyz
		res := resolveImageFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "family", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, c.Project, res[1], userAgent); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/family/%s", res[1]), nil
		}
		if builtInProject != "" {
			// check the families GCP provides
			if ok, err := resolveImageFamilyExists(c, builtInProject, res[1], userAgent); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/family/%s", builtInProject, res[1]), nil
			}
		}
	}
	return "", fmt.Errorf("Could not find image or family %s", name)
}

// resolveImageRefToRelativeURI takes the output of resolveImage and coerces it
// into a relative URI. In the event that a global/images/IMAGE or
// global/images/family/FAMILY reference is returned from resolveImage,
// providerProject will be used as the project for the self_link.
func resolveImageRefToRelativeURI(providerProject, name string) (string, error) {
	switch {
	case resolveImageLink.MatchString(name): // https://www.googleapis.com/compute/v1/projects/xyz/global/images/xyz
		namePath, err := getRelativePath(name)
		if err != nil {
			return "", err
		}

		return namePath, nil
	case resolveImageProjectImage.MatchString(name): // projects/xyz/global/images/xyz
		return name, nil
	case resolveImageProjectFamily.MatchString(name): // projects/xyz/global/images/family/xyz
		res := resolveImageProjectFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
	case resolveImageGlobalImage.MatchString(name): // global/images/xyz
		res := resolveImageGlobalImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/%s", providerProject, res[1]), nil
	case resolveImageGlobalFamily.MatchString(name): // global/images/family/xyz
		res := resolveImageGlobalFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/family/%s", providerProject, res[1]), nil
	}
	return "", fmt.Errorf("Could not expand image or family %q into a relative URI", name)

}
<commit_msg>Add rocky-linux-cloud to the builtin project to search for images if image name contains rocky-linux Closes #11184 (#5773)<commit_after>package google

import (
	"fmt"
	"regexp"
	"strings"

	"google.golang.org/api/googleapi"
)

const (
	resolveImageFamilyRegex = "[-_a-zA-Z0-9]*"
	resolveImageImageRegex  = "[-_a-zA-Z0-9]*"
)

var (
	resolveImageProjectImage           = regexp.MustCompile(fmt.Sprintf("projects/(%s)/global/images/(%s)$", ProjectRegex, resolveImageImageRegex))
	resolveImageProjectFamily          = regexp.MustCompile(fmt.Sprintf("projects/(%s)/global/images/family/(%s)$", ProjectRegex, resolveImageFamilyRegex))
	resolveImageGlobalImage            = regexp.MustCompile(fmt.Sprintf("^global/images/(%s)$", resolveImageImageRegex))
	resolveImageGlobalFamily           = regexp.MustCompile(fmt.Sprintf("^global/images/family/(%s)$", resolveImageFamilyRegex))
	resolveImageFamilyFamily           = regexp.MustCompile(fmt.Sprintf("^family/(%s)$", resolveImageFamilyRegex))
	resolveImageProjectImageShorthand  = regexp.MustCompile(fmt.Sprintf("^(%s)/(%s)$", ProjectRegex, resolveImageImageRegex))
	resolveImageProjectFamilyShorthand = regexp.MustCompile(fmt.Sprintf("^(%s)/(%s)$", ProjectRegex, resolveImageFamilyRegex))
	resolveImageFamily                 = regexp.MustCompile(fmt.Sprintf("^(%s)$", resolveImageFamilyRegex))
	resolveImageImage                  = regexp.MustCompile(fmt.Sprintf("^(%s)$", resolveImageImageRegex))
	resolveImageLink                   = regexp.MustCompile(fmt.Sprintf("^https://www.googleapis.com/compute/[a-z0-9]+/projects/(%s)/global/images/(%s)", ProjectRegex, resolveImageImageRegex))

	windowsSqlImage         = regexp.MustCompile("^sql-(?:server-)?([0-9]{4})-([a-z]+)-windows-(?:server-)?([0-9]{4})(?:-r([0-9]+))?-dc-v[0-9]+$")
	canonicalUbuntuLtsImage = regexp.MustCompile("^ubuntu-(minimal-)?([0-9]+)-")
	cosLtsImage             = regexp.MustCompile("^cos-([0-9]+)-")
)

// built-in projects to look for images/families containing the string
// on the left in
var imageMap = map[string]string{
	"centos":      "centos-cloud",
	"coreos":      "coreos-cloud",
	"debian":      "debian-cloud",
	"opensuse":    "opensuse-cloud",
	"rhel":        "rhel-cloud",
	"rocky-linux": "rocky-linux-cloud",
	"sles":        "suse-cloud",
	"ubuntu":      "ubuntu-os-cloud",
	"windows":     "windows-cloud",
	"windows-sql": "windows-sql-cloud",
}

func resolveImageImageExists(c *Config, project, name, userAgent string) (bool, error) {
	if _, err := c.NewComputeClient(userAgent).Images.Get(project, name).Do(); err == nil {
		return true, nil
	} else if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
		return false, nil
	} else {
		return false, fmt.Errorf("Error checking if image %s exists: %s", name, err)
	}
}

func resolveImageFamilyExists(c *Config, project, name, userAgent string) (bool, error) {
	if _, err := c.NewComputeClient(userAgent).Images.GetFromFamily(project, name).Do(); err == nil {
		return true, nil
	} else if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
		return false, nil
	} else {
		return false, fmt.Errorf("Error checking if family %s exists: %s", name, err)
	}
}

func sanityTestRegexMatches(expected int, got []string, regexType, name string) error {
	if len(got)-1 != expected { // subtract one, index zero is the entire matched expression
		return fmt.Errorf("Expected %d %s regex matches, got %d for %s", expected, regexType, len(got)-1, name)
	}
	return nil
}

// If the given name is a URL, return it.
// If it's in the form projects/{project}/global/images/{image}, return it
// If it's in the form projects/{project}/global/images/family/{family}, return it
// If it's in the form global/images/{image}, return it
// If it's in the form global/images/family/{family}, return it
// If it's in the form family/{family}, check if it's a family in the current project. If it is, return it as global/images/family/{family}.
//    If not, check if it could be a GCP-provided family, and if it exists. If it does, return it as projects/{project}/global/images/family/{family}.
// If it's in the form {project}/{family-or-image}, check if it's an image in the named project. If it is, return it as projects/{project}/global/images/{image}.
//    If not, check if it's a family in the named project. If it is, return it as projects/{project}/global/images/family/{family}.
// If it's in the form {family-or-image}, check if it's an image in the current project. If it is, return it as global/images/{image}.
//    If not, check if it could be a GCP-provided image, and if it exists. If it does, return it as projects/{project}/global/images/{image}.
//    If not, check if it's a family in the current project. If it is, return it as global/images/family/{family}.
//    If not, check if it could be a GCP-provided family, and if it exists. If it does, return it as projects/{project}/global/images/family/{family}
func resolveImage(c *Config, project, name, userAgent string) (string, error) {
	var builtInProject string
	for k, v := range imageMap {
		if strings.Contains(name, k) {
			builtInProject = v
			break
		}
	}
	switch {
	case resolveImageLink.MatchString(name): // https://www.googleapis.com/compute/v1/projects/xyz/global/images/xyz
		return name, nil
	case resolveImageProjectImage.MatchString(name): // projects/xyz/global/images/xyz
		res := resolveImageProjectImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/%s", res[1], res[2]), nil
	case resolveImageProjectFamily.MatchString(name): // projects/xyz/global/images/family/xyz
		res := resolveImageProjectFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
	case resolveImageGlobalImage.MatchString(name): // global/images/xyz
		res := resolveImageGlobalImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("global/images/%s", res[1]), nil
	case resolveImageGlobalFamily.MatchString(name): // global/images/family/xyz
		res := resolveImageGlobalFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("global/images/family/%s", res[1]), nil
	case resolveImageFamilyFamily.MatchString(name): // family/xyz
		res := resolveImageFamilyFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "family family", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, project, res[1], userAgent); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/family/%s", res[1]), nil
		}
		if builtInProject != "" {
			if ok, err := resolveImageFamilyExists(c, builtInProject, res[1], userAgent); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/family/%s", builtInProject, res[1]), nil
			}
		}
	case resolveImageProjectImageShorthand.MatchString(name): // xyz/xyz
		res := resolveImageProjectImageShorthand.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project image shorthand", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageImageExists(c, res[1], res[2], userAgent); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("projects/%s/global/images/%s", res[1], res[2]), nil
		}
		fallthrough // check if it's a family
	case resolveImageProjectFamilyShorthand.MatchString(name): // xyz/xyz
		res := resolveImageProjectFamilyShorthand.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family shorthand", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, res[1], res[2], userAgent); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
		}
	case resolveImageImage.MatchString(name): // xyz
		res := resolveImageImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "image", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageImageExists(c, project, res[1], userAgent); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/%s", res[1]), nil
		}
		if builtInProject != "" {
			// check the images GCP provides
			if ok, err := resolveImageImageExists(c, builtInProject, res[1], userAgent); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/%s", builtInProject, res[1]), nil
			}
		}
		fallthrough // check if the name is a family, instead of an image
	case resolveImageFamily.MatchString(name): // xyz
		res := resolveImageFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "family", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, c.Project, res[1], userAgent); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/family/%s", res[1]), nil
		}
		if builtInProject != "" {
			// check the families GCP provides
			if ok, err := resolveImageFamilyExists(c, builtInProject, res[1], userAgent); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/family/%s", builtInProject, res[1]), nil
			}
		}
	}
	return "", fmt.Errorf("Could not find image or family %s", name)
}

// resolveImageRefToRelativeURI takes the output of resolveImage and coerces it
// into a relative URI. In the event that a global/images/IMAGE or
// global/images/family/FAMILY reference is returned from resolveImage,
// providerProject will be used as the project for the self_link.
func resolveImageRefToRelativeURI(providerProject, name string) (string, error) {
	switch {
	case resolveImageLink.MatchString(name): // https://www.googleapis.com/compute/v1/projects/xyz/global/images/xyz
		namePath, err := getRelativePath(name)
		if err != nil {
			return "", err
		}

		return namePath, nil
	case resolveImageProjectImage.MatchString(name): // projects/xyz/global/images/xyz
		return name, nil
	case resolveImageProjectFamily.MatchString(name): // projects/xyz/global/images/family/xyz
		res := resolveImageProjectFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
	case resolveImageGlobalImage.MatchString(name): // global/images/xyz
		res := resolveImageGlobalImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/%s", providerProject, res[1]), nil
	case resolveImageGlobalFamily.MatchString(name): // global/images/family/xyz
		res := resolveImageGlobalFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/family/%s", providerProject, res[1]), nil
	}
	return "", fmt.Errorf("Could not expand image or family %q into a relative URI", name)

}
<|endoftext|>


<commit_before>package itest

import (
	"context"
	"fmt"
	"testing"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcutil"
	"github.com/lightningnetwork/lnd/lnrpc"
	"github.com/lightningnetwork/lnd/lnrpc/invoicesrpc"
	"github.com/lightningnetwork/lnd/lntest"
	"github.com/lightningnetwork/lnd/lntypes"
)

func testMultiHopHtlcClaims(net *lntest.NetworkHarness, t *harnessTest) {

	type testCase struct {
		name string
		test func(net *lntest.NetworkHarness, t *harnessTest, alice,
			bob *lntest.HarnessNode, c commitType)
	}

	subTests := []testCase{
		{
			// bob: outgoing our commit timeout
			// carol: incoming their commit watch and see timeout
			name: "local force close immediate expiry",
			test: testMultiHopHtlcLocalTimeout,
		},
		{
			// bob: outgoing watch and see, they sweep on chain
			// carol: incoming our commit, know preimage
			name: "receiver chain claim",
			test: testMultiHopReceiverChainClaim,
		},
		{
			// bob: outgoing our commit watch and see timeout
			// carol: incoming their commit watch and see timeout
			name: "local force close on-chain htlc timeout",
			test: testMultiHopLocalForceCloseOnChainHtlcTimeout,
		},
		{
			// bob: outgoing their commit watch and see timeout
			// carol: incoming our commit watch and see timeout
			name: "remote force close on-chain htlc timeout",
			test: testMultiHopRemoteForceCloseOnChainHtlcTimeout,
		},
		{
			// bob: outgoing our commit watch and see, they sweep on chain
			// bob: incoming our commit watch and learn preimage
			// carol: incoming their commit know preimage
			name: "local chain claim",
			test: testMultiHopHtlcLocalChainClaim,
		},
		{
			// bob: outgoing their commit watch and see, they sweep on chain
			// bob: incoming their commit watch and learn preimage
			// carol: incoming our commit know preimage
			name: "remote chain claim",
			test: testMultiHopHtlcRemoteChainClaim,
		},
	}

	commitTypes := []commitType{
		commitTypeLegacy,
		commitTypeAnchors,
	}

	for _, commitType := range commitTypes {
		testName := fmt.Sprintf("committype=%v", commitType.String())

		commitType := commitType
		success := t.t.Run(testName, func(t *testing.T) {
			ht := newHarnessTest(t, net)

			args := commitType.Args()
			alice, err := net.NewNode("Alice", args)
			if err != nil {
				t.Fatalf("unable to create new node: %v", err)
			}
			defer shutdownAndAssert(net, ht, alice)

			bob, err := net.NewNode("Bob", args)
			if err != nil {
				t.Fatalf("unable to create new node: %v", err)
			}
			defer shutdownAndAssert(net, ht, bob)

			ctxb := context.Background()
			ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
			if err := net.ConnectNodes(ctxt, alice, bob); err != nil {
				t.Fatalf("unable to connect alice to bob: %v", err)
			}

			for _, subTest := range subTests {
				subTest := subTest

				success := ht.t.Run(subTest.name, func(t *testing.T) {
					ht := newHarnessTest(t, net)

					// Start each test with the default
					// static fee estimate.
					net.SetFeeEstimate(12500)

					subTest.test(net, ht, alice, bob, commitType)
				})
				if !success {
					return
				}
			}
		})
		if !success {
			return
		}
	}
}

// waitForInvoiceAccepted waits until the specified invoice moved to the
// accepted state by the node.
func waitForInvoiceAccepted(t *harnessTest, node *lntest.HarnessNode,
	payHash lntypes.Hash) {

	ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout)
	defer cancel()
	invoiceUpdates, err := node.SubscribeSingleInvoice(ctx,
		&invoicesrpc.SubscribeSingleInvoiceRequest{
			RHash: payHash[:],
		},
	)
	if err != nil {
		t.Fatalf("subscribe single invoice: %v", err)
	}

	for {
		update, err := invoiceUpdates.Recv()
		if err != nil {
			t.Fatalf("invoice update err: %v", err)
		}
		if update.State == lnrpc.Invoice_ACCEPTED {
			break
		}
	}
}

// checkPaymentStatus asserts that the given node list a payment with the given
// preimage has the expected status.
func checkPaymentStatus(ctxt context.Context, node *lntest.HarnessNode,
	preimage lntypes.Preimage, status lnrpc.Payment_PaymentStatus) error {

	req := &lnrpc.ListPaymentsRequest{
		IncludeIncomplete: true,
	}
	paymentsResp, err := node.ListPayments(ctxt, req)
	if err != nil {
		return fmt.Errorf("error when obtaining Alice payments: %v",
			err)
	}

	payHash := preimage.Hash()
	var found bool
	for _, p := range paymentsResp.Payments {
		if p.PaymentHash != payHash.String() {
			continue
		}

		found = true
		if p.Status != status {
			return fmt.Errorf("expected payment status "+
				"%v, got %v", status, p.Status)
		}

		switch status {

		// If this expected status is SUCCEEDED, we expect the final preimage.
		case lnrpc.Payment_SUCCEEDED:
			if p.PaymentPreimage != preimage.String() {
				return fmt.Errorf("preimage doesn't match: %v vs %v",
					p.PaymentPreimage, preimage.String())
			}

		// Otherwise we expect an all-zero preimage.
		default:
			if p.PaymentPreimage != (lntypes.Preimage{}).String() {
				return fmt.Errorf("expected zero preimage, got %v",
					p.PaymentPreimage)
			}
		}

	}

	if !found {
		return fmt.Errorf("payment with payment hash %v not found "+
			"in response", payHash)
	}

	return nil
}

func createThreeHopNetwork(t *harnessTest, net *lntest.NetworkHarness,
	alice, bob *lntest.HarnessNode, carolHodl bool, c commitType) (
	*lnrpc.ChannelPoint, *lnrpc.ChannelPoint, *lntest.HarnessNode) {

	ctxb := context.Background()

	ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
	err := net.EnsureConnected(ctxt, alice, bob)
	if err != nil {
		t.Fatalf("unable to connect peers: %v", err)
	}

	// Make sure there are enough utxos for anchoring.
	for i := 0; i < 2; i++ {
		ctxt, _ = context.WithTimeout(context.Background(), defaultTimeout)
		err = net.SendCoins(ctxt, btcutil.SatoshiPerBitcoin, alice)
		if err != nil {
			t.Fatalf("unable to send coins to Alice: %v", err)
		}

		ctxt, _ = context.WithTimeout(context.Background(), defaultTimeout)
		err = net.SendCoins(ctxt, btcutil.SatoshiPerBitcoin, bob)
		if err != nil {
			t.Fatalf("unable to send coins to Bob: %v", err)
		}
	}

	// We'll start the test by creating a channel between Alice and Bob,
	// which will act as the first leg for out multi-hop HTLC.
	const chanAmt = 1000000
	ctxt, _ = context.WithTimeout(ctxb, channelOpenTimeout)
	aliceChanPoint := openChannelAndAssert(
		ctxt, t, net, alice, bob,
		lntest.OpenChannelParams{
			Amt: chanAmt,
		},
	)

	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
	err = alice.WaitForNetworkChannelOpen(ctxt, aliceChanPoint)
	if err != nil {
		t.Fatalf("alice didn't report channel: %v", err)
	}

	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
	err = bob.WaitForNetworkChannelOpen(ctxt, aliceChanPoint)
	if err != nil {
		t.Fatalf("bob didn't report channel: %v", err)
	}

	// Next, we'll create a new node "carol" and have Bob connect to her. If
	// the carolHodl flag is set, we'll make carol always hold onto the
	// HTLC, this way it'll force Bob to go to chain to resolve the HTLC.
	carolFlags := c.Args()
	if carolHodl {
		carolFlags = append(carolFlags, "--hodl.exit-settle")
	}
	carol, err := net.NewNode("Carol", carolFlags)
	if err != nil {
		t.Fatalf("unable to create new node: %v", err)
	}
	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
	if err := net.ConnectNodes(ctxt, bob, carol); err != nil {
		t.Fatalf("unable to connect bob to carol: %v", err)
	}

	// Make sure Carol has enough utxos for anchoring. Because the anchor by
	// itself often doesn't meet the dust limit, a utxo from the wallet
	// needs to be attached as an additional input. This can still lead to a
	// positively-yielding transaction.
	for i := 0; i < 2; i++ {
		ctxt, _ = context.WithTimeout(context.Background(), defaultTimeout)
		err = net.SendCoins(ctxt, btcutil.SatoshiPerBitcoin, carol)
		if err != nil {
			t.Fatalf("unable to send coins to Alice: %v", err)
		}
	}

	// We'll then create a channel from Bob to Carol. After this channel is
	// open, our topology looks like:  A -> B -> C.
	ctxt, _ = context.WithTimeout(ctxb, channelOpenTimeout)
	bobChanPoint := openChannelAndAssert(
		ctxt, t, net, bob, carol,
		lntest.OpenChannelParams{
			Amt: chanAmt,
		},
	)
	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
	err = bob.WaitForNetworkChannelOpen(ctxt, bobChanPoint)
	if err != nil {
		t.Fatalf("alice didn't report channel: %v", err)
	}
	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
	err = carol.WaitForNetworkChannelOpen(ctxt, bobChanPoint)
	if err != nil {
		t.Fatalf("bob didn't report channel: %v", err)
	}
	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
	err = alice.WaitForNetworkChannelOpen(ctxt, bobChanPoint)
	if err != nil {
		t.Fatalf("bob didn't report channel: %v", err)
	}

	return aliceChanPoint, bobChanPoint, carol
}

// assertAllTxesSpendFrom asserts that all txes in the list spend from the given
// tx.
func assertAllTxesSpendFrom(t *harnessTest, txes []*wire.MsgTx,
	prevTxid chainhash.Hash) {

	for _, tx := range txes {
		if tx.TxIn[0].PreviousOutPoint.Hash != prevTxid {
			t.Fatalf("tx %v did not spend from %v",
				tx.TxHash(), prevTxid)
		}
	}
}
<commit_msg>itest: fix incorrect error message<commit_after>package itest

import (
	"context"
	"fmt"
	"testing"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcutil"
	"github.com/lightningnetwork/lnd/lnrpc"
	"github.com/lightningnetwork/lnd/lnrpc/invoicesrpc"
	"github.com/lightningnetwork/lnd/lntest"
	"github.com/lightningnetwork/lnd/lntypes"
)

func testMultiHopHtlcClaims(net *lntest.NetworkHarness, t *harnessTest) {

	type testCase struct {
		name string
		test func(net *lntest.NetworkHarness, t *harnessTest, alice,
			bob *lntest.HarnessNode, c commitType)
	}

	subTests := []testCase{
		{
			// bob: outgoing our commit timeout
			// carol: incoming their commit watch and see timeout
			name: "local force close immediate expiry",
			test: testMultiHopHtlcLocalTimeout,
		},
		{
			// bob: outgoing watch and see, they sweep on chain
			// carol: incoming our commit, know preimage
			name: "receiver chain claim",
			test: testMultiHopReceiverChainClaim,
		},
		{
			// bob: outgoing our commit watch and see timeout
			// carol: incoming their commit watch and see timeout
			name: "local force close on-chain htlc timeout",
			test: testMultiHopLocalForceCloseOnChainHtlcTimeout,
		},
		{
			// bob: outgoing their commit watch and see timeout
			// carol: incoming our commit watch and see timeout
			name: "remote force close on-chain htlc timeout",
			test: testMultiHopRemoteForceCloseOnChainHtlcTimeout,
		},
		{
			// bob: outgoing our commit watch and see, they sweep on chain
			// bob: incoming our commit watch and learn preimage
			// carol: incoming their commit know preimage
			name: "local chain claim",
			test: testMultiHopHtlcLocalChainClaim,
		},
		{
			// bob: outgoing their commit watch and see, they sweep on chain
			// bob: incoming their commit watch and learn preimage
			// carol: incoming our commit know preimage
			name: "remote chain claim",
			test: testMultiHopHtlcRemoteChainClaim,
		},
	}

	commitTypes := []commitType{
		commitTypeLegacy,
		commitTypeAnchors,
	}

	for _, commitType := range commitTypes {
		testName := fmt.Sprintf("committype=%v", commitType.String())

		commitType := commitType
		success := t.t.Run(testName, func(t *testing.T) {
			ht := newHarnessTest(t, net)

			args := commitType.Args()
			alice, err := net.NewNode("Alice", args)
			if err != nil {
				t.Fatalf("unable to create new node: %v", err)
			}
			defer shutdownAndAssert(net, ht, alice)

			bob, err := net.NewNode("Bob", args)
			if err != nil {
				t.Fatalf("unable to create new node: %v", err)
			}
			defer shutdownAndAssert(net, ht, bob)

			ctxb := context.Background()
			ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
			if err := net.ConnectNodes(ctxt, alice, bob); err != nil {
				t.Fatalf("unable to connect alice to bob: %v", err)
			}

			for _, subTest := range subTests {
				subTest := subTest

				success := ht.t.Run(subTest.name, func(t *testing.T) {
					ht := newHarnessTest(t, net)

					// Start each test with the default
					// static fee estimate.
					net.SetFeeEstimate(12500)

					subTest.test(net, ht, alice, bob, commitType)
				})
				if !success {
					return
				}
			}
		})
		if !success {
			return
		}
	}
}

// waitForInvoiceAccepted waits until the specified invoice moved to the
// accepted state by the node.
func waitForInvoiceAccepted(t *harnessTest, node *lntest.HarnessNode,
	payHash lntypes.Hash) {

	ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout)
	defer cancel()
	invoiceUpdates, err := node.SubscribeSingleInvoice(ctx,
		&invoicesrpc.SubscribeSingleInvoiceRequest{
			RHash: payHash[:],
		},
	)
	if err != nil {
		t.Fatalf("subscribe single invoice: %v", err)
	}

	for {
		update, err := invoiceUpdates.Recv()
		if err != nil {
			t.Fatalf("invoice update err: %v", err)
		}
		if update.State == lnrpc.Invoice_ACCEPTED {
			break
		}
	}
}

// checkPaymentStatus asserts that the given node list a payment with the given
// preimage has the expected status.
func checkPaymentStatus(ctxt context.Context, node *lntest.HarnessNode,
	preimage lntypes.Preimage, status lnrpc.Payment_PaymentStatus) error {

	req := &lnrpc.ListPaymentsRequest{
		IncludeIncomplete: true,
	}
	paymentsResp, err := node.ListPayments(ctxt, req)
	if err != nil {
		return fmt.Errorf("error when obtaining Alice payments: %v",
			err)
	}

	payHash := preimage.Hash()
	var found bool
	for _, p := range paymentsResp.Payments {
		if p.PaymentHash != payHash.String() {
			continue
		}

		found = true
		if p.Status != status {
			return fmt.Errorf("expected payment status "+
				"%v, got %v", status, p.Status)
		}

		switch status {

		// If this expected status is SUCCEEDED, we expect the final preimage.
		case lnrpc.Payment_SUCCEEDED:
			if p.PaymentPreimage != preimage.String() {
				return fmt.Errorf("preimage doesn't match: %v vs %v",
					p.PaymentPreimage, preimage.String())
			}

		// Otherwise we expect an all-zero preimage.
		default:
			if p.PaymentPreimage != (lntypes.Preimage{}).String() {
				return fmt.Errorf("expected zero preimage, got %v",
					p.PaymentPreimage)
			}
		}

	}

	if !found {
		return fmt.Errorf("payment with payment hash %v not found "+
			"in response", payHash)
	}

	return nil
}

func createThreeHopNetwork(t *harnessTest, net *lntest.NetworkHarness,
	alice, bob *lntest.HarnessNode, carolHodl bool, c commitType) (
	*lnrpc.ChannelPoint, *lnrpc.ChannelPoint, *lntest.HarnessNode) {

	ctxb := context.Background()

	ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
	err := net.EnsureConnected(ctxt, alice, bob)
	if err != nil {
		t.Fatalf("unable to connect peers: %v", err)
	}

	// Make sure there are enough utxos for anchoring.
	for i := 0; i < 2; i++ {
		ctxt, _ = context.WithTimeout(context.Background(), defaultTimeout)
		err = net.SendCoins(ctxt, btcutil.SatoshiPerBitcoin, alice)
		if err != nil {
			t.Fatalf("unable to send coins to Alice: %v", err)
		}

		ctxt, _ = context.WithTimeout(context.Background(), defaultTimeout)
		err = net.SendCoins(ctxt, btcutil.SatoshiPerBitcoin, bob)
		if err != nil {
			t.Fatalf("unable to send coins to Bob: %v", err)
		}
	}

	// We'll start the test by creating a channel between Alice and Bob,
	// which will act as the first leg for out multi-hop HTLC.
	const chanAmt = 1000000
	ctxt, _ = context.WithTimeout(ctxb, channelOpenTimeout)
	aliceChanPoint := openChannelAndAssert(
		ctxt, t, net, alice, bob,
		lntest.OpenChannelParams{
			Amt: chanAmt,
		},
	)

	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
	err = alice.WaitForNetworkChannelOpen(ctxt, aliceChanPoint)
	if err != nil {
		t.Fatalf("alice didn't report channel: %v", err)
	}

	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
	err = bob.WaitForNetworkChannelOpen(ctxt, aliceChanPoint)
	if err != nil {
		t.Fatalf("bob didn't report channel: %v", err)
	}

	// Next, we'll create a new node "carol" and have Bob connect to her. If
	// the carolHodl flag is set, we'll make carol always hold onto the
	// HTLC, this way it'll force Bob to go to chain to resolve the HTLC.
	carolFlags := c.Args()
	if carolHodl {
		carolFlags = append(carolFlags, "--hodl.exit-settle")
	}
	carol, err := net.NewNode("Carol", carolFlags)
	if err != nil {
		t.Fatalf("unable to create new node: %v", err)
	}
	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
	if err := net.ConnectNodes(ctxt, bob, carol); err != nil {
		t.Fatalf("unable to connect bob to carol: %v", err)
	}

	// Make sure Carol has enough utxos for anchoring. Because the anchor by
	// itself often doesn't meet the dust limit, a utxo from the wallet
	// needs to be attached as an additional input. This can still lead to a
	// positively-yielding transaction.
	for i := 0; i < 2; i++ {
		ctxt, _ = context.WithTimeout(context.Background(), defaultTimeout)
		err = net.SendCoins(ctxt, btcutil.SatoshiPerBitcoin, carol)
		if err != nil {
			t.Fatalf("unable to send coins to Carol: %v", err)
		}
	}

	// We'll then create a channel from Bob to Carol. After this channel is
	// open, our topology looks like:  A -> B -> C.
	ctxt, _ = context.WithTimeout(ctxb, channelOpenTimeout)
	bobChanPoint := openChannelAndAssert(
		ctxt, t, net, bob, carol,
		lntest.OpenChannelParams{
			Amt: chanAmt,
		},
	)
	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
	err = bob.WaitForNetworkChannelOpen(ctxt, bobChanPoint)
	if err != nil {
		t.Fatalf("alice didn't report channel: %v", err)
	}
	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
	err = carol.WaitForNetworkChannelOpen(ctxt, bobChanPoint)
	if err != nil {
		t.Fatalf("bob didn't report channel: %v", err)
	}
	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
	err = alice.WaitForNetworkChannelOpen(ctxt, bobChanPoint)
	if err != nil {
		t.Fatalf("bob didn't report channel: %v", err)
	}

	return aliceChanPoint, bobChanPoint, carol
}

// assertAllTxesSpendFrom asserts that all txes in the list spend from the given
// tx.
func assertAllTxesSpendFrom(t *harnessTest, txes []*wire.MsgTx,
	prevTxid chainhash.Hash) {

	for _, tx := range txes {
		if tx.TxIn[0].PreviousOutPoint.Hash != prevTxid {
			t.Fatalf("tx %v did not spend from %v",
				tx.TxHash(), prevTxid)
		}
	}
}
<|endoftext|>


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

// Package gsutil is collection of helper functions for storage tools.
package gsutil

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"kythe.io/kythe/go/services/graphstore"
	"kythe.io/kythe/go/storage/inmemory"
)

// Handler returns a graphstore.Service based on the given specification.
// See also: Register(string, Handler).
type Handler func(spec string) (graphstore.Service, error)

var (
	handlers = map[string]Handler{
		"in-memory": func(_ string) (graphstore.Service, error) {
			return new(inmemory.GraphStore), nil
		},
	}
	defaultHandlerKind string
)

// Register exposes the given Handler to ParseGraphStore.  Each string starting
// with kind+":" will be passed to the given Handler.  A kind can only be
// registered once.
func Register(kind string, h Handler) {
	if _, exists := handlers[kind]; exists {
		log.Fatalf("gsutil Handler for kind %q already exists", kind)
	}
	handlers[kind] = h
}

// RegisterDefault gives ParseGraphStore a fallback kind if not given any
// "____:" prefix.  A default can only be set once.
func RegisterDefault(kind string) {
	if defaultHandlerKind != "" {
		log.Fatalf("default gsutil Handler kind already registered as %q", defaultHandlerKind)
	}
	defaultHandlerKind = kind
}

type gsFlag struct {
	gs *graphstore.Service
}

// String implements part of the flag.Value interface.
func (f *gsFlag) String() string { return fmt.Sprintf("%T", *f.gs) }

// Set implements part of the flag.Value interface.
func (f *gsFlag) Set(str string) (err error) {
	*f.gs, err = ParseGraphStore(str)
	return
}

// Flag defines a GraphStore flag with the specified name and usage string.
func Flag(gs *graphstore.Service, name, usage string) {
	if gs == nil {
		log.Fatal("GraphStoreFlag given nil GraphStore pointer")
	}
	f := gsFlag{gs: gs}
	flag.Var(&f, name, usage)
}

// ParseGraphStore returns a GraphStore for the given specification.
func ParseGraphStore(str string) (graphstore.Service, error) {
	str = strings.TrimSpace(str)
	split := strings.SplitN(str, ":", 2)
	var kind, spec string
	if len(split) == 2 {
		spec = split[1]
		kind = split[0]
	} else {
		spec = str
		switch {
		case spec == "" || spec == "in-memory":
			kind = "in-memory"
		case defaultHandlerKind != "":
			kind = defaultHandlerKind
		default:
			return nil, fmt.Errorf("unknown GraphStore: %q", str)
		}
	}

	h, ok := handlers[kind]
	if !ok {
		return nil, fmt.Errorf("no gsutil Handler registered for kind %q", kind)
	}
	return h(spec)
}

// EnsureGracefulExit will try to close each gs when notified of an Interrupt,
// SIGTERM, or Kill signal and immediately exit the program unsuccessfully. Any
// errors will be logged. This function should only be called once and closing
// the GraphStores manually is still needed when the program does not receive a
// signal to quit.
func EnsureGracefulExit(gs ...graphstore.Service) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)
	signal.Notify(c, os.Kill)
	go func() {
		sig := <-c
		log.Printf("graphstore: signal %v", sig)
		for _, g := range gs {
			LogClose(context.Background(), g)
		}
		os.Exit(1)
	}()
}

// LogClose closes gs and logs any resulting error.
func LogClose(ctx context.Context, gs graphstore.Service) {
	if err := gs.Close(ctx); err != nil {
		log.Printf("GraphStore failed to close: %v", err)
	}
}
<commit_msg>Fix a panic in the gsFlag String implementation.<commit_after>/*
 * Copyright 2014 Google Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Package gsutil is collection of helper functions for storage tools.
package gsutil

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"kythe.io/kythe/go/services/graphstore"
	"kythe.io/kythe/go/storage/inmemory"
)

// Handler returns a graphstore.Service based on the given specification.
// See also: Register(string, Handler).
type Handler func(spec string) (graphstore.Service, error)

var (
	handlers = map[string]Handler{
		"in-memory": func(_ string) (graphstore.Service, error) {
			return new(inmemory.GraphStore), nil
		},
	}
	defaultHandlerKind string
)

// Register exposes the given Handler to ParseGraphStore.  Each string starting
// with kind+":" will be passed to the given Handler.  A kind can only be
// registered once.
func Register(kind string, h Handler) {
	if _, exists := handlers[kind]; exists {
		log.Fatalf("gsutil Handler for kind %q already exists", kind)
	}
	handlers[kind] = h
}

// RegisterDefault gives ParseGraphStore a fallback kind if not given any
// "____:" prefix.  A default can only be set once.
func RegisterDefault(kind string) {
	if defaultHandlerKind != "" {
		log.Fatalf("default gsutil Handler kind already registered as %q", defaultHandlerKind)
	}
	defaultHandlerKind = kind
}

type gsFlag struct {
	gs *graphstore.Service
}

// String implements part of the flag.Value interface.
func (f *gsFlag) String() string {
	if f.gs == nil {
		return "<graphstore>"
	}
	return fmt.Sprintf("%T", *f.gs)
}

// Set implements part of the flag.Value interface.
func (f *gsFlag) Set(str string) (err error) {
	*f.gs, err = ParseGraphStore(str)
	return
}

// Flag defines a GraphStore flag with the specified name and usage string.
func Flag(gs *graphstore.Service, name, usage string) {
	if gs == nil {
		log.Fatal("GraphStoreFlag given nil GraphStore pointer")
	}
	f := gsFlag{gs: gs}
	flag.Var(&f, name, usage)
}

// ParseGraphStore returns a GraphStore for the given specification.
func ParseGraphStore(str string) (graphstore.Service, error) {
	str = strings.TrimSpace(str)
	split := strings.SplitN(str, ":", 2)
	var kind, spec string
	if len(split) == 2 {
		spec = split[1]
		kind = split[0]
	} else {
		spec = str
		switch {
		case spec == "" || spec == "in-memory":
			kind = "in-memory"
		case defaultHandlerKind != "":
			kind = defaultHandlerKind
		default:
			return nil, fmt.Errorf("unknown GraphStore: %q", str)
		}
	}

	h, ok := handlers[kind]
	if !ok {
		return nil, fmt.Errorf("no gsutil Handler registered for kind %q", kind)
	}
	return h(spec)
}

// EnsureGracefulExit will try to close each gs when notified of an Interrupt,
// SIGTERM, or Kill signal and immediately exit the program unsuccessfully. Any
// errors will be logged. This function should only be called once and closing
// the GraphStores manually is still needed when the program does not receive a
// signal to quit.
func EnsureGracefulExit(gs ...graphstore.Service) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)
	signal.Notify(c, os.Kill)
	go func() {
		sig := <-c
		log.Printf("graphstore: signal %v", sig)
		for _, g := range gs {
			LogClose(context.Background(), g)
		}
		os.Exit(1)
	}()
}

// LogClose closes gs and logs any resulting error.
func LogClose(ctx context.Context, gs graphstore.Service) {
	if err := gs.Close(ctx); err != nil {
		log.Printf("GraphStore failed to close: %v", err)
	}
}
<|endoftext|>


<commit_before>package elasticorm

import (
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"strings"

	"github.com/pkg/errors"
	elastic "gopkg.in/olivere/elastic.v5"
)

// NewDatastore returns a fresh instance of an elasticorm datastore
func NewDatastore(esc *elastic.Client, options ...DatastoreOptFunc) (*Datastore, error) {
	ds := &Datastore{
		elasticClient: esc,
		Ctx:           context.Background(),
	}
	var err error
	for _, opt := range options {
		err = opt(ds)
	}
	return ds, err
}

// DatastoreOptFunc is used as a parameter to NewDatastore and provides a way of configuration
type DatastoreOptFunc func(*Datastore) error

// Datstore is an instance to communicate easy with elasticsearch
// It is meant to be used for one struct and helps storing and retrieving it in/from elasticsearch
// It leverages the great elastic package from olivere
type Datastore struct {
	elasticClient   *elastic.Client
	Ctx             context.Context
	indexName       string
	goType          reflect.Type
	idFieldName     string          // the name of the structs field to store the ID
	typeName        string          // in elasticsearch
	indexDefinition IndexDefinition // in elasticsearch
}

// EnsureIndexExists checks wether the needed index for this datastore exists. It it doesn't it gets created
// the name of the datastore is determined by the structs name (+ plural s)
func (ds *Datastore) EnsureIndexExists() error {
	if ds.indexName == `` {
		return errors.New(`EnsureIndexExists failed, because no index name is defined`)
	}
	exists, err := ds.elasticClient.
		IndexExists(ds.indexName).
		Do(ds.Ctx)

	if exists || err != nil {
		return err
	}

	err = ds.createIndex()
	if err != nil {
		return err
	}
	ds.elasticClient.Refresh().Do(ds.Ctx)
	return nil
}

func (ds *Datastore) createIndex() error {
	JSON, err := json.Marshal(ds.indexDefinition)
	if err != nil {
		return err
	}
	ack, err := ds.elasticClient.
		CreateIndex(ds.indexName).
		BodyString(string(JSON)).
		Do(ds.Ctx)
	if err != nil || !ack.Acknowledged {
		return errors.Wrap(
			err, fmt.Sprintf("creating elasticsearch index %s failed - %s", ds.indexName, string(JSON)),
		)
	}
	return nil
}

// ForStruct generates a DatastoreOptFunc
// It is used to generate a default mapping, type name and index name by analysing a provided struct
func ForStruct(i interface{}) DatastoreOptFunc {
	return func(ds *Datastore) error {
		ds.goType = reflect.TypeOf(i)
		typeName, err := typeNameFromStruct(i)
		if err != nil {
			return err
		}
		ds.typeName = typeName
		ds.indexName = typeName + `s`
		indexDefinition, err := NewIndexDefinition(
			AddMappingFromStruct(ds.typeName, i),
		)
		if err != nil {
			return err
		}
		ds.indexDefinition = indexDefinition
		// TODO configure ID
		ds.idFieldName = `ID`
		return nil
	}
}

func typeNameFromStruct(i interface{}) (string, error) {
	typeName := getType(i)
	if typeName == `` {
		return ``, errors.New(`Could not determine type name from struct`)
	}

	typeName = strings.ToLower(typeName)
	typeName = strings.TrimLeft(typeName, `*`)
	return typeName, nil
}

func getType(myvar interface{}) string {
	return nameOfType(reflect.TypeOf(myvar))
}

func nameOfType(t reflect.Type) string {
	if t.Kind() == reflect.Ptr {
		return "*" + t.Elem().Name()
	} else {
		return t.Name()
	}
}

func (ds *Datastore) Create(o interface{}) error {
	err := ds.isSaveableType(o)
	if err != nil {
		return err
	}
	put, err := ds.elasticClient.Index().
		Index(ds.indexName).
		Type(ds.typeName).
		BodyJson(o).
		Do(ds.Ctx)

	if err != nil {
		return err
	}
	if !put.Created {
		return ErrCreationFailed
	}
	return ds.setID(o, put.Id)
}

func (ds *Datastore) isSaveableType(i interface{}) error {
	rt := reflect.TypeOf(i)
	if nameOfType(rt) != nameOfType(ds.goType) {
		return errors.Wrapf(
			ErrInvalidType,
			`create failed for %s, expected %s`, nameOfType(rt), nameOfType(ds.goType),
		)
	}
	rv := reflect.ValueOf(i)
	if rv.Kind() != reflect.Ptr || rv.IsNil() {
		return errors.Wrap(ErrInvalidType, `no pointer given`)
	}
	return nil
}

func (ds *Datastore) Find(ID string, result interface{}) error {
	err := ds.isSaveableType(result)
	if err != nil {
		return err
	}
	res, err := ds.elasticClient.Get().
		Index(ds.indexName).
		Type(ds.typeName).
		Id(ID).
		Do(ds.Ctx)

	if !res.Found || elastic.IsNotFound(err) {
		return ErrNotFound
	}
	if err != nil {
		return err
	}

	return ds.decodeElasticResponse(res.Source, res.Id, result)
}

func (ds *Datastore) Update(o interface{}) error {
	err := ds.isSaveableType(o)
	if err != nil {
		return err
	}
	ID, err := ds.getID(o)
	if err != nil {
		return err
	}
	if ID == `` {
		return errors.New(`can't save struct with empty ID`)
	}
	_, err = ds.elasticClient.Update().
		Index(ds.indexName).
		Type(ds.typeName).
		Id(ID).
		Doc(o).
		Do(ds.Ctx)

	return err
}

func (ds *Datastore) FindOneBy(fieldName string, value interface{}, result interface{}) error {
	elasticFieldName, err := ds.indexDefinition.elasticFieldName(ds.typeName, fieldName)
	if err != nil {
		return err
	}
	res, err := ds.elasticClient.Search().
		Index(ds.indexName).
		Query(elastic.NewBoolQuery().Filter(elastic.NewTermQuery(elasticFieldName, value))).
		From(0).Size(1).
		Do(ds.Ctx)
	if err != nil {
		return err
	}
	if res.TotalHits() < 1 {
		return ErrNotFound
	}
	return ds.decodeElasticResponse(res.Hits.Hits[0].Source, res.Hits.Hits[0].Id, result)
}

func (ds *Datastore) setID(o interface{}, ID string) error {
	eo := reflect.ValueOf(o).Elem()
	if eo.Kind() != reflect.Struct {
		return errors.Wrap(ErrInvalidType, `setID failed`)
	}
	idField := eo.FieldByName(ds.idFieldName)
	if !idField.IsValid() || !idField.CanSet() || idField.Kind() != reflect.String {
		return ErrInvalidIDField
	}
	idField.SetString(ID)
	return nil
}

func (ds *Datastore) getID(o interface{}) (string, error) {
	eo := reflect.ValueOf(o).Elem()
	if eo.Kind() != reflect.Struct {
		return ``, errors.Wrap(ErrInvalidType, `getID failed`)
	}
	idField := eo.FieldByName(ds.idFieldName)
	if !idField.IsValid() || idField.Kind() != reflect.String {
		return ``, ErrInvalidIDField
	}
	return idField.String(), nil
}

func (ds *Datastore) decodeElasticResponse(source *json.RawMessage, ID string, o interface{}) error {
	err := json.Unmarshal(*source, o)
	if err != nil {
		return err
	}
	ds.setID(o, ID)
	return nil
}
<commit_msg>[FEATURE] CleanUp, Refresh and EnsureIndexDoesntExist helpers<commit_after>package elasticorm

import (
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"strings"

	"github.com/pkg/errors"
	elastic "gopkg.in/olivere/elastic.v5"
)

// NewDatastore returns a fresh instance of an elasticorm datastore
func NewDatastore(esc *elastic.Client, options ...DatastoreOptFunc) (*Datastore, error) {
	ds := &Datastore{
		elasticClient: esc,
		Ctx:           context.Background(),
	}
	var err error
	for _, opt := range options {
		err = opt(ds)
	}
	return ds, err
}

// DatastoreOptFunc is used as a parameter to NewDatastore and provides a way of configuration
type DatastoreOptFunc func(*Datastore) error

// Datstore is an instance to communicate easy with elasticsearch
// It is meant to be used for one struct and helps storing and retrieving it in/from elasticsearch
// It leverages the great elastic package from olivere
type Datastore struct {
	elasticClient   *elastic.Client
	Ctx             context.Context
	indexName       string
	goType          reflect.Type
	idFieldName     string          // the name of the structs field to store the ID
	typeName        string          // in elasticsearch
	indexDefinition IndexDefinition // in elasticsearch
}

// EnsureIndexExists checks wether the needed index for this datastore exists. It it doesn't it gets created
// the name of the datastore is determined by the structs name (+ plural s)
func (ds *Datastore) EnsureIndexExists() error {
	if ds.indexName == `` {
		return errors.New(`EnsureIndexExists failed, because no index name is defined`)
	}
	exists, err := ds.elasticClient.
		IndexExists(ds.indexName).
		Do(ds.Ctx)

	if exists || err != nil {
		return err
	}

	err = ds.createIndex()
	if err != nil {
		return err
	}
	ds.Refresh()
	return nil
}

func (ds *Datastore) createIndex() error {
	JSON, err := json.Marshal(ds.indexDefinition)
	if err != nil {
		return err
	}
	ack, err := ds.elasticClient.
		CreateIndex(ds.indexName).
		BodyString(string(JSON)).
		Do(ds.Ctx)
	if err != nil || !ack.Acknowledged {
		return errors.Wrap(
			err, fmt.Sprintf("creating elasticsearch index %s failed - %s", ds.indexName, string(JSON)),
		)
	}
	return nil
}

func (ds *Datastore) EnsureIndexDoesntExist() error {
	if ds.indexName == `` {
		return errors.New(`EnsureIndexDoesntExists failed, because no index name is defined`)
	}
	exists, err := ds.elasticClient.
		IndexExists(ds.indexName).
		Do(ds.Ctx)

	if !exists || err != nil {
		return err
	}
	res, err := ds.elasticClient.
		DeleteIndex(ds.indexName).
		Do(context.Background())
	if err != nil || !res.Acknowledged {
		return errors.Wrap(
			err, fmt.Sprintf("deleting elasticsearch index %s failed", ds.indexName),
		)
	}
	return nil
}

func (ds *Datastore) CleanUp() error {
	err := ds.EnsureIndexDoesntExist()
	if err != nil {
		return err
	}
	err = ds.EnsureIndexExists()
	if err != nil {
		return err
	}
	return ds.Refresh()
}

func (ds *Datastore) Refresh() error {
	_, err := ds.elasticClient.Refresh().Do(ds.Ctx)
	return err
}

// ForStruct generates a DatastoreOptFunc
// It is used to generate a default mapping, type name and index name by analysing a provided struct
func ForStruct(i interface{}) DatastoreOptFunc {
	return func(ds *Datastore) error {
		ds.goType = reflect.TypeOf(i)
		typeName, err := typeNameFromStruct(i)
		if err != nil {
			return err
		}
		ds.typeName = typeName
		ds.indexName = typeName + `s`
		indexDefinition, err := NewIndexDefinition(
			AddMappingFromStruct(ds.typeName, i),
		)
		if err != nil {
			return err
		}
		ds.indexDefinition = indexDefinition
		// TODO configure ID
		ds.idFieldName = `ID`
		return nil
	}
}

func typeNameFromStruct(i interface{}) (string, error) {
	typeName := getType(i)
	if typeName == `` {
		return ``, errors.New(`Could not determine type name from struct`)
	}

	typeName = strings.ToLower(typeName)
	typeName = strings.TrimLeft(typeName, `*`)
	return typeName, nil
}

func getType(myvar interface{}) string {
	return nameOfType(reflect.TypeOf(myvar))
}

func nameOfType(t reflect.Type) string {
	if t.Kind() == reflect.Ptr {
		return "*" + t.Elem().Name()
	} else {
		return t.Name()
	}
}

func (ds *Datastore) Create(o interface{}) error {
	err := ds.isSaveableType(o)
	if err != nil {
		return err
	}
	put, err := ds.elasticClient.Index().
		Index(ds.indexName).
		Type(ds.typeName).
		BodyJson(o).
		Do(ds.Ctx)

	if err != nil {
		return err
	}
	if !put.Created {
		return ErrCreationFailed
	}
	return ds.setID(o, put.Id)
}

func (ds *Datastore) isSaveableType(i interface{}) error {
	rt := reflect.TypeOf(i)
	if nameOfType(rt) != nameOfType(ds.goType) {
		return errors.Wrapf(
			ErrInvalidType,
			`create failed for %s, expected %s`, nameOfType(rt), nameOfType(ds.goType),
		)
	}
	rv := reflect.ValueOf(i)
	if rv.Kind() != reflect.Ptr || rv.IsNil() {
		return errors.Wrap(ErrInvalidType, `no pointer given`)
	}
	return nil
}

func (ds *Datastore) Find(ID string, result interface{}) error {
	err := ds.isSaveableType(result)
	if err != nil {
		return err
	}
	res, err := ds.elasticClient.Get().
		Index(ds.indexName).
		Type(ds.typeName).
		Id(ID).
		Do(ds.Ctx)

	if elastic.IsNotFound(err) || (res != nil && !res.Found) {
		return ErrNotFound
	}
	if err != nil {
		return err
	}

	return ds.decodeElasticResponse(res.Source, res.Id, result)
}

func (ds *Datastore) Update(o interface{}) error {
	err := ds.isSaveableType(o)
	if err != nil {
		return err
	}
	ID, err := ds.getID(o)
	if err != nil {
		return err
	}
	if ID == `` {
		return errors.New(`can't save struct with empty ID`)
	}
	_, err = ds.elasticClient.Update().
		Index(ds.indexName).
		Type(ds.typeName).
		Id(ID).
		Doc(o).
		Do(ds.Ctx)

	return err
}

func (ds *Datastore) FindOneBy(fieldName string, value interface{}, result interface{}) error {
	elasticFieldName, err := ds.indexDefinition.elasticFieldName(ds.typeName, fieldName)
	if err != nil {
		return err
	}
	res, err := ds.elasticClient.Search().
		Index(ds.indexName).
		Query(elastic.NewBoolQuery().Filter(elastic.NewTermQuery(elasticFieldName, value))).
		From(0).Size(1).
		Do(ds.Ctx)
	if err != nil {
		return err
	}
	if res.TotalHits() < 1 {
		return ErrNotFound
	}
	return ds.decodeElasticResponse(res.Hits.Hits[0].Source, res.Hits.Hits[0].Id, result)
}

func (ds *Datastore) setID(o interface{}, ID string) error {
	eo := reflect.ValueOf(o).Elem()
	if eo.Kind() != reflect.Struct {
		return errors.Wrap(ErrInvalidType, `setID failed`)
	}
	idField := eo.FieldByName(ds.idFieldName)
	if !idField.IsValid() || !idField.CanSet() || idField.Kind() != reflect.String {
		return ErrInvalidIDField
	}
	idField.SetString(ID)
	return nil
}

func (ds *Datastore) getID(o interface{}) (string, error) {
	eo := reflect.ValueOf(o).Elem()
	if eo.Kind() != reflect.Struct {
		return ``, errors.Wrap(ErrInvalidType, `getID failed`)
	}
	idField := eo.FieldByName(ds.idFieldName)
	if !idField.IsValid() || idField.Kind() != reflect.String {
		return ``, ErrInvalidIDField
	}
	return idField.String(), nil
}

func (ds *Datastore) decodeElasticResponse(source *json.RawMessage, ID string, o interface{}) error {
	err := json.Unmarshal(*source, o)
	if err != nil {
		return err
	}
	ds.setID(o, ID)
	return nil
}
<|endoftext|>


<commit_before>package emailnotifier

import (
	"bytes"
	"fmt"
	"html/template"
	"socialapi/config"
	"socialapi/workers/notification/models"
)

var (
	mainTemplateFile        = "../socialapi/workers/emailnotifier/templates/main.tmpl"
	footerTemplateFile      = "../socialapi/workers/emailnotifier/templates/footer.tmpl"
	contentTemplateFile     = "../socialapi/workers/emailnotifier/templates/content.tmpl"
	gravatarTemplateFile    = "../socialapi/workers/emailnotifier/templates/gravatar.tmpl"
	groupTemplateFile       = "../socialapi/workers/emailnotifier/templates/group.tmpl"
	previewTemplateFile     = "../socialapi/workers/emailnotifier/templates/preview.tmpl"
	objectTemplateFile      = "../socialapi/workers/emailnotifier/templates/object.tmpl"
	unsubscribeTemplateFile = "../socialapi/workers/emailnotifier/templates/unsubscribe.tmpl"
)

type NotificationContainer struct {
	Activity        *models.NotificationActivity
	Content         *models.NotificationContent
	Notification    *models.Notification
	Message         string
	Slug            string
	ActivityMessage string
	ObjectType      string
	Group           GroupContent
}

type MailContent struct {
	TurnOffLink  string
	CurrentDate  string
	FirstName    string
	Description  string
	Size         int
	Uri          string
	ActorContact UserContact
	// event content
	ActivityTime string
	Avatar       string
	Sender       string
	Action       string
	ContentLink  string
	Group        GroupContent
	Message      string
	Slug         string
	ObjectType   string
}

type GroupContent struct {
	Slug string
	Name string
}

func renderTemplate(uc *UserContact, nc *NotificationContainer) (string, error) {
	t := template.Must(template.ParseFiles(
		mainTemplateFile, footerTemplateFile, contentTemplateFile,
		gravatarTemplateFile, unsubscribeTemplateFile))
	mc, err := buildMailContent(uc, nc)
	t = appendPreviewTemplate(t, nc)
	t = appendGroupTemplate(t, nc)

	if err != nil {
		return "", err
	}

	var doc bytes.Buffer
	if err := t.Execute(&doc, mc); err != nil {
		return "", err
	}

	return doc.String(), nil
}

func buildMailContent(uc *UserContact, nc *NotificationContainer) (*MailContent, error) {
	mc := &MailContent{
		FirstName:  uc.FirstName,
		Size:       20,
		Slug:       nc.Slug,
		Message:    nc.Message,
		Action:     nc.ActivityMessage,
		ObjectType: nc.ObjectType,
		Group:      nc.Group,
	}

	mc.CurrentDate, mc.ActivityTime = prepareDateTime(nc)
	mc.Uri = config.Get().Uri

	actor, err := fetchUserContact(nc.Activity.ActorId)
	if err != nil {
		return nil, fmt.Errorf("an error occurred while retrieving actor details", err)
	}
	mc.ActorContact = *actor

	return mc, nil
}

func appendGroupTemplate(t *template.Template, nc *NotificationContainer) *template.Template {
	var groupTemplate *template.Template
	if nc.Group.Name == "" || nc.Group.Slug == "koding" {
		groupTemplate = getEmptyTemplate()
	} else {
		groupTemplate = template.Must(
			template.ParseFiles(groupTemplateFile))
	}

	t.AddParseTree("group", groupTemplate.Tree)

	return t
}

func appendPreviewTemplate(t *template.Template, nc *NotificationContainer) *template.Template {
	var previewTemplate, objectTemplate *template.Template
	if nc.Message == "" {
		previewTemplate = getEmptyTemplate()
		objectTemplate = getEmptyTemplate()
	} else {
		previewTemplate = template.Must(template.ParseFiles(previewTemplateFile))
		objectTemplate = template.Must(template.ParseFiles(objectTemplateFile))
	}

	t.AddParseTree("preview", previewTemplate.Tree)
	t.AddParseTree("object", objectTemplate.Tree)

	return t
}

func getEmptyTemplate() *template.Template {
	return template.Must(template.New("").Parse(""))
}

func prepareDateTime(nc *NotificationContainer) (string, string) {
	layoutDate := "Jan 02"
	layoutTime := "15:04"

	return nc.Notification.ActivatedAt.Format(layoutDate),
		nc.Notification.ActivatedAt.Format(layoutTime)
}

func prepareSubject(nc *NotificationContainer) string {
	t, err := nc.Content.GetContentType()
	if err != nil {
		return ""
	}

	return t.GetDefinition()
}
<commit_msg>Notification: MailContent is updated with escaped recipient email address, token and content type information<commit_after>package emailnotifier

import (
	"bytes"
	"fmt"
	"html/template"
	"net/url"
	"socialapi/config"
	"socialapi/workers/notification/models"
)

var (
	mainTemplateFile        = "../socialapi/workers/emailnotifier/templates/main.tmpl"
	footerTemplateFile      = "../socialapi/workers/emailnotifier/templates/footer.tmpl"
	contentTemplateFile     = "../socialapi/workers/emailnotifier/templates/content.tmpl"
	gravatarTemplateFile    = "../socialapi/workers/emailnotifier/templates/gravatar.tmpl"
	groupTemplateFile       = "../socialapi/workers/emailnotifier/templates/group.tmpl"
	previewTemplateFile     = "../socialapi/workers/emailnotifier/templates/preview.tmpl"
	objectTemplateFile      = "../socialapi/workers/emailnotifier/templates/object.tmpl"
	unsubscribeTemplateFile = "../socialapi/workers/emailnotifier/templates/unsubscribe.tmpl"
)

type NotificationContainer struct {
	Activity        *models.NotificationActivity
	Content         *models.NotificationContent
	Notification    *models.Notification
	Message         string
	Slug            string
	ActivityMessage string
	ObjectType      string
	Group           GroupContent
	Token           string
}

type MailContent struct {
	CurrentDate  string
	FirstName    string
	Description  string
	Size         int
	Uri          string
	ActorContact UserContact
	// event content
	ActivityTime string
	Avatar       string
	Sender       string
	Action       string
	ContentLink  string
	Group        GroupContent
	Message      string
	Slug         string
	ObjectType   string
	Token        string
	ContentType  string
	Recipient    string
}

type GroupContent struct {
	Slug string
	Name string
}

func renderTemplate(uc *UserContact, nc *NotificationContainer) (string, error) {
	t := template.Must(template.ParseFiles(
		mainTemplateFile, footerTemplateFile, contentTemplateFile,
		gravatarTemplateFile, unsubscribeTemplateFile))
	mc, err := buildMailContent(uc, nc)
	t = appendPreviewTemplate(t, nc)
	t = appendGroupTemplate(t, nc)

	if err != nil {
		return "", err
	}

	var doc bytes.Buffer
	if err := t.Execute(&doc, mc); err != nil {
		return "", err
	}

	return doc.String(), nil
}

func buildMailContent(uc *UserContact, nc *NotificationContainer) (*MailContent, error) {
	mc := &MailContent{
		FirstName:   uc.FirstName,
		Size:        20,
		Slug:        nc.Slug,
		Message:     nc.Message,
		Action:      nc.ActivityMessage,
		ObjectType:  nc.ObjectType,
		Group:       nc.Group,
		Token:       nc.Token,
		ContentType: nc.Content.TypeConstant,
		Recipient:   url.QueryEscape(uc.Email),
	}

	mc.CurrentDate, mc.ActivityTime = prepareDateTime(nc)
	mc.Uri = config.Get().Uri

	actor, err := fetchUserContact(nc.Activity.ActorId)
	if err != nil {
		return nil, fmt.Errorf("an error occurred while retrieving actor details", err)
	}
	mc.ActorContact = *actor

	return mc, nil
}

func appendGroupTemplate(t *template.Template, nc *NotificationContainer) *template.Template {
	var groupTemplate *template.Template
	if nc.Group.Name == "" || nc.Group.Slug == "koding" {
		groupTemplate = getEmptyTemplate()
	} else {
		groupTemplate = template.Must(
			template.ParseFiles(groupTemplateFile))
	}

	t.AddParseTree("group", groupTemplate.Tree)

	return t
}

func appendPreviewTemplate(t *template.Template, nc *NotificationContainer) *template.Template {
	var previewTemplate, objectTemplate *template.Template
	if nc.Message == "" {
		previewTemplate = getEmptyTemplate()
		objectTemplate = getEmptyTemplate()
	} else {
		previewTemplate = template.Must(template.ParseFiles(previewTemplateFile))
		objectTemplate = template.Must(template.ParseFiles(objectTemplateFile))
	}

	t.AddParseTree("preview", previewTemplate.Tree)
	t.AddParseTree("object", objectTemplate.Tree)

	return t
}

func getEmptyTemplate() *template.Template {
	return template.Must(template.New("").Parse(""))
}

func prepareDateTime(nc *NotificationContainer) (string, string) {
	layoutDate := "Jan 02"
	layoutTime := "15:04"

	return nc.Notification.ActivatedAt.Format(layoutDate),
		nc.Notification.ActivatedAt.Format(layoutTime)
}

func prepareSubject(nc *NotificationContainer) string {
	t, err := nc.Content.GetContentType()
	if err != nil {
		return ""
	}

	return t.GetDefinition()
}
<|endoftext|>


<commit_before>package db

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/shieldproject/shield/plugin"
)

type Agent struct {
	UUID          string `json:"uuid"            mbus:"uuid"`
	Name          string `json:"name"            mbus:"name"`
	Address       string `json:"address"         mbus:"address"`
	Version       string `json:"version"         mbus:"version"`
	Hidden        bool   `json:"hidden"          mbus:"hidden"`
	LastSeenAt    int64  `json:"last_seen_at"    mbus:"last_seen_at"`
	LastCheckedAt int64  `json:"last_checked_at" mbus:"last_checked_at"`
	LastError     string `json:"last_error"      mbus:"last_error"`
	Status        string `json:"status"          mbus:"status"`

	Meta    map[string]interface{} `json:"metadata,omitempty"`
	RawMeta string                 `json:"-"`
}

func (a *Agent) Metadata() (map[string]interface{}, error) {
	raw := make(map[string]interface{})
	if a.RawMeta == "" {
		return raw, nil
	}
	return raw, json.Unmarshal([]byte(a.RawMeta), &raw)
}

type AgentFilter struct {
	UUID        string
	ExactMatch  bool
	Address     string
	Name        string
	Status      string
	SkipHidden  bool
	SkipVisible bool

	InflateMetadata bool
}

func (f *AgentFilter) Query() (string, []interface{}) {
	wheres := []string{"a.uuid = a.uuid"}
	var args []interface{}

	if f.UUID != "" {
		if f.ExactMatch {
			wheres = []string{"a.uuid = ?"}
			args = []interface{}{f.UUID}
		} else {
			wheres = []string{"a.uuid LIKE ? ESCAPE '/'"}
			args = []interface{}{PatternPrefix(f.UUID)}
		}
	}

	if f.SkipHidden || f.SkipVisible {
		wheres = append(wheres, "a.hidden = ?")
		if f.SkipHidden {
			args = append(args, false)
		} else {
			args = append(args, true)
		}
	}

	if f.Status != "" {
		wheres = append(wheres, "a.status = ?")
		args = append(args, f.Status)
	}

	if f.Address != "" {
		wheres = append(wheres, "a.address = ?")
		args = append(args, f.Address)
	}

	if f.Name != "" {
		wheres = append(wheres, "a.name = ?")
		args = append(args, f.Name)
	}

	return `
	   SELECT a.uuid, a.name, a.address, a.version,
	          a.hidden, a.last_seen_at, a.last_checked_at, a.status,
	          a.metadata, a.last_error

	     FROM agents a

	    WHERE ` + strings.Join(wheres, " AND ") + `
	 ORDER BY a.name DESC, a.uuid ASC`, args
}

func (db *DB) GetAllAgents(filter *AgentFilter) ([]*Agent, error) {
	if filter == nil {
		filter = &AgentFilter{}
	}

	l := []*Agent{}
	query, args := filter.Query()
	db.exclusive.Lock()
	defer db.exclusive.Unlock()
	r, err := db.query(query, args...)
	if err != nil {
		return l, err
	}
	defer r.Close()

	for r.Next() {
		agent := &Agent{}

		var seen, checked *int64
		if err = r.Scan(
			&agent.UUID, &agent.Name, &agent.Address, &agent.Version,
			&agent.Hidden, &seen, &checked, &agent.Status, &agent.RawMeta,
			&agent.LastError); err != nil {
			return l, err
		}
		if seen != nil {
			agent.LastSeenAt = *seen
		}
		if checked != nil {
			agent.LastCheckedAt = *checked
		}

		if filter.InflateMetadata {
			agent.Meta, _ = agent.Metadata()
		}

		l = append(l, agent)
	}

	return l, nil
}

func (db *DB) GetAgent(id string) (*Agent, error) {
	all, err := db.GetAllAgents(&AgentFilter{UUID: id, ExactMatch: true})
	if err != nil {
		return nil, err
	}
	if len(all) == 0 {
		return nil, nil
	}
	return all[0], nil
}

func (db *DB) GetAgentByAddress(address string) (*Agent, error) {
	all, err := db.GetAllAgents(&AgentFilter{Address: address})
	if err != nil {
		return nil, err
	}
	if len(all) == 0 {
		return nil, nil
	}
	return all[0], nil
}

func (db *DB) GetAgentPluginMetadata(addr, name string) (*plugin.PluginInfo, error) {
	db.exclusive.Lock()
	defer db.exclusive.Unlock()
	r, err := db.query(`SELECT metadata FROM agents WHERE address = ?`, addr)
	if err != nil {
		return nil, err
	}
	defer r.Close()

	if !r.Next() {
		return nil, nil
	}

	var (
		raw  []byte
		meta struct {
			Plugins map[string]plugin.PluginInfo `json:"plugins"`
		}
	)
	if err = r.Scan(&raw); err != nil {
		return nil, err
	}
	if err = json.Unmarshal(raw, &meta); err != nil {
		return nil, err
	}

	if p, ok := meta.Plugins[name]; ok {
		return &p, nil
	}
	return nil, nil
}

func (db *DB) PreRegisterAgent(host, name string, port int) error {
	address := fmt.Sprintf("%s:%d", host, port)
	existing, err := db.GetAllAgents(&AgentFilter{
		Name: name,
	})
	if err != nil {
		return err
	}

	if len(existing) > 0 {
		return db.Exec(`
		  UPDATE agents
		     SET address      = ?,
		         last_seen_at = ?
		   WHERE name = ?`,
			address, time.Now().Unix(), name,
		)
	}

	id := RandomID()

	err = db.Exec(`
	   INSERT INTO agents (uuid, name, address, hidden, status, last_seen_at)
	               VALUES (?,    ?,    ?,       ?,      ?,      ?)`,
		id, name, address, false, "pending", time.Now().Unix(),
	)
	if err != nil {
		return err
	}

	agent, err := db.GetAgent(id)
	if err != nil {
		return err
	}
	if agent == nil {
		return fmt.Errorf("failed to retrieve newly-inserted agent [%s]: not found in database.", id)
	}

	db.sendCreateObjectEvent(agent, "*")
	return nil
}

func (db *DB) UpdateAgent(agent *Agent) error {
	return db.Exec(
		`UPDATE agents SET name            = ?,
		                   address         = ?,
		                   version         = ?,
		                   status          = ?,
		                   hidden          = ?,
		                   metadata        = ?,
		                   last_checked_at = ?,
		                   last_seen_at    = ?,
		                   last_error      = ?
		        WHERE uuid = ?`,
		agent.Name, agent.Address, agent.Version, agent.Status, agent.Hidden, agent.RawMeta,
		agent.LastCheckedAt, agent.LastSeenAt, agent.LastError,
		agent.UUID)
}

func (db *DB) DeleteAgent(agent *Agent) error {
	return db.exclusively(func() error {
		n, err := db.count(`SELECT uuid FROM jobs WHERE agent = ?`, agent.Address)
		if err != nil {
			return fmt.Errorf("unable to determine if agent can be deleted: %s", err)
		}
		if n > 0 {
			return fmt.Errorf("agent is still referenced by configured data protection jobs")
		}

		return db.exec(`DELETE FROM agents WHERE uuid = ?`, agent.UUID)
	})
}
<commit_msg>Fix Agent Deletion<commit_after>package db

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/shieldproject/shield/plugin"
)

type Agent struct {
	UUID          string `json:"uuid"            mbus:"uuid"`
	Name          string `json:"name"            mbus:"name"`
	Address       string `json:"address"         mbus:"address"`
	Version       string `json:"version"         mbus:"version"`
	Hidden        bool   `json:"hidden"          mbus:"hidden"`
	LastSeenAt    int64  `json:"last_seen_at"    mbus:"last_seen_at"`
	LastCheckedAt int64  `json:"last_checked_at" mbus:"last_checked_at"`
	LastError     string `json:"last_error"      mbus:"last_error"`
	Status        string `json:"status"          mbus:"status"`

	Meta    map[string]interface{} `json:"metadata,omitempty"`
	RawMeta string                 `json:"-"`
}

func (a *Agent) Metadata() (map[string]interface{}, error) {
	raw := make(map[string]interface{})
	if a.RawMeta == "" {
		return raw, nil
	}
	return raw, json.Unmarshal([]byte(a.RawMeta), &raw)
}

type AgentFilter struct {
	UUID        string
	ExactMatch  bool
	Address     string
	Name        string
	Status      string
	SkipHidden  bool
	SkipVisible bool

	InflateMetadata bool
}

func (f *AgentFilter) Query() (string, []interface{}) {
	wheres := []string{"a.uuid = a.uuid"}
	var args []interface{}

	if f.UUID != "" {
		if f.ExactMatch {
			wheres = []string{"a.uuid = ?"}
			args = []interface{}{f.UUID}
		} else {
			wheres = []string{"a.uuid LIKE ? ESCAPE '/'"}
			args = []interface{}{PatternPrefix(f.UUID)}
		}
	}

	if f.SkipHidden || f.SkipVisible {
		wheres = append(wheres, "a.hidden = ?")
		if f.SkipHidden {
			args = append(args, false)
		} else {
			args = append(args, true)
		}
	}

	if f.Status != "" {
		wheres = append(wheres, "a.status = ?")
		args = append(args, f.Status)
	}

	if f.Address != "" {
		wheres = append(wheres, "a.address = ?")
		args = append(args, f.Address)
	}

	if f.Name != "" {
		wheres = append(wheres, "a.name = ?")
		args = append(args, f.Name)
	}

	return `
	   SELECT a.uuid, a.name, a.address, a.version,
	          a.hidden, a.last_seen_at, a.last_checked_at, a.status,
	          a.metadata, a.last_error

	     FROM agents a

	    WHERE ` + strings.Join(wheres, " AND ") + `
	 ORDER BY a.name DESC, a.uuid ASC`, args
}

func (db *DB) GetAllAgents(filter *AgentFilter) ([]*Agent, error) {
	if filter == nil {
		filter = &AgentFilter{}
	}

	l := []*Agent{}
	query, args := filter.Query()
	db.exclusive.Lock()
	defer db.exclusive.Unlock()
	r, err := db.query(query, args...)
	if err != nil {
		return l, err
	}
	defer r.Close()

	for r.Next() {
		agent := &Agent{}

		var seen, checked *int64
		if err = r.Scan(
			&agent.UUID, &agent.Name, &agent.Address, &agent.Version,
			&agent.Hidden, &seen, &checked, &agent.Status, &agent.RawMeta,
			&agent.LastError); err != nil {
			return l, err
		}
		if seen != nil {
			agent.LastSeenAt = *seen
		}
		if checked != nil {
			agent.LastCheckedAt = *checked
		}

		if filter.InflateMetadata {
			agent.Meta, _ = agent.Metadata()
		}

		l = append(l, agent)
	}

	return l, nil
}

func (db *DB) GetAgent(id string) (*Agent, error) {
	all, err := db.GetAllAgents(&AgentFilter{UUID: id, ExactMatch: true})
	if err != nil {
		return nil, err
	}
	if len(all) == 0 {
		return nil, nil
	}
	return all[0], nil
}

func (db *DB) GetAgentByAddress(address string) (*Agent, error) {
	all, err := db.GetAllAgents(&AgentFilter{Address: address})
	if err != nil {
		return nil, err
	}
	if len(all) == 0 {
		return nil, nil
	}
	return all[0], nil
}

func (db *DB) GetAgentPluginMetadata(addr, name string) (*plugin.PluginInfo, error) {
	db.exclusive.Lock()
	defer db.exclusive.Unlock()
	r, err := db.query(`SELECT metadata FROM agents WHERE address = ?`, addr)
	if err != nil {
		return nil, err
	}
	defer r.Close()

	if !r.Next() {
		return nil, nil
	}

	var (
		raw  []byte
		meta struct {
			Plugins map[string]plugin.PluginInfo `json:"plugins"`
		}
	)
	if err = r.Scan(&raw); err != nil {
		return nil, err
	}
	if err = json.Unmarshal(raw, &meta); err != nil {
		return nil, err
	}

	if p, ok := meta.Plugins[name]; ok {
		return &p, nil
	}
	return nil, nil
}

func (db *DB) PreRegisterAgent(host, name string, port int) error {
	address := fmt.Sprintf("%s:%d", host, port)
	existing, err := db.GetAllAgents(&AgentFilter{
		Name: name,
	})
	if err != nil {
		return err
	}

	if len(existing) > 0 {
		return db.Exec(`
		  UPDATE agents
		     SET address      = ?,
		         last_seen_at = ?
		   WHERE name = ?`,
			address, time.Now().Unix(), name,
		)
	}

	id := RandomID()

	err = db.Exec(`
	   INSERT INTO agents (uuid, name, address, hidden, status, last_seen_at)
	               VALUES (?,    ?,    ?,       ?,      ?,      ?)`,
		id, name, address, false, "pending", time.Now().Unix(),
	)
	if err != nil {
		return err
	}

	agent, err := db.GetAgent(id)
	if err != nil {
		return err
	}
	if agent == nil {
		return fmt.Errorf("failed to retrieve newly-inserted agent [%s]: not found in database.", id)
	}

	db.sendCreateObjectEvent(agent, "*")
	return nil
}

func (db *DB) UpdateAgent(agent *Agent) error {
	return db.Exec(
		`UPDATE agents SET name            = ?,
		                   address         = ?,
		                   version         = ?,
		                   status          = ?,
		                   hidden          = ?,
		                   metadata        = ?,
		                   last_checked_at = ?,
		                   last_seen_at    = ?,
		                   last_error      = ?
		        WHERE uuid = ?`,
		agent.Name, agent.Address, agent.Version, agent.Status, agent.Hidden, agent.RawMeta,
		agent.LastCheckedAt, agent.LastSeenAt, agent.LastError,
		agent.UUID)
}

func (db *DB) DeleteAgent(agent *Agent) error {
	return db.exclusively(func() error {
		n, err := db.count(`SELECT uuid FROM targets WHERE agent = ?`, agent.Address)
		if err != nil {
			return fmt.Errorf("unable to determine if agent can be deleted: %s", err)
		}
		if n > 0 {
			return fmt.Errorf("agent is still referenced by configured data systems")
		}

		n, err = db.count(`SELECT uuid FROM stores WHERE agent = ?`, agent.Address)
		if err != nil {
			return fmt.Errorf("unable to determine if agent can be deleted: %s", err)
		}
		if n > 0 {
			return fmt.Errorf("agent is still referenced by configured storage systems")
		}

		return db.exec(`DELETE FROM agents WHERE uuid = ?`, agent.UUID)
	})
}
<|endoftext|>


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

// The virtualbox provider fetches the configuration from the /Ignition/Config
// guest property.

package virtualbox

// We want at least this warning, since the default C behavior of
// assuming int foo(int) is totally broken.

// #cgo CFLAGS: -Werror=implicit-function-declaration
// #include <linux/vbox_err.h>
// #include <stdlib.h>
// #include "virtualbox.h"
import "C"

import (
	"bytes"
	"fmt"
	"unsafe"

	"github.com/coreos/ignition/v2/config/shared/errors"
	"github.com/coreos/ignition/v2/config/v3_4_experimental/types"
	"github.com/coreos/ignition/v2/internal/providers/util"
	"github.com/coreos/ignition/v2/internal/resource"

	"github.com/coreos/vcontext/report"
)

const (
	configProperty         = "/Ignition/Config"
	configEncodingProperty = "/Ignition/Config/Encoding"
)

func FetchConfig(f *resource.Fetcher) (types.Config, report.Report, error) {
	f.Logger.Debug("reading Ignition config from VirtualBox guest property")

	// for forward compatibility, check an encoding property analogous
	// to vmware's ignition.config.data.encoding, and fail if it's
	// present and non-empty
	encoding, err := fetchProperty(configEncodingProperty)
	if err != nil {
		return types.Config{}, report.Report{}, err
	}
	if len(encoding) > 0 {
		return types.Config{}, report.Report{}, fmt.Errorf("unsupported %q value %q", configEncodingProperty, encoding)
	}

	config, err := fetchProperty(configProperty)
	if err != nil {
		return types.Config{}, report.Report{}, err
	}
	if config == nil {
		f.Logger.Info("VirtualBox guest property %q does not exist; assuming no config", configProperty)
		return types.Config{}, report.Report{}, errors.ErrEmpty
	}
	return util.ParseConfig(f.Logger, config)
}

func fetchProperty(name string) ([]byte, error) {
	cName := C.CString(name)
	defer C.free(unsafe.Pointer(cName))
	buf := unsafe.Pointer(nil)
	defer C.free(buf)
	var size C.size_t

	ret, errno := C.virtualbox_get_guest_property(cName, &buf, &size)
	if ret != C.VINF_SUCCESS {
		if ret == C.VERR_GENERAL_FAILURE && errno != nil {
			return nil, fmt.Errorf("fetching VirtualBox guest property %q: %w", name, errno)
		}
		// see <linux/vbox_err.h>
		return nil, fmt.Errorf("fetching VirtualBox guest property %q: error %d", name, ret)
	}
	if buf == nil {
		return nil, nil
	}
	// properties are double-NUL-terminated
	return bytes.TrimRight(C.GoBytes(buf, C.int(size)), "\x00"), nil
}
<commit_msg>providers/virtualbox: fix reading properties with flags<commit_after>// Copyright 2021 Red Hat, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// The virtualbox provider fetches the configuration from the /Ignition/Config
// guest property.

package virtualbox

// We want at least this warning, since the default C behavior of
// assuming int foo(int) is totally broken.

// #cgo CFLAGS: -Werror=implicit-function-declaration
// #include <linux/vbox_err.h>
// #include <stdlib.h>
// #include "virtualbox.h"
import "C"

import (
	"bytes"
	"fmt"
	"unsafe"

	"github.com/coreos/ignition/v2/config/shared/errors"
	"github.com/coreos/ignition/v2/config/v3_4_experimental/types"
	"github.com/coreos/ignition/v2/internal/providers/util"
	"github.com/coreos/ignition/v2/internal/resource"

	"github.com/coreos/vcontext/report"
)

const (
	configProperty         = "/Ignition/Config"
	configEncodingProperty = "/Ignition/Config/Encoding"
)

func FetchConfig(f *resource.Fetcher) (types.Config, report.Report, error) {
	f.Logger.Debug("reading Ignition config from VirtualBox guest property")

	// for forward compatibility, check an encoding property analogous
	// to vmware's ignition.config.data.encoding, and fail if it's
	// present and non-empty
	encoding, err := fetchProperty(configEncodingProperty)
	if err != nil {
		return types.Config{}, report.Report{}, err
	}
	if len(encoding) > 0 {
		return types.Config{}, report.Report{}, fmt.Errorf("unsupported %q value %q", configEncodingProperty, encoding)
	}

	config, err := fetchProperty(configProperty)
	if err != nil {
		return types.Config{}, report.Report{}, err
	}
	if config == nil {
		f.Logger.Info("VirtualBox guest property %q does not exist; assuming no config", configProperty)
		return types.Config{}, report.Report{}, errors.ErrEmpty
	}
	return util.ParseConfig(f.Logger, config)
}

func fetchProperty(name string) ([]byte, error) {
	cName := C.CString(name)
	defer C.free(unsafe.Pointer(cName))
	buf := unsafe.Pointer(nil)
	defer C.free(buf)
	var size C.size_t

	ret, errno := C.virtualbox_get_guest_property(cName, &buf, &size)
	if ret != C.VINF_SUCCESS {
		if ret == C.VERR_GENERAL_FAILURE && errno != nil {
			return nil, fmt.Errorf("fetching VirtualBox guest property %q: %w", name, errno)
		}
		// see <linux/vbox_err.h>
		return nil, fmt.Errorf("fetching VirtualBox guest property %q: error %d", name, ret)
	}
	if buf == nil {
		return nil, nil
	}
	// property format: <data> NUL <flags> NUL
	// return only the data; ignore the flags
	s := C.GoBytes(buf, C.int(size))
	len := bytes.IndexByte(s, 0)
	if len == -1 {
		return nil, fmt.Errorf("VirtualBox guest property %q is not NUL-terminated", name)
	}
	return s[0:len], nil
}
<|endoftext|>