<commit_before>package cluster

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

	"github.com/pkg/errors"

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

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

			if ok {
				connected = true
				break
			}

			time.Sleep(500 * time.Millisecond)
		}

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

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

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

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

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

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

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

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

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

	return Connect(address, cert, false)
}

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

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

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

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

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

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

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

	return nil
}

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

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

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

	"github.com/pkg/errors"

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

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

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

			if ok {
				connected = true
				break
			}

			time.Sleep(500 * time.Millisecond)
		}

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

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

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

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

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

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

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

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

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

	return Connect(address, cert, false)
}

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

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

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

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

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

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

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

	return nil
}

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

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


<commit_before>// +build linux,cgo,!agent

package db

import (
	"database/sql"
	"encoding/json"
	"fmt"

	"github.com/pkg/errors"

	"github.com/lxc/lxd/shared/api"
)

// GetNetworkACLs returns the names of existing Network ACLs.
func (c *Cluster) GetNetworkACLs(project string) ([]string, error) {
	q := `SELECT name FROM networks_acls
		WHERE project_id = (SELECT id FROM projects WHERE name = ? LIMIT 1)
		ORDER BY id
	`
	inargs := []interface{}{project}

	var name string
	outfmt := []interface{}{name}
	result, err := queryScan(c, q, inargs, outfmt)
	if err != nil {
		return nil, err
	}

	response := []string{}
	for _, r := range result {
		response = append(response, r[0].(string))
	}

	return response, nil
}

// GetNetworkACL returns the Network ACL with the given name in the given project.
func (c *Cluster) GetNetworkACL(projectName string, name string) (int64, *api.NetworkACL, error) {
	var id int64 = int64(-1)
	var ingressJSON string
	var egressJSON string

	acl := api.NetworkACL{
		NetworkACLPost: api.NetworkACLPost{
			Name: name,
		},
	}

	q := `
		SELECT id, description, ingress, egress
		FROM networks_acls
		WHERE project_id = (SELECT id FROM projects WHERE name = ? LIMIT 1) AND name=?
		LIMIT 1
	`
	arg1 := []interface{}{projectName, name}
	arg2 := []interface{}{&id, &acl.Description, &ingressJSON, &egressJSON}

	err := dbQueryRowScan(c, q, arg1, arg2)
	if err != nil {
		if err == sql.ErrNoRows {
			return -1, nil, ErrNoSuchObject
		}

		return -1, nil, err
	}

	acl.Ingress = []api.NetworkACLRule{}
	if ingressJSON != "" {
		err = json.Unmarshal([]byte(ingressJSON), &acl.Ingress)
		if err != nil {
			return -1, nil, errors.Wrapf(err, "Failed unmarshalling ingress rules")
		}
	}

	acl.Egress = []api.NetworkACLRule{}
	if egressJSON != "" {
		err = json.Unmarshal([]byte(egressJSON), &acl.Egress)
		if err != nil {
			return -1, nil, errors.Wrapf(err, "Failed unmarshalling egress rules")
		}
	}

	acl.Config, err = c.networkACLConfig(id)
	if err != nil {
		return -1, nil, errors.Wrapf(err, "Failed loading config")
	}

	return id, &acl, nil
}

// networkACLConfig returns the config map of the Network ACL with the given ID.
func (c *Cluster) networkACLConfig(id int64) (map[string]string, error) {
	var key, value string
	query := `
		SELECT key, value
		FROM networks_acls_config
		WHERE network_acl_id=?
	`
	inargs := []interface{}{id}
	outfmt := []interface{}{key, value}
	results, err := queryScan(c, query, inargs, outfmt)
	if err != nil {
		return nil, err
	}

	config := make(map[string]string, len(results))

	for _, r := range results {
		key = r[0].(string)
		value = r[1].(string)

		_, found := config[key]
		if found {
			return nil, fmt.Errorf("Duplicate config row found for key %q for network ACL ID %d", key, id)
		}

		config[key] = value
	}

	return config, nil
}

// CreateNetworkACL creates a new Network ACL.
func (c *Cluster) CreateNetworkACL(projectName string, info *api.NetworkACLsPost) (int64, error) {
	var id int64
	var err error
	var ingressJSON, egressJSON []byte

	if info.Ingress != nil {
		ingressJSON, err = json.Marshal(info.Ingress)
		if err != nil {
			return -1, errors.Wrapf(err, "Failed marshalling ingress rules")
		}
	}

	if info.Egress != nil {
		egressJSON, err = json.Marshal(info.Egress)
		if err != nil {
			return -1, errors.Wrapf(err, "Failed marshalling egress rules")
		}
	}

	err = c.Transaction(func(tx *ClusterTx) error {
		// Insert a new Network ACL record.
		result, err := tx.tx.Exec(`
			INSERT INTO networks_acls (project_id, name, description, ingress, egress)
			VALUES ((SELECT id FROM projects WHERE name = ? LIMIT 1), ?, ?, ?, ?)
		`, projectName, info.Name, info.Description, string(ingressJSON), string(egressJSON))
		if err != nil {
			return err
		}

		id, err := result.LastInsertId()
		if err != nil {
			return err
		}

		err = networkACLConfigAdd(tx.tx, id, info.Config)
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		id = -1
	}

	return id, err
}

// networkACLConfigAdd inserts Network ACL config keys.
func networkACLConfigAdd(tx *sql.Tx, id int64, config map[string]string) error {
	sql := "INSERT INTO networks_acls_config (network_acl_id, key, value) VALUES(?, ?, ?)"
	stmt, err := tx.Prepare(sql)
	if err != nil {
		return err
	}
	defer stmt.Close()

	for k, v := range config {
		if v == "" {
			continue
		}

		_, err = stmt.Exec(id, k, v)
		if err != nil {
			return errors.Wrapf(err, "Failed inserting config")
		}
	}

	return nil
}

// UpdateNetworkACL updates the Network ACL with the given ID.
func (c *Cluster) UpdateNetworkACL(id int64, config *api.NetworkACLPut) error {
	var err error
	var ingressJSON, egressJSON []byte

	if config.Ingress != nil {
		ingressJSON, err = json.Marshal(config.Ingress)
		if err != nil {
			return errors.Wrapf(err, "Failed marshalling ingress rules")
		}
	}

	if config.Egress != nil {
		egressJSON, err = json.Marshal(config.Egress)
		if err != nil {
			return errors.Wrapf(err, "Failed marshalling egress rules")
		}
	}

	return c.Transaction(func(tx *ClusterTx) error {
		_, err := tx.tx.Exec(`
			UPDATE networks_acls
			SET description=?, ingress = ?, egress = ?
			WHERE id=?
		`, config.Description, ingressJSON, egressJSON, id)
		if err != nil {
			return err
		}

		err = networkACLConfigUpdate(tx.tx, id, config.Config)
		if err != nil {
			return err
		}

		return nil
	})
}

// networkACLConfigUpdate updates Network ACL config keys.
func networkACLConfigUpdate(tx *sql.Tx, id int64, config map[string]string) error {
	_, err := tx.Exec("DELETE FROM networks_acls_config WHERE network_acl_id=?", id)
	if err != nil {
		return err
	}

	str := fmt.Sprintf("INSERT INTO networks_acls_config (network_acl_id, key, value) VALUES(?, ?, ?)")
	stmt, err := tx.Prepare(str)
	if err != nil {
		return err
	}
	defer stmt.Close()

	for k, v := range config {
		if v == "" {
			continue
		}

		_, err = stmt.Exec(id, k, v)
		if err != nil {
			return err
		}
	}

	return nil
}

// RenameNetworkACL renames a Network ACL.
func (c *Cluster) RenameNetworkACL(id int64, newName string) error {
	return c.Transaction(func(tx *ClusterTx) error {
		_, err := tx.tx.Exec("UPDATE networks_acls SET name=? WHERE id=?", newName, id)
		return err
	})
}

// DeleteNetworkACL deletes the Network ACL.
func (c *Cluster) DeleteNetworkACL(id int64) error {
	return c.Transaction(func(tx *ClusterTx) error {
		_, err := tx.tx.Exec("DELETE FROM networks_acls WHERE id=?", id)
		return err
	})
}
<commit_msg>lxd/db/network/acls: Makes slice allocation more efficient in GetNetworkACLs<commit_after>// +build linux,cgo,!agent

package db

import (
	"database/sql"
	"encoding/json"
	"fmt"

	"github.com/pkg/errors"

	"github.com/lxc/lxd/shared/api"
)

// GetNetworkACLs returns the names of existing Network ACLs.
func (c *Cluster) GetNetworkACLs(project string) ([]string, error) {
	q := `SELECT name FROM networks_acls
		WHERE project_id = (SELECT id FROM projects WHERE name = ? LIMIT 1)
		ORDER BY id
	`
	inargs := []interface{}{project}

	var name string
	outfmt := []interface{}{name}
	result, err := queryScan(c, q, inargs, outfmt)
	if err != nil {
		return nil, err
	}

	response := make([]string, 0, len(result))
	for _, r := range result {
		response = append(response, r[0].(string))
	}

	return response, nil
}

// GetNetworkACL returns the Network ACL with the given name in the given project.
func (c *Cluster) GetNetworkACL(projectName string, name string) (int64, *api.NetworkACL, error) {
	var id int64 = int64(-1)
	var ingressJSON string
	var egressJSON string

	acl := api.NetworkACL{
		NetworkACLPost: api.NetworkACLPost{
			Name: name,
		},
	}

	q := `
		SELECT id, description, ingress, egress
		FROM networks_acls
		WHERE project_id = (SELECT id FROM projects WHERE name = ? LIMIT 1) AND name=?
		LIMIT 1
	`
	arg1 := []interface{}{projectName, name}
	arg2 := []interface{}{&id, &acl.Description, &ingressJSON, &egressJSON}

	err := dbQueryRowScan(c, q, arg1, arg2)
	if err != nil {
		if err == sql.ErrNoRows {
			return -1, nil, ErrNoSuchObject
		}

		return -1, nil, err
	}

	acl.Ingress = []api.NetworkACLRule{}
	if ingressJSON != "" {
		err = json.Unmarshal([]byte(ingressJSON), &acl.Ingress)
		if err != nil {
			return -1, nil, errors.Wrapf(err, "Failed unmarshalling ingress rules")
		}
	}

	acl.Egress = []api.NetworkACLRule{}
	if egressJSON != "" {
		err = json.Unmarshal([]byte(egressJSON), &acl.Egress)
		if err != nil {
			return -1, nil, errors.Wrapf(err, "Failed unmarshalling egress rules")
		}
	}

	acl.Config, err = c.networkACLConfig(id)
	if err != nil {
		return -1, nil, errors.Wrapf(err, "Failed loading config")
	}

	return id, &acl, nil
}

// networkACLConfig returns the config map of the Network ACL with the given ID.
func (c *Cluster) networkACLConfig(id int64) (map[string]string, error) {
	var key, value string
	query := `
		SELECT key, value
		FROM networks_acls_config
		WHERE network_acl_id=?
	`
	inargs := []interface{}{id}
	outfmt := []interface{}{key, value}
	results, err := queryScan(c, query, inargs, outfmt)
	if err != nil {
		return nil, err
	}

	config := make(map[string]string, len(results))

	for _, r := range results {
		key = r[0].(string)
		value = r[1].(string)

		_, found := config[key]
		if found {
			return nil, fmt.Errorf("Duplicate config row found for key %q for network ACL ID %d", key, id)
		}

		config[key] = value
	}

	return config, nil
}

// CreateNetworkACL creates a new Network ACL.
func (c *Cluster) CreateNetworkACL(projectName string, info *api.NetworkACLsPost) (int64, error) {
	var id int64
	var err error
	var ingressJSON, egressJSON []byte

	if info.Ingress != nil {
		ingressJSON, err = json.Marshal(info.Ingress)
		if err != nil {
			return -1, errors.Wrapf(err, "Failed marshalling ingress rules")
		}
	}

	if info.Egress != nil {
		egressJSON, err = json.Marshal(info.Egress)
		if err != nil {
			return -1, errors.Wrapf(err, "Failed marshalling egress rules")
		}
	}

	err = c.Transaction(func(tx *ClusterTx) error {
		// Insert a new Network ACL record.
		result, err := tx.tx.Exec(`
			INSERT INTO networks_acls (project_id, name, description, ingress, egress)
			VALUES ((SELECT id FROM projects WHERE name = ? LIMIT 1), ?, ?, ?, ?)
		`, projectName, info.Name, info.Description, string(ingressJSON), string(egressJSON))
		if err != nil {
			return err
		}

		id, err := result.LastInsertId()
		if err != nil {
			return err
		}

		err = networkACLConfigAdd(tx.tx, id, info.Config)
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		id = -1
	}

	return id, err
}

// networkACLConfigAdd inserts Network ACL config keys.
func networkACLConfigAdd(tx *sql.Tx, id int64, config map[string]string) error {
	sql := "INSERT INTO networks_acls_config (network_acl_id, key, value) VALUES(?, ?, ?)"
	stmt, err := tx.Prepare(sql)
	if err != nil {
		return err
	}
	defer stmt.Close()

	for k, v := range config {
		if v == "" {
			continue
		}

		_, err = stmt.Exec(id, k, v)
		if err != nil {
			return errors.Wrapf(err, "Failed inserting config")
		}
	}

	return nil
}

// UpdateNetworkACL updates the Network ACL with the given ID.
func (c *Cluster) UpdateNetworkACL(id int64, config *api.NetworkACLPut) error {
	var err error
	var ingressJSON, egressJSON []byte

	if config.Ingress != nil {
		ingressJSON, err = json.Marshal(config.Ingress)
		if err != nil {
			return errors.Wrapf(err, "Failed marshalling ingress rules")
		}
	}

	if config.Egress != nil {
		egressJSON, err = json.Marshal(config.Egress)
		if err != nil {
			return errors.Wrapf(err, "Failed marshalling egress rules")
		}
	}

	return c.Transaction(func(tx *ClusterTx) error {
		_, err := tx.tx.Exec(`
			UPDATE networks_acls
			SET description=?, ingress = ?, egress = ?
			WHERE id=?
		`, config.Description, ingressJSON, egressJSON, id)
		if err != nil {
			return err
		}

		err = networkACLConfigUpdate(tx.tx, id, config.Config)
		if err != nil {
			return err
		}

		return nil
	})
}

// networkACLConfigUpdate updates Network ACL config keys.
func networkACLConfigUpdate(tx *sql.Tx, id int64, config map[string]string) error {
	_, err := tx.Exec("DELETE FROM networks_acls_config WHERE network_acl_id=?", id)
	if err != nil {
		return err
	}

	str := fmt.Sprintf("INSERT INTO networks_acls_config (network_acl_id, key, value) VALUES(?, ?, ?)")
	stmt, err := tx.Prepare(str)
	if err != nil {
		return err
	}
	defer stmt.Close()

	for k, v := range config {
		if v == "" {
			continue
		}

		_, err = stmt.Exec(id, k, v)
		if err != nil {
			return err
		}
	}

	return nil
}

// RenameNetworkACL renames a Network ACL.
func (c *Cluster) RenameNetworkACL(id int64, newName string) error {
	return c.Transaction(func(tx *ClusterTx) error {
		_, err := tx.tx.Exec("UPDATE networks_acls SET name=? WHERE id=?", newName, id)
		return err
	})
}

// DeleteNetworkACL deletes the Network ACL.
func (c *Cluster) DeleteNetworkACL(id int64) error {
	return c.Transaction(func(tx *ClusterTx) error {
		_, err := tx.tx.Exec("DELETE FROM networks_acls WHERE id=?", id)
		return err
	})
}
<|endoftext|>


<commit_before>package gorocksdb

import (
	"testing"

	"github.com/facebookgo/ensure"
)

func TestMergeOperator(t *testing.T) {
	var (
		givenKey    = []byte("hello")
		givenVal1   = []byte("foo")
		givenVal2   = []byte("bar")
		givenMerged = []byte("foobar")
	)
	merger := &mockMergeOperator{
		fullMerge: func(key, existingValue []byte, operands [][]byte) ([]byte, bool) {
			ensure.DeepEqual(&fatalAsError{t}, key, givenKey)
			ensure.DeepEqual(&fatalAsError{t}, existingValue, givenVal1)
			ensure.DeepEqual(&fatalAsError{t}, operands, [][]byte{givenVal2})
			return givenMerged, true
		},
	}
	db := newTestDB(t, "TestMergeOperator", func(opts *Options) {
		opts.SetMergeOperator(merger)
	})
	defer db.Close()

	wo := NewDefaultWriteOptions()
	ensure.Nil(t, db.Put(wo, givenKey, givenVal1))
	ensure.Nil(t, db.Merge(wo, givenKey, givenVal2))

	ro := NewDefaultReadOptions()
	v1, err := db.Get(ro, givenKey)
	defer v1.Free()
	ensure.Nil(t, err)
	ensure.DeepEqual(t, v1.Data(), givenMerged)
}

type mockMergeOperator struct {
	fullMerge    func(key, existingValue []byte, operands [][]byte) ([]byte, bool)
	partialMerge func(key, leftOperand, rightOperand []byte) ([]byte, bool)
}

func (m *mockMergeOperator) Name() string { return "gorocksdb.test" }
func (m *mockMergeOperator) FullMerge(key, existingValue []byte, operands [][]byte) ([]byte, bool) {
	return m.fullMerge(key, existingValue, operands)
}
func (m *mockMergeOperator) PartialMerge(key, leftOperand, rightOperand []byte) ([]byte, bool) {
	return m.partialMerge(key, leftOperand, rightOperand)
}
<commit_msg>trigger compaction in test to ensure that the merge operator is executed<commit_after>package gorocksdb

import (
	"testing"

	"github.com/facebookgo/ensure"
)

func TestMergeOperator(t *testing.T) {
	var (
		givenKey    = []byte("hello")
		givenVal1   = []byte("foo")
		givenVal2   = []byte("bar")
		givenMerged = []byte("foobar")
	)
	merger := &mockMergeOperator{
		fullMerge: func(key, existingValue []byte, operands [][]byte) ([]byte, bool) {
			ensure.DeepEqual(&fatalAsError{t}, key, givenKey)
			ensure.DeepEqual(&fatalAsError{t}, existingValue, givenVal1)
			ensure.DeepEqual(&fatalAsError{t}, operands, [][]byte{givenVal2})
			return givenMerged, true
		},
	}
	db := newTestDB(t, "TestMergeOperator", func(opts *Options) {
		opts.SetMergeOperator(merger)
	})
	defer db.Close()

	wo := NewDefaultWriteOptions()
	ensure.Nil(t, db.Put(wo, givenKey, givenVal1))
	ensure.Nil(t, db.Merge(wo, givenKey, givenVal2))

    // trigger a compaction to ensure that a merge is performed
	db.CompactRange(Range{nil, nil})

	ro := NewDefaultReadOptions()
	v1, err := db.Get(ro, givenKey)
	defer v1.Free()
	ensure.Nil(t, err)
	ensure.DeepEqual(t, v1.Data(), givenMerged)
}

type mockMergeOperator struct {
	fullMerge    func(key, existingValue []byte, operands [][]byte) ([]byte, bool)
	partialMerge func(key, leftOperand, rightOperand []byte) ([]byte, bool)
}

func (m *mockMergeOperator) Name() string { return "gorocksdb.test" }
func (m *mockMergeOperator) FullMerge(key, existingValue []byte, operands [][]byte) ([]byte, bool) {
	return m.fullMerge(key, existingValue, operands)
}
func (m *mockMergeOperator) PartialMerge(key, leftOperand, rightOperand []byte) ([]byte, bool) {
	return m.partialMerge(key, leftOperand, rightOperand)
}
<|endoftext|>


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

package core

// Definition of labels supported in MetricSet.

var (
	LabelMetricSetType = LabelDescriptor{
		Key:         "type",
		Description: "Type of the metrics set (container, pod, namespace, node, cluster)",
	}
	MetricSetTypeSystemContainer = "sys_container"
	MetricSetTypePodContainer    = "pod_container"
	MetricSetTypePod             = "pod"
	MetricSetTypeNamespace       = "ns"
	MetricSetTypeNode            = "node"
	MetricSetTypeCluster         = "cluster"

	LabelPodId = LabelDescriptor{
		Key:         "pod_id",
		Description: "The unique ID of the pod",
	}
	LabelPodName = LabelDescriptor{
		Key:         "pod_name",
		Description: "The name of the pod",
	}
	// Deprecated label
	LabelPodNamespace = LabelDescriptor{
		Key:         "pod_namespace",
		Description: "The namespace of the pod",
	}
	LabelNamespaceName = LabelDescriptor{
		Key:         "namespace_name",
		Description: "The name of the namespace",
	}
	LabelPodNamespaceUID = LabelDescriptor{
		Key:         "namespace_id",
		Description: "The UID of namespace of the pod",
	}
	LabelContainerName = LabelDescriptor{
		Key:         "container_name",
		Description: "User-provided name of the container or full container name for system containers",
	}
	LabelLabels = LabelDescriptor{
		Key:         "labels",
		Description: "Comma-separated list of user-provided labels",
	}
	LabelNodename = LabelDescriptor{
		Key:         "nodename",
		Description: "nodename where the container ran",
	}
	LabelHostname = LabelDescriptor{
		Key:         "hostname",
		Description: "Hostname where the container ran",
	}
	LabelResourceID = LabelDescriptor{
		Key:         "resource_id",
		Description: "Identifier(s) specific to a metric",
	}
	LabelHostID = LabelDescriptor{
		Key:         "host_id",
		Description: "Identifier specific to a host. Set by cloud provider or user",
	}
	LabelContainerBaseImage = LabelDescriptor{
		Key:         "container_base_image",
		Description: "User-defined image name that is run inside the container",
	}
	// The label is populated only for GCM
	LabelCustomMetricName = LabelDescriptor{
		Key:         "custom_metric_name",
		Description: "User-defined name of the exported custom metric",
	}
	LabelGCEResourceID = LabelDescriptor{
		Key:         "compute.googleapis.com/resource_id",
		Description: "Resource id for nodes specific for GCE.",
	}
	LabelGCEResourceType = LabelDescriptor{
		Key:         "compute.googleapis.com/resource_type",
		Description: "Resource types for nodes specific for GCE.",
	}
)

type LabelDescriptor struct {
	// Key to use for the label.
	Key string `json:"key,omitempty"`

	// Description of the label.
	Description string `json:"description,omitempty"`
}

var commonLabels = []LabelDescriptor{
	LabelNodename,
	LabelHostname,
	LabelHostID,
}

var containerLabels = []LabelDescriptor{
	LabelContainerName,
	LabelContainerBaseImage,
}

var podLabels = []LabelDescriptor{
	LabelPodName,
	LabelPodId,
	LabelPodNamespace,
	// TODO: Enable label once it is supported by NamespaceEnricher
	//	LabelPodNamespaceUID,
	LabelLabels,
}

var metricLabels = []LabelDescriptor{
	LabelResourceID,
}

var customMetricLabels = []LabelDescriptor{
	LabelCustomMetricName,
}

// Labels exported to GCM. The number of labels that can be exported to GCM is limited by 10.
var gcmLabels = []LabelDescriptor{
	LabelMetricSetType,
	LabelPodName,
	LabelNamespaceName,
	LabelHostname,
	LabelHostID,
	LabelContainerName,
	LabelContainerBaseImage,
	LabelCustomMetricName,
}

var gcmNodeAutoscalingLabels = []LabelDescriptor{
	LabelGCEResourceID,
	LabelGCEResourceType,
	LabelHostname,
}

func CommonLabels() []LabelDescriptor {
	result := make([]LabelDescriptor, len(commonLabels))
	copy(result, commonLabels)
	return result
}

func ContainerLabels() []LabelDescriptor {
	result := make([]LabelDescriptor, len(containerLabels))
	copy(result, containerLabels)
	return result
}

func PodLabels() []LabelDescriptor {
	result := make([]LabelDescriptor, len(podLabels))
	copy(result, podLabels)
	return result
}

func MetricLabels() []LabelDescriptor {
	result := make([]LabelDescriptor, len(metricLabels)+len(customMetricLabels))
	copy(result, metricLabels)
	copy(result, customMetricLabels)
	return result
}

func SupportedLabels() []LabelDescriptor {
	result := CommonLabels()
	result = append(result, PodLabels()...)
	return append(result, MetricLabels()...)
}

func GcmLabels() map[string]LabelDescriptor {
	result := make(map[string]LabelDescriptor, len(gcmLabels))
	for _, l := range gcmLabels {
		result[l.Key] = l
	}
	return result
}
func GcmNodeAutoscalingLabels() map[string]LabelDescriptor {
	result := make(map[string]LabelDescriptor, len(gcmNodeAutoscalingLabels))
	for _, l := range gcmNodeAutoscalingLabels {
		result[l.Key] = l
	}
	return result
}
<commit_msg>Enable namespace_id for GKE<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package core

// Definition of labels supported in MetricSet.

var (
	LabelMetricSetType = LabelDescriptor{
		Key:         "type",
		Description: "Type of the metrics set (container, pod, namespace, node, cluster)",
	}
	MetricSetTypeSystemContainer = "sys_container"
	MetricSetTypePodContainer    = "pod_container"
	MetricSetTypePod             = "pod"
	MetricSetTypeNamespace       = "ns"
	MetricSetTypeNode            = "node"
	MetricSetTypeCluster         = "cluster"

	LabelPodId = LabelDescriptor{
		Key:         "pod_id",
		Description: "The unique ID of the pod",
	}
	LabelPodName = LabelDescriptor{
		Key:         "pod_name",
		Description: "The name of the pod",
	}
	// Deprecated label
	LabelPodNamespace = LabelDescriptor{
		Key:         "pod_namespace",
		Description: "The namespace of the pod",
	}
	LabelNamespaceName = LabelDescriptor{
		Key:         "namespace_name",
		Description: "The name of the namespace",
	}
	LabelPodNamespaceUID = LabelDescriptor{
		Key:         "namespace_id",
		Description: "The UID of namespace of the pod",
	}
	LabelContainerName = LabelDescriptor{
		Key:         "container_name",
		Description: "User-provided name of the container or full container name for system containers",
	}
	LabelLabels = LabelDescriptor{
		Key:         "labels",
		Description: "Comma-separated list of user-provided labels",
	}
	LabelNodename = LabelDescriptor{
		Key:         "nodename",
		Description: "nodename where the container ran",
	}
	LabelHostname = LabelDescriptor{
		Key:         "hostname",
		Description: "Hostname where the container ran",
	}
	LabelResourceID = LabelDescriptor{
		Key:         "resource_id",
		Description: "Identifier(s) specific to a metric",
	}
	LabelHostID = LabelDescriptor{
		Key:         "host_id",
		Description: "Identifier specific to a host. Set by cloud provider or user",
	}
	LabelContainerBaseImage = LabelDescriptor{
		Key:         "container_base_image",
		Description: "User-defined image name that is run inside the container",
	}
	// The label is populated only for GCM
	LabelCustomMetricName = LabelDescriptor{
		Key:         "custom_metric_name",
		Description: "User-defined name of the exported custom metric",
	}
	LabelGCEResourceID = LabelDescriptor{
		Key:         "compute.googleapis.com/resource_id",
		Description: "Resource id for nodes specific for GCE.",
	}
	LabelGCEResourceType = LabelDescriptor{
		Key:         "compute.googleapis.com/resource_type",
		Description: "Resource types for nodes specific for GCE.",
	}
)

type LabelDescriptor struct {
	// Key to use for the label.
	Key string `json:"key,omitempty"`

	// Description of the label.
	Description string `json:"description,omitempty"`
}

var commonLabels = []LabelDescriptor{
	LabelNodename,
	LabelHostname,
	LabelHostID,
}

var containerLabels = []LabelDescriptor{
	LabelContainerName,
	LabelContainerBaseImage,
}

var podLabels = []LabelDescriptor{
	LabelPodName,
	LabelPodId,
	LabelPodNamespace,
	LabelPodNamespaceUID,
	LabelLabels,
}

var metricLabels = []LabelDescriptor{
	LabelResourceID,
}

var customMetricLabels = []LabelDescriptor{
	LabelCustomMetricName,
}

// Labels exported to GCM. The number of labels that can be exported to GCM is limited by 10.
var gcmLabels = []LabelDescriptor{
	LabelMetricSetType,
	LabelPodName,
	LabelNamespaceName,
	LabelHostname,
	LabelHostID,
	LabelContainerName,
	LabelContainerBaseImage,
	LabelCustomMetricName,
}

var gcmNodeAutoscalingLabels = []LabelDescriptor{
	LabelGCEResourceID,
	LabelGCEResourceType,
	LabelHostname,
}

func CommonLabels() []LabelDescriptor {
	result := make([]LabelDescriptor, len(commonLabels))
	copy(result, commonLabels)
	return result
}

func ContainerLabels() []LabelDescriptor {
	result := make([]LabelDescriptor, len(containerLabels))
	copy(result, containerLabels)
	return result
}

func PodLabels() []LabelDescriptor {
	result := make([]LabelDescriptor, len(podLabels))
	copy(result, podLabels)
	return result
}

func MetricLabels() []LabelDescriptor {
	result := make([]LabelDescriptor, len(metricLabels)+len(customMetricLabels))
	copy(result, metricLabels)
	copy(result, customMetricLabels)
	return result
}

func SupportedLabels() []LabelDescriptor {
	result := CommonLabels()
	result = append(result, PodLabels()...)
	return append(result, MetricLabels()...)
}

func GcmLabels() map[string]LabelDescriptor {
	result := make(map[string]LabelDescriptor, len(gcmLabels))
	for _, l := range gcmLabels {
		result[l.Key] = l
	}
	return result
}
func GcmNodeAutoscalingLabels() map[string]LabelDescriptor {
	result := make(map[string]LabelDescriptor, len(gcmNodeAutoscalingLabels))
	for _, l := range gcmNodeAutoscalingLabels {
		result[l.Key] = l
	}
	return result
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package rest

import (
	"k8s.io/apiserver/pkg/registry/generic"
	"k8s.io/apiserver/pkg/registry/rest"
	genericapiserver "k8s.io/apiserver/pkg/server"
	serverstorage "k8s.io/apiserver/pkg/server/storage"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/apis/autoscaling"
	autoscalingapiv1 "k8s.io/kubernetes/pkg/apis/autoscaling/v1"
	autoscalingapiv2alpha1 "k8s.io/kubernetes/pkg/apis/autoscaling/v2alpha1"
	horizontalpodautoscalerstore "k8s.io/kubernetes/pkg/registry/autoscaling/horizontalpodautoscaler/storage"
)

type RESTStorageProvider struct{}

func (p RESTStorageProvider) NewRESTStorage(apiResourceConfigSource serverstorage.APIResourceConfigSource, restOptionsGetter generic.RESTOptionsGetter) (genericapiserver.APIGroupInfo, bool) {
	apiGroupInfo := genericapiserver.NewDefaultAPIGroupInfo(autoscaling.GroupName, api.Registry, api.Scheme, api.ParameterCodec, api.Codecs)

	if apiResourceConfigSource.AnyResourcesForVersionEnabled(autoscalingapiv1.SchemeGroupVersion) {
		apiGroupInfo.VersionedResourcesStorageMap[autoscalingapiv1.SchemeGroupVersion.Version] = p.v1Storage(apiResourceConfigSource, restOptionsGetter)
		apiGroupInfo.GroupMeta.GroupVersion = autoscalingapiv1.SchemeGroupVersion
	}
	if apiResourceConfigSource.AnyResourcesForVersionEnabled(autoscalingapiv2alpha1.SchemeGroupVersion) {
		apiGroupInfo.VersionedResourcesStorageMap[autoscalingapiv2alpha1.SchemeGroupVersion.Version] = p.v2alpha1Storage(apiResourceConfigSource, restOptionsGetter)
		apiGroupInfo.GroupMeta.GroupVersion = autoscalingapiv2alpha1.SchemeGroupVersion
	}

	return apiGroupInfo, true
}

func (p RESTStorageProvider) v1Storage(apiResourceConfigSource serverstorage.APIResourceConfigSource, restOptionsGetter generic.RESTOptionsGetter) map[string]rest.Storage {
	version := autoscalingapiv1.SchemeGroupVersion

	storage := map[string]rest.Storage{}
	if apiResourceConfigSource.ResourceEnabled(version.WithResource("horizontalpodautoscalers")) {
		hpaStorage, hpaStatusStorage := horizontalpodautoscalerstore.NewREST(restOptionsGetter)
		storage["horizontalpodautoscalers"] = hpaStorage
		storage["horizontalpodautoscalers/status"] = hpaStatusStorage
	}
	return storage
}

func (p RESTStorageProvider) v2alpha1Storage(apiResourceConfigSource serverstorage.APIResourceConfigSource, restOptionsGetter generic.RESTOptionsGetter) map[string]rest.Storage {
	version := autoscalingapiv2alpha1.SchemeGroupVersion

	storage := map[string]rest.Storage{}
	if apiResourceConfigSource.ResourceEnabled(version.WithResource("horizontalpodautoscalers")) {
		hpaStorage, hpaStatusStorage := horizontalpodautoscalerstore.NewREST(restOptionsGetter)
		storage["horizontalpodautoscalers"] = hpaStorage
		storage["horizontalpodautoscalers/status"] = hpaStatusStorage
	}
	return storage
}

func (p RESTStorageProvider) GroupName() string {
	return autoscaling.GroupName
}
<commit_msg>UPSTREAM: 45741: Fix discovery version for autoscaling to be v1<commit_after>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package rest

import (
	"k8s.io/apiserver/pkg/registry/generic"
	"k8s.io/apiserver/pkg/registry/rest"
	genericapiserver "k8s.io/apiserver/pkg/server"
	serverstorage "k8s.io/apiserver/pkg/server/storage"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/apis/autoscaling"
	autoscalingapiv1 "k8s.io/kubernetes/pkg/apis/autoscaling/v1"
	autoscalingapiv2alpha1 "k8s.io/kubernetes/pkg/apis/autoscaling/v2alpha1"
	horizontalpodautoscalerstore "k8s.io/kubernetes/pkg/registry/autoscaling/horizontalpodautoscaler/storage"
)

type RESTStorageProvider struct{}

func (p RESTStorageProvider) NewRESTStorage(apiResourceConfigSource serverstorage.APIResourceConfigSource, restOptionsGetter generic.RESTOptionsGetter) (genericapiserver.APIGroupInfo, bool) {
	apiGroupInfo := genericapiserver.NewDefaultAPIGroupInfo(autoscaling.GroupName, api.Registry, api.Scheme, api.ParameterCodec, api.Codecs)

	if apiResourceConfigSource.AnyResourcesForVersionEnabled(autoscalingapiv2alpha1.SchemeGroupVersion) {
		apiGroupInfo.VersionedResourcesStorageMap[autoscalingapiv2alpha1.SchemeGroupVersion.Version] = p.v2alpha1Storage(apiResourceConfigSource, restOptionsGetter)
		apiGroupInfo.GroupMeta.GroupVersion = autoscalingapiv2alpha1.SchemeGroupVersion
	}
	if apiResourceConfigSource.AnyResourcesForVersionEnabled(autoscalingapiv1.SchemeGroupVersion) {
		apiGroupInfo.VersionedResourcesStorageMap[autoscalingapiv1.SchemeGroupVersion.Version] = p.v1Storage(apiResourceConfigSource, restOptionsGetter)
		apiGroupInfo.GroupMeta.GroupVersion = autoscalingapiv1.SchemeGroupVersion
	}

	return apiGroupInfo, true
}

func (p RESTStorageProvider) v1Storage(apiResourceConfigSource serverstorage.APIResourceConfigSource, restOptionsGetter generic.RESTOptionsGetter) map[string]rest.Storage {
	version := autoscalingapiv1.SchemeGroupVersion

	storage := map[string]rest.Storage{}
	if apiResourceConfigSource.ResourceEnabled(version.WithResource("horizontalpodautoscalers")) {
		hpaStorage, hpaStatusStorage := horizontalpodautoscalerstore.NewREST(restOptionsGetter)
		storage["horizontalpodautoscalers"] = hpaStorage
		storage["horizontalpodautoscalers/status"] = hpaStatusStorage
	}
	return storage
}

func (p RESTStorageProvider) v2alpha1Storage(apiResourceConfigSource serverstorage.APIResourceConfigSource, restOptionsGetter generic.RESTOptionsGetter) map[string]rest.Storage {
	version := autoscalingapiv2alpha1.SchemeGroupVersion

	storage := map[string]rest.Storage{}
	if apiResourceConfigSource.ResourceEnabled(version.WithResource("horizontalpodautoscalers")) {
		hpaStorage, hpaStatusStorage := horizontalpodautoscalerstore.NewREST(restOptionsGetter)
		storage["horizontalpodautoscalers"] = hpaStorage
		storage["horizontalpodautoscalers/status"] = hpaStatusStorage
	}
	return storage
}

func (p RESTStorageProvider) GroupName() string {
	return autoscaling.GroupName
}
<|endoftext|>


<commit_before>package models

import (
	"errors"
	"net/http"
	"strconv"
	"time"

	"github.com/lib/pq"
	"github.com/thermokarst/bactdb/router"
)

// An Observation is a lookup type
type Observation struct {
	Id                int64       `json:"id,omitempty"`
	ObservationName   string      `db:"observation_name" json:"observationName"`
	ObservationTypeId int64       `db:"observation_type_id" json:"observationTypeId"`
	CreatedAt         time.Time   `db:"created_at" json:"createdAt"`
	UpdatedAt         time.Time   `db:"updated_at" json:"updatedAt"`
	DeletedAt         pq.NullTime `db:"deleted_at" json:"deletedAt"`
}

func NewObservation() *Observation {
	return &Observation{
		ObservationName: "Test Observation",
	}
}

type ObservationsService interface {
	// Get an observation
	Get(id int64) (*Observation, error)

	// List all observations
	List(opt *ObservationListOptions) ([]*Observation, error)

	// Create an observation
	Create(observation *Observation) (bool, error)

	// Update an observation
	Update(id int64, Observation *Observation) (updated bool, err error)

	// Delete an observation
	Delete(id int64) (deleted bool, err error)
}

var (
	ErrObservationNotFound = errors.New("observation not found")
)

type observationsService struct {
	client *Client
}

func (s *observationsService) Get(id int64) (*Observation, error) {
	strId := strconv.FormatInt(id, 10)

	url, err := s.client.url(router.Observation, map[string]string{"Id": strId}, nil)
	if err != nil {
		return nil, err
	}

	req, err := s.client.NewRequest("GET", url.String(), nil)
	if err != nil {
		return nil, err
	}

	var observation *Observation
	_, err = s.client.Do(req, &observation)
	if err != nil {
		return nil, err
	}

	return observation, nil
}

func (s *observationsService) Create(observation *Observation) (bool, error) {
	url, err := s.client.url(router.CreateObservation, nil, nil)
	if err != nil {
		return false, err
	}

	req, err := s.client.NewRequest("POST", url.String(), observation)
	if err != nil {
		return false, err
	}

	resp, err := s.client.Do(req, &observation)
	if err != nil {
		return false, err
	}

	return resp.StatusCode == http.StatusCreated, nil
}

type ObservationListOptions struct {
	ListOptions
}

func (s *observationsService) List(opt *ObservationListOptions) ([]*Observation, error) {
	url, err := s.client.url(router.Observations, nil, opt)
	if err != nil {
		return nil, err
	}

	req, err := s.client.NewRequest("GET", url.String(), nil)
	if err != nil {
		return nil, err
	}

	var observations []*Observation
	_, err = s.client.Do(req, &observations)
	if err != nil {
		return nil, err
	}

	return observations, nil
}

func (s *observationsService) Update(id int64, observation *Observation) (bool, error) {
	strId := strconv.FormatInt(id, 10)

	url, err := s.client.url(router.UpdateObservation, map[string]string{"Id": strId}, nil)
	if err != nil {
		return false, err
	}

	req, err := s.client.NewRequest("PUT", url.String(), observation)
	if err != nil {
		return false, err
	}

	resp, err := s.client.Do(req, &observation)
	if err != nil {
		return false, err
	}

	return resp.StatusCode == http.StatusOK, nil
}

func (s *observationsService) Delete(id int64) (bool, error) {
	strId := strconv.FormatInt(id, 10)

	url, err := s.client.url(router.DeleteObservation, map[string]string{"Id": strId}, nil)
	if err != nil {
		return false, err
	}

	req, err := s.client.NewRequest("DELETE", url.String(), nil)
	if err != nil {
		return false, err
	}

	var observation *Observation
	resp, err := s.client.Do(req, &observation)
	if err != nil {
		return false, err
	}

	return resp.StatusCode == http.StatusOK, nil
}

type MockObservationsService struct {
	Get_    func(id int64) (*Observation, error)
	List_   func(opt *ObservationListOptions) ([]*Observation, error)
	Create_ func(observation *Observation) (bool, error)
	Update_ func(id int64, observation *Observation) (bool, error)
	Delete_ func(id int64) (bool, error)
}

var _ObservationsService = &MockObservationsService{}

func (s *MockObservationsService) Get(id int64) (*Observation, error) {
	if s.Get_ == nil {
		return nil, nil
	}
	return s.Get_(id)
}

func (s *MockObservationsService) Create(observation *Observation) (bool, error) {
	if s.Create_ == nil {
		return false, nil
	}
	return s.Create_(observation)
}

func (s *MockObservationsService) List(opt *ObservationListOptions) ([]*Observation, error) {
	if s.List_ == nil {
		return nil, nil
	}
	return s.List_(opt)
}

func (s *MockObservationsService) Update(id int64, observation *Observation) (bool, error) {
	if s.Update_ == nil {
		return false, nil
	}
	return s.Update_(id, observation)
}

func (s *MockObservationsService) Delete(id int64) (bool, error) {
	if s.Delete_ == nil {
		return false, nil
	}
	return s.Delete_(id)
}
<commit_msg>Typo in Mock Observations Service check.<commit_after>package models

import (
	"errors"
	"net/http"
	"strconv"
	"time"

	"github.com/lib/pq"
	"github.com/thermokarst/bactdb/router"
)

// An Observation is a lookup type
type Observation struct {
	Id                int64       `json:"id,omitempty"`
	ObservationName   string      `db:"observation_name" json:"observationName"`
	ObservationTypeId int64       `db:"observation_type_id" json:"observationTypeId"`
	CreatedAt         time.Time   `db:"created_at" json:"createdAt"`
	UpdatedAt         time.Time   `db:"updated_at" json:"updatedAt"`
	DeletedAt         pq.NullTime `db:"deleted_at" json:"deletedAt"`
}

func NewObservation() *Observation {
	return &Observation{
		ObservationName: "Test Observation",
	}
}

type ObservationsService interface {
	// Get an observation
	Get(id int64) (*Observation, error)

	// List all observations
	List(opt *ObservationListOptions) ([]*Observation, error)

	// Create an observation
	Create(observation *Observation) (bool, error)

	// Update an observation
	Update(id int64, Observation *Observation) (updated bool, err error)

	// Delete an observation
	Delete(id int64) (deleted bool, err error)
}

var (
	ErrObservationNotFound = errors.New("observation not found")
)

type observationsService struct {
	client *Client
}

func (s *observationsService) Get(id int64) (*Observation, error) {
	strId := strconv.FormatInt(id, 10)

	url, err := s.client.url(router.Observation, map[string]string{"Id": strId}, nil)
	if err != nil {
		return nil, err
	}

	req, err := s.client.NewRequest("GET", url.String(), nil)
	if err != nil {
		return nil, err
	}

	var observation *Observation
	_, err = s.client.Do(req, &observation)
	if err != nil {
		return nil, err
	}

	return observation, nil
}

func (s *observationsService) Create(observation *Observation) (bool, error) {
	url, err := s.client.url(router.CreateObservation, nil, nil)
	if err != nil {
		return false, err
	}

	req, err := s.client.NewRequest("POST", url.String(), observation)
	if err != nil {
		return false, err
	}

	resp, err := s.client.Do(req, &observation)
	if err != nil {
		return false, err
	}

	return resp.StatusCode == http.StatusCreated, nil
}

type ObservationListOptions struct {
	ListOptions
}

func (s *observationsService) List(opt *ObservationListOptions) ([]*Observation, error) {
	url, err := s.client.url(router.Observations, nil, opt)
	if err != nil {
		return nil, err
	}

	req, err := s.client.NewRequest("GET", url.String(), nil)
	if err != nil {
		return nil, err
	}

	var observations []*Observation
	_, err = s.client.Do(req, &observations)
	if err != nil {
		return nil, err
	}

	return observations, nil
}

func (s *observationsService) Update(id int64, observation *Observation) (bool, error) {
	strId := strconv.FormatInt(id, 10)

	url, err := s.client.url(router.UpdateObservation, map[string]string{"Id": strId}, nil)
	if err != nil {
		return false, err
	}

	req, err := s.client.NewRequest("PUT", url.String(), observation)
	if err != nil {
		return false, err
	}

	resp, err := s.client.Do(req, &observation)
	if err != nil {
		return false, err
	}

	return resp.StatusCode == http.StatusOK, nil
}

func (s *observationsService) Delete(id int64) (bool, error) {
	strId := strconv.FormatInt(id, 10)

	url, err := s.client.url(router.DeleteObservation, map[string]string{"Id": strId}, nil)
	if err != nil {
		return false, err
	}

	req, err := s.client.NewRequest("DELETE", url.String(), nil)
	if err != nil {
		return false, err
	}

	var observation *Observation
	resp, err := s.client.Do(req, &observation)
	if err != nil {
		return false, err
	}

	return resp.StatusCode == http.StatusOK, nil
}

type MockObservationsService struct {
	Get_    func(id int64) (*Observation, error)
	List_   func(opt *ObservationListOptions) ([]*Observation, error)
	Create_ func(observation *Observation) (bool, error)
	Update_ func(id int64, observation *Observation) (bool, error)
	Delete_ func(id int64) (bool, error)
}

var _ ObservationsService = &MockObservationsService{}

func (s *MockObservationsService) Get(id int64) (*Observation, error) {
	if s.Get_ == nil {
		return nil, nil
	}
	return s.Get_(id)
}

func (s *MockObservationsService) Create(observation *Observation) (bool, error) {
	if s.Create_ == nil {
		return false, nil
	}
	return s.Create_(observation)
}

func (s *MockObservationsService) List(opt *ObservationListOptions) ([]*Observation, error) {
	if s.List_ == nil {
		return nil, nil
	}
	return s.List_(opt)
}

func (s *MockObservationsService) Update(id int64, observation *Observation) (bool, error) {
	if s.Update_ == nil {
		return false, nil
	}
	return s.Update_(id, observation)
}

func (s *MockObservationsService) Delete(id int64) (bool, error) {
	if s.Delete_ == nil {
		return false, nil
	}
	return s.Delete_(id)
}
<|endoftext|>


<commit_before>package models_test

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

	"bitbucket.org/kardianos/osext"
	. "github.com/gernest/lora/models"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Project", func() {
	var (
		base           string
		err            error
		project        Project
		basePath       string
		currentProject *Project
		baseProject    Project
	)
	baseProject = Project{
		Id:           1,
		Title:        "my new tushabe site",
		Name:         "pasiansi",
		Theme:        "loraina",
		PublishDir:   "www",
		BaseUrl:      "http://yourSiteHere",
		LanguageCode: "en-us",
		Pages: []Page{
			{Id: 1, Title: "home", Content: "## hello home", Slug: "slug", Draft: false},
			{Id: 2, Title: "about", Content: "## hello about", Slug: "slug", Draft: false},
			{Id: 3, Title: "products", Content: "## hello products", Slug: "slug", Draft: false},
			{Id: 4, Title: "contact", Content: "## hello contact", Slug: "slug", Draft: false},
		},
	}
	basePath, _ = osext.ExecutableFolder()
	base = filepath.Join(path.Dir(strings.TrimSuffix(basePath, "/")), "fixtures")
	Describe("InitializeProject", func() {
		AfterEach(func() {
			clearAll(filepath.Join(base, "projects"))
		})
		PContext(" Given the base path ", func() {

			It("Shold populate with initial values", func() {
				p, _ := NewLoraProject(base, "mchele", "", "")
				Expect(base).Should(Equal(p.BaseDir))
			})

		})
		Context("Without base path", func() {

			It("should have default values ", func() {
				p := new(Project)
				err = p.Initialize(project.BaseDir, "unga", "", "")
				Expect(err).Should(HaveOccurred())
			})
		})

	})
	Describe("GenerateScaffold", func() {
		BeforeEach(func() {
			currentProject = new(Project)
			_ = currentProject.Initialize(base, "kilimahewa", "", "")

		})
		AfterEach(func() {
			_ = currentProject.Clean()
		})

		It("should generate project", func() {
			err = currentProject.GenScaffold()
			projectPath := base + "/projects/" + "kilimahewa"
			file, _ := os.Stat(currentProject.ProjectPath)

			Expect(err).ShouldNot(HaveOccurred())
			Expect(currentProject.ProjectPath).Should(Equal(projectPath))
			Expect(file.IsDir()).Should(BeTrue())
		})
	})

	Describe("Clean", func() {
		BeforeEach(func() {
			currentProject = new(Project)
			_ = currentProject.Initialize(base, "bigbite", "", "")
			_ = currentProject.GenScaffold()

		})

		It("should remove generated files", func() {
			err = currentProject.Clean()
			file, _ := os.Stat(currentProject.ProjectPath)

			Expect(err).ShouldNot(HaveOccurred())
			Expect(file).Should(BeNil())
		})

	})
	Describe("LoadConfigFile", func() {
		BeforeEach(func() {
			currentProject = new(Project)
			_ = currentProject.Initialize(base, "pasiansi", "", "")
			_ = currentProject.GenScaffold()
		})
		AfterEach(func() {
			_ = currentProject.Clean()
		})
		It("Loads config file", func() {
			err = currentProject.LoadConfigFile()

			Expect(err).ShouldNot(HaveOccurred())
		})
		It("should have correct confg values", func() {
			_ = currentProject.LoadConfigFile()

			Expect(currentProject.Name).Should(Equal(baseProject.Name))
			Expect(currentProject.Title).Should(Equal(baseProject.Title))
			Expect(currentProject.Id).Should(Equal(baseProject.Id))
			Expect(currentProject.PublishDir).Should(Equal(baseProject.PublishDir))
			Expect(currentProject.BaseUrl).Should(Equal(baseProject.BaseUrl))
		})
	})
	Describe("SaveConfigFile", func() {
		BeforeEach(func() {
			_ = currentProject.Initialize(base, "pasiansi", "", "")
			_ = currentProject.GenScaffold()
			_ = currentProject.LoadConfigFile()
		})
		AfterEach(func() {
			_ = currentProject.Clean()
		})
		It("Should save", func() {
			err = currentProject.SaveConfigFile()

			Expect(err).ShouldNot(HaveOccurred())
		})
		It("Should update the values", func() {
			currentProject.Title = "Fuck ISIS"
			_ = currentProject.SaveConfigFile()
			_ = currentProject.LoadConfigFile()

			Expect(currentProject.Title).Should(Equal("Fuck ISIS"))
		})
	})
	Describe("Install template", func() {
		BeforeEach(func() {
			_ = currentProject.Initialize(base, "yoyo", "", "")
			_ = currentProject.GenScaffold()
		})
		It("It tick", func() {
			err := currentProject.InstallTemplate("", "")
			Expect(err).ShouldNot(HaveOccurred())
		})
	})
	Describe("Install template", func() {
		BeforeEach(func() {
			_ = currentProject.Initialize(base, "yoyo", "", "")
			_ = currentProject.GenScaffold()
		})
		It("ticks", func() {
			err := currentProject.InstallTheme("")
			Expect(err).ShouldNot(HaveOccurred())
		})
		It("No theme", func() {
			err := currentProject.InstallTheme("nouma")
			Expect(err).Should(HaveOccurred())
		})
	})

	Describe("Page", func() {
		BeforeEach(func() {
			_ = currentProject.Initialize(base, "pasiansi", "", "")
			_ = currentProject.GenScaffold()
			_ = currentProject.LoadConfigFile()
		})
		AfterEach(func() {
			_ = currentProject.Clean()
		})
		It("Should generate a new page", func() {
			p := currentProject.Pages[0]

			Expect(p.Generate(currentProject)).Should(Succeed())
			Expect(p.ContentPath).ShouldNot(BeEmpty())
		})
	})

})

func clearAll(s string) {
	fmt.Printf("cleaning %s \n", s)
	_ = os.RemoveAll(s)
}
<commit_msg>fix test suite<commit_after>package models_test

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

	"bitbucket.org/kardianos/osext"
	. "github.com/gernest/lora/models"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Project", func() {
	var (
		base           string
		err            error
		project        Project
		basePath       string
		currentProject *Project
		baseProject    Project
	)
	baseProject = Project{
		Id:           1,
		Title:        "my new tushabe site",
		Name:         "pasiansi",
		Theme:        "loraina",
		PublishDir:   "www",
		BaseUrl:      "http://yourSiteHere",
		LanguageCode: "en-us",
		Pages: []Page{
			{Id: 1, Title: "home", Content: "## hello home", Slug: "slug", Draft: false},
			{Id: 2, Title: "about", Content: "## hello about", Slug: "slug", Draft: false},
			{Id: 3, Title: "products", Content: "## hello products", Slug: "slug", Draft: false},
			{Id: 4, Title: "contact", Content: "## hello contact", Slug: "slug", Draft: false},
		},
	}
	basePath, _ = osext.ExecutableFolder()
	base = filepath.Join(path.Dir(strings.TrimSuffix(basePath, "/")), "fixtures")
	Describe("InitializeProject", func() {
		AfterEach(func() {
			clearAll(filepath.Join(base, "projects"))
		})
		PContext(" Given the base path ", func() {

			It("Shold populate with initial values", func() {
				p, _ := NewLoraProject(base, "mchele", "", "")
				Expect(base).Should(Equal(p.BaseDir))
			})

		})
		Context("Without base path", func() {

			It("should have default values ", func() {
				p := new(Project)
				err = p.Initialize(project.BaseDir, "unga", "", "")
				Expect(err).Should(HaveOccurred())
			})
		})

	})
	Describe("GenerateScaffold", func() {
		BeforeEach(func() {
			currentProject = new(Project)
			_ = currentProject.Initialize(base, "kilimahewa", "", "")

		})
		AfterEach(func() {
			_ = currentProject.Clean()
		})

		It("should generate project", func() {
			err = currentProject.GenScaffold()
			projectPath := base + "/projects/" + "kilimahewa"
			file, _ := os.Stat(currentProject.ProjectPath)

			Expect(err).ShouldNot(HaveOccurred())
			Expect(currentProject.ProjectPath).Should(Equal(projectPath))
			Expect(file.IsDir()).Should(BeTrue())
		})
	})

	Describe("Clean", func() {
		BeforeEach(func() {
			currentProject = new(Project)
			_ = currentProject.Initialize(base, "bigbite", "", "")
			_ = currentProject.GenScaffold()

		})

		It("should remove generated files", func() {
			err = currentProject.Clean()
			file, _ := os.Stat(currentProject.ProjectPath)

			Expect(err).ShouldNot(HaveOccurred())
			Expect(file).Should(BeNil())
		})

	})
	Describe("LoadConfigFile", func() {
		BeforeEach(func() {
			currentProject = new(Project)
			_ = currentProject.Initialize(base, "pasiansi", "", "")
			_ = currentProject.GenScaffold()
		})
		AfterEach(func() {
			_ = currentProject.Clean()
		})
		It("Loads config file", func() {
			err = currentProject.LoadConfigFile()

			Expect(err).ShouldNot(HaveOccurred())
		})
		It("should have correct confg values", func() {
			_ = currentProject.LoadConfigFile()

			Expect(currentProject.Name).Should(Equal(baseProject.Name))
			Expect(currentProject.Title).Should(Equal(baseProject.Title))
			Expect(currentProject.Id).Should(Equal(baseProject.Id))
			Expect(currentProject.PublishDir).Should(Equal(baseProject.PublishDir))
			Expect(currentProject.BaseUrl).Should(Equal(baseProject.BaseUrl))
		})
	})
	Describe("SaveConfigFile", func() {
		BeforeEach(func() {
			_ = currentProject.Initialize(base, "pasiansi", "", "")
			_ = currentProject.GenScaffold()
			_ = currentProject.LoadConfigFile()
		})
		AfterEach(func() {
			_ = currentProject.Clean()
		})
		It("Should save", func() {
			err = currentProject.SaveConfigFile()

			Expect(err).ShouldNot(HaveOccurred())
		})
		It("Should update the values", func() {
			currentProject.Title = "Fuck ISIS"
			_ = currentProject.SaveConfigFile()
			_ = currentProject.LoadConfigFile()

			Expect(currentProject.Title).Should(Equal("Fuck ISIS"))
		})
	})
	Describe("Install template", func() {
		BeforeEach(func() {
			_ = currentProject.Initialize(base, "yoyo", "", "")
			_ = currentProject.GenScaffold()
		})
		It("It tick", func() {
			err := currentProject.InstallTemplate("", "")
			Expect(err).ShouldNot(HaveOccurred())
		})
	})
	Describe("Install template", func() {
		BeforeEach(func() {
			_ = currentProject.Initialize(base, "yoyo", "", "")
			_ = currentProject.GenScaffold()
		})
		AfterEach(func() {
			_ = currentProject.Clean()
		})
		It("ticks", func() {
			err := currentProject.InstallTheme("")
			Expect(err).ShouldNot(HaveOccurred())
		})
		It("No theme", func() {
			err := currentProject.InstallTheme("nouma")
			Expect(err).Should(HaveOccurred())
		})
	})

	Describe("Page", func() {
		BeforeEach(func() {
			_ = currentProject.Initialize(base, "pasiansi", "", "")
			_ = currentProject.GenScaffold()
			_ = currentProject.LoadConfigFile()
		})
		AfterEach(func() {
			_ = currentProject.Clean()
		})
		It("Should generate a new page", func() {
			p := currentProject.Pages[0]

			Expect(p.Generate(currentProject)).Should(Succeed())
			Expect(p.ContentPath).ShouldNot(BeEmpty())
		})
	})

})

func clearAll(s string) {
	fmt.Printf("cleaning %s \n", s)
	_ = os.RemoveAll(s)
}
<|endoftext|>


<commit_before>package loader

import (
	"context"
	"strconv"
	"strings"

	"github.com/go-ggz/ggz/helper"
	"github.com/go-ggz/ggz/model"
	"github.com/go-ggz/ggz/module/loader/lru"
	"github.com/go-ggz/ggz/module/loader/memory"

	"gopkg.in/nicksrandall/dataloader.v5"
)

var (
	// Cache for dataloader
	Cache dataloader.Cache
	// UserIDCache for user cache from ID
	UserIDCache *dataloader.Loader
)

// NewEngine for initialize cache engine
func NewEngine(driver, prefix string, expire int) error {
	switch driver {
	case "lru":
		Cache = lru.NewEngine(prefix)
	case "memory":
		Cache = memory.NewEngine(prefix, expire)
	default:
		Cache = dataloader.NewCache()
	}

	// load cache
	initLoader()

	return nil
}

const sep = ":"

// GetCacheKey get cache key for data loader
func GetCacheKey(module string, id interface{}) string {
	var str string
	switch v := id.(type) {
	case int64:
		str = strconv.FormatInt(v, 10)
	case string:
		str = v
	}
	return module + sep + str
}

// GetCacheID get cache id for model id
func GetCacheID(key string) (interface{}, error) {
	strs := strings.Split(key, sep)

	return strs[1], nil
}

func initLoader() {
	UserIDCache = dataloader.NewBatchedLoader(userBatch, dataloader.WithCache(Cache))
}

func userBatch(ctx context.Context, keys dataloader.Keys) []*dataloader.Result {
	var results []*dataloader.Result
	id, _ := helper.GetCacheID(keys[0].String())

	user, err := model.GetUserByID(id.(int64))

	results = append(results, &dataloader.Result{
		Data:  user,
		Error: err,
	})

	return results
}

// GetUserFromLoader get user cache
func GetUserFromLoader(ctx context.Context, id interface{}) (*model.User, error) {
	key := GetCacheKey("user", id)
	userCache, err := UserIDCache.Load(ctx, dataloader.StringKey(key))()

	if err != nil {
		return nil, err
	}

	return userCache.(*model.User), nil
}
<commit_msg>fix: rename user loader<commit_after>package loader

import (
	"context"
	"strconv"
	"strings"

	"github.com/go-ggz/ggz/helper"
	"github.com/go-ggz/ggz/model"
	"github.com/go-ggz/ggz/module/loader/lru"
	"github.com/go-ggz/ggz/module/loader/memory"

	"gopkg.in/nicksrandall/dataloader.v5"
)

var (
	// Cache for dataloader
	Cache dataloader.Cache
	// UserIDCache for user cache from ID
	UserCache *dataloader.Loader
)

// NewEngine for initialize cache engine
func NewEngine(driver, prefix string, expire int) error {
	switch driver {
	case "lru":
		Cache = lru.NewEngine(prefix)
	case "memory":
		Cache = memory.NewEngine(prefix, expire)
	default:
		Cache = dataloader.NewCache()
	}

	// load cache
	initLoader()

	return nil
}

const sep = ":"

// GetCacheKey get cache key for data loader
func GetCacheKey(module string, id interface{}) string {
	var str string
	switch v := id.(type) {
	case int64:
		str = strconv.FormatInt(v, 10)
	case string:
		str = v
	}
	return module + sep + str
}

// GetCacheID get cache id for model id
func GetCacheID(key string) (interface{}, error) {
	strs := strings.Split(key, sep)

	return strs[1], nil
}

func initLoader() {
	UserCache = dataloader.NewBatchedLoader(userBatch, dataloader.WithCache(Cache))
}

func userBatch(ctx context.Context, keys dataloader.Keys) []*dataloader.Result {
	var results []*dataloader.Result
	id, _ := helper.GetCacheID(keys[0].String())

	user, err := model.GetUserByID(id.(int64))

	results = append(results, &dataloader.Result{
		Data:  user,
		Error: err,
	})

	return results
}

// GetUserFromLoader get user cache
func GetUserFromLoader(ctx context.Context, id interface{}) (*model.User, error) {
	key := GetCacheKey("user", id)
	userCache, err := UserCache.Load(ctx, dataloader.StringKey(key))()

	if err != nil {
		return nil, err
	}

	return userCache.(*model.User), nil
}
<|endoftext|>


<commit_before>package loader

import (
	"context"
	"strconv"
	"strings"

	"github.com/go-ggz/ggz/model"
	"github.com/go-ggz/ggz/module/loader/lru"
	"github.com/go-ggz/ggz/module/loader/memory"

	"gopkg.in/nicksrandall/dataloader.v5"
)

var (
	// Cache for dataloader
	Cache dataloader.Cache
	// UserIDCache for user cache from ID
	UserCache *dataloader.Loader
)

// NewEngine for initialize cache engine
func NewEngine(driver, prefix string, expire int) error {
	switch driver {
	case "lru":
		Cache = lru.NewEngine(prefix)
	case "memory":
		Cache = memory.NewEngine(prefix, expire)
	default:
		Cache = dataloader.NewCache()
	}

	// load cache
	initLoader()

	return nil
}

const sep = ":"

// GetCacheKey get cache key for data loader
func GetCacheKey(module string, id interface{}) string {
	var str string
	switch v := id.(type) {
	case int64:
		str = strconv.FormatInt(v, 10)
	case string:
		str = v
	}
	return module + sep + str
}

// GetCacheID get cache id for model id
func GetCacheID(key string) (interface{}, error) {
	strs := strings.Split(key, sep)

	return strs[1], nil
}

func initLoader() {
	UserCache = dataloader.NewBatchedLoader(userBatch, dataloader.WithCache(Cache))
}

func userBatch(ctx context.Context, keys dataloader.Keys) []*dataloader.Result {
	var results []*dataloader.Result
	id, _ := GetCacheID(keys[0].String())

	user, err := model.GetUserByID(id.(int64))

	results = append(results, &dataloader.Result{
		Data:  user,
		Error: err,
	})

	return results
}

// GetUserFromLoader get user cache
func GetUserFromLoader(ctx context.Context, id interface{}) (*model.User, error) {
	key := GetCacheKey("user", id)
	userCache, err := UserCache.Load(ctx, dataloader.StringKey(key))()

	if err != nil {
		return nil, err
	}

	return userCache.(*model.User), nil
}
<commit_msg>fix: unexport cache func.<commit_after>package loader

import (
	"context"
	"strconv"
	"strings"

	"github.com/go-ggz/ggz/model"
	"github.com/go-ggz/ggz/module/loader/lru"
	"github.com/go-ggz/ggz/module/loader/memory"

	"gopkg.in/nicksrandall/dataloader.v5"
)

var (
	// Cache for dataloader
	Cache dataloader.Cache
	// UserIDCache for user cache from ID
	UserCache *dataloader.Loader
)

const sep = ":"

func initLoader() {
	UserCache = dataloader.NewBatchedLoader(userBatch, dataloader.WithCache(Cache))
}

func getCacheKey(module string, id interface{}) string {
	var str string
	switch v := id.(type) {
	case int64:
		str = strconv.FormatInt(v, 10)
	case string:
		str = v
	}
	return module + sep + str
}

func getCacheID(key string) (interface{}, error) {
	strs := strings.Split(key, sep)

	return strs[1], nil
}

// NewEngine for initialize cache engine
func NewEngine(driver, prefix string, expire int) error {
	switch driver {
	case "lru":
		Cache = lru.NewEngine(prefix)
	case "memory":
		Cache = memory.NewEngine(prefix, expire)
	default:
		Cache = dataloader.NewCache()
	}

	// load cache
	initLoader()

	return nil
}

func userBatch(ctx context.Context, keys dataloader.Keys) []*dataloader.Result {
	var results []*dataloader.Result
	id, _ := getCacheID(keys[0].String())

	user, err := model.GetUserByID(id.(int64))

	results = append(results, &dataloader.Result{
		Data:  user,
		Error: err,
	})

	return results
}

// GetUserFromLoader get user cache
func GetUserFromLoader(ctx context.Context, id interface{}) (*model.User, error) {
	key := getCacheKey("user", id)
	userCache, err := UserCache.Load(ctx, dataloader.StringKey(key))()

	if err != nil {
		return nil, err
	}

	return userCache.(*model.User), nil
}
<|endoftext|>


<commit_before>package gateway

import (
	"errors"
	"sync"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
)

// rpcID is an 8-byte signature that is added to all RPCs to tell the gatway
// what to do with the RPC.
type rpcID [8]byte

// String returns a string representation of an rpcID.
func (id rpcID) String() string {
	for i := range id {
		if id[i] == 0 {
			id[i] = ' '
		}
	}
	return string(id[:])
}

// handlerName truncates a string to 8 bytes. If len(name) < 8, the remaining
// bytes are 0. A handlerName is specified at the beginning of each network
// call, indicating which function should handle the connection.
func handlerName(name string) (id rpcID) {
	copy(id[:], name)
	return
}

// managedRPC calls an RPC on the given address. managedRPC cannot be called on
// an address that the Gateway is not connected to.
func (g *Gateway) managedRPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	g.mu.RLock()
	peer, ok := g.peers[addr]
	g.mu.RUnlock()
	if !ok {
		return errors.New("can't call RPC on unconnected peer " + string(addr))
	}

	conn, err := peer.open()
	if err != nil {
		return err
	}
	defer conn.Close()

	// write header
	if err := encoding.WriteObject(conn, handlerName(name)); err != nil {
		return err
	}
	// call fn
	return fn(conn)
}

// RPC calls an RPC on the given address. RPC cannot be called on an address
// that the Gateway is not connected to.
func (g *Gateway) RPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	if err := g.threads.Add(); err != nil {
		return err
	}
	defer g.threads.Done()
	return g.managedRPC(addr, name, fn)
}

// RegisterRPC registers an RPCFunc as a handler for a given identifier. To
// call an RPC, use gateway.RPC, supplying the same identifier given to
// RegisterRPC. Identifiers should always use PascalCase. The first 8
// characters of an identifier should be unique, as the identifier used
// internally is truncated to 8 bytes.
func (g *Gateway) RegisterRPC(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; ok {
		build.Critical("RPC already registered: " + name)
	}
	g.handlers[handlerName(name)] = fn
}

// UnregisterRPC unregisters an RPC and removes the corresponding RPCFunc from
// g.handlers. Future calls to the RPC by peers will fail.
func (g *Gateway) UnregisterRPC(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; !ok {
		build.Critical("RPC not registered: " + name)
	}
	delete(g.handlers, handlerName(name))
}

// RegisterConnectCall registers a name and RPCFunc to be called on a peer
// upon connecting.
func (g *Gateway) RegisterConnectCall(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; ok {
		build.Critical("ConnectCall already registered: " + name)
	}
	g.initRPCs[name] = fn
}

// UnregisterConnectCall unregisters an on-connect call and removes the
// corresponding RPCFunc from g.initRPCs. Future connections to peers will not
// trigger the RPC to be called on them.
func (g *Gateway) UnregisterConnectCall(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; !ok {
		build.Critical("ConnectCall not registered: " + name)
	}
	delete(g.initRPCs, name)
}

// threadedListenPeer listens for new streams on a peer connection and serves them via
// threadedHandleConn.
func (g *Gateway) threadedListenPeer(p *peer) {
	// threadedListenPeer registers to the peerTG instead of the primary thread
	// group because peer connections can be lifetime in length, but can also
	// be short-lived. The fact that they can be lifetime means that they can't
	// call threads.Add as they will block calls to threads.Flush. The fact
	// that they can be short-lived means that threads.OnStop is not a good
	// tool for closing out the threads. Instead, they register to peerTG,
	// which is cleanly closed upon gateway shutdown but will not block any
	// calls to threads.Flush()
	if g.peerTG.Add() != nil {
		return
	}
	defer g.peerTG.Done()

	// Spin up a goroutine to listen for a shutdown signal from both the peer
	// and from the gateway. In the event of either, close the muxado session.
	connClosedChan := make(chan struct{})
	peerCloseChan := make(chan struct{})
	go func() {
		// Signal that the muxado session has been successfully closed, and
		// that this goroutine has terminated.
		defer close(connClosedChan)

		// Listen for a stop signal.
		select {
		case <-g.threads.StopChan():
		case <-peerCloseChan:
		}

		// Can't call Disconnect because it could return sync.ErrStopped.
		g.mu.Lock()
		delete(g.peers, p.NetAddress)
		g.mu.Unlock()
		if err := p.sess.Close(); err != nil {
			g.log.Debugf("WARN: error disconnecting from peer %q: %v", p.NetAddress, err)
		}
	}()

	for {
		conn, err := p.accept()
		if err != nil {
			g.log.Debugln("Peer connection closed:", p.NetAddress)
			break
		}

		// it is the handler's responsibility to close the connection
		go g.threadedHandleConn(conn)
	}
	// Signal that the goroutine can shutdown.
	close(peerCloseChan)
	// Wait for confirmation that the goroutine has shut down before returning
	// and releasing the threadgroup registration.
	<-connClosedChan
}

// threadedHandleConn reads header data from a connection, then routes it to the
// appropriate handler for further processing.
func (g *Gateway) threadedHandleConn(conn modules.PeerConn) {
	defer conn.Close()
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	var id rpcID
	if err := encoding.ReadObject(conn, &id, 8); err != nil {
		return
	}
	// call registered handler for this ID
	g.mu.RLock()
	fn, ok := g.handlers[id]
	g.mu.RUnlock()
	if !ok {
		g.log.Debugf("WARN: incoming conn %v requested unknown RPC \"%v\"", conn.RPCAddr(), id)
		return
	}
	g.log.Debugf("INFO: incoming conn %v requested RPC \"%v\"", conn.RPCAddr(), id)

	// call fn
	err := fn(conn)
	// don't log benign errors
	if err == modules.ErrDuplicateTransactionSet || err == modules.ErrBlockKnown {
		err = nil
	}
	if err != nil {
		g.log.Debugf("WARN: incoming RPC \"%v\" from conn %v failed: %v", id, conn.RPCAddr(), err)
	}
}

// Broadcast calls an RPC on all of the specified peers. The calls are run in
// parallel. Broadcasts are restricted to "one-way" RPCs, which simply write an
// object and disconnect. This is why Broadcast takes an interface{} instead of
// an RPCFunc.
func (g *Gateway) Broadcast(name string, obj interface{}, peers []modules.Peer) {
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	g.log.Printf("INFO: broadcasting RPC %q to %v peers", name, len(peers))

	// only encode obj once, instead of using WriteObject
	enc := encoding.Marshal(obj)
	fn := func(conn modules.PeerConn) error {
		return encoding.WritePrefix(conn, enc)
	}

	var wg sync.WaitGroup
	for _, p := range peers {
		wg.Add(1)
		go func(addr modules.NetAddress) {
			defer wg.Done()
			err := g.managedRPC(addr, name, fn)
			if err != nil {
				g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed (attempting again in 10 seconds): %v", name, addr, err)
				// try one more time before giving up
				select {
				case <-time.After(10 * time.Second):
				case <-g.threads.StopChan():
					return
				}
				err := g.managedRPC(addr, name, fn)
				if err != nil {
					g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed twice: %v", name, addr, err)
				}
			}
		}(p.NetAddress)
	}
	wg.Wait()
}
<commit_msg>Broadcast log is now a debug log<commit_after>package gateway

import (
	"errors"
	"sync"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
)

// rpcID is an 8-byte signature that is added to all RPCs to tell the gatway
// what to do with the RPC.
type rpcID [8]byte

// String returns a string representation of an rpcID.
func (id rpcID) String() string {
	for i := range id {
		if id[i] == 0 {
			id[i] = ' '
		}
	}
	return string(id[:])
}

// handlerName truncates a string to 8 bytes. If len(name) < 8, the remaining
// bytes are 0. A handlerName is specified at the beginning of each network
// call, indicating which function should handle the connection.
func handlerName(name string) (id rpcID) {
	copy(id[:], name)
	return
}

// managedRPC calls an RPC on the given address. managedRPC cannot be called on
// an address that the Gateway is not connected to.
func (g *Gateway) managedRPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	g.mu.RLock()
	peer, ok := g.peers[addr]
	g.mu.RUnlock()
	if !ok {
		return errors.New("can't call RPC on unconnected peer " + string(addr))
	}

	conn, err := peer.open()
	if err != nil {
		return err
	}
	defer conn.Close()

	// write header
	if err := encoding.WriteObject(conn, handlerName(name)); err != nil {
		return err
	}
	// call fn
	return fn(conn)
}

// RPC calls an RPC on the given address. RPC cannot be called on an address
// that the Gateway is not connected to.
func (g *Gateway) RPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	if err := g.threads.Add(); err != nil {
		return err
	}
	defer g.threads.Done()
	return g.managedRPC(addr, name, fn)
}

// RegisterRPC registers an RPCFunc as a handler for a given identifier. To
// call an RPC, use gateway.RPC, supplying the same identifier given to
// RegisterRPC. Identifiers should always use PascalCase. The first 8
// characters of an identifier should be unique, as the identifier used
// internally is truncated to 8 bytes.
func (g *Gateway) RegisterRPC(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; ok {
		build.Critical("RPC already registered: " + name)
	}
	g.handlers[handlerName(name)] = fn
}

// UnregisterRPC unregisters an RPC and removes the corresponding RPCFunc from
// g.handlers. Future calls to the RPC by peers will fail.
func (g *Gateway) UnregisterRPC(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; !ok {
		build.Critical("RPC not registered: " + name)
	}
	delete(g.handlers, handlerName(name))
}

// RegisterConnectCall registers a name and RPCFunc to be called on a peer
// upon connecting.
func (g *Gateway) RegisterConnectCall(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; ok {
		build.Critical("ConnectCall already registered: " + name)
	}
	g.initRPCs[name] = fn
}

// UnregisterConnectCall unregisters an on-connect call and removes the
// corresponding RPCFunc from g.initRPCs. Future connections to peers will not
// trigger the RPC to be called on them.
func (g *Gateway) UnregisterConnectCall(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; !ok {
		build.Critical("ConnectCall not registered: " + name)
	}
	delete(g.initRPCs, name)
}

// threadedListenPeer listens for new streams on a peer connection and serves them via
// threadedHandleConn.
func (g *Gateway) threadedListenPeer(p *peer) {
	// threadedListenPeer registers to the peerTG instead of the primary thread
	// group because peer connections can be lifetime in length, but can also
	// be short-lived. The fact that they can be lifetime means that they can't
	// call threads.Add as they will block calls to threads.Flush. The fact
	// that they can be short-lived means that threads.OnStop is not a good
	// tool for closing out the threads. Instead, they register to peerTG,
	// which is cleanly closed upon gateway shutdown but will not block any
	// calls to threads.Flush()
	if g.peerTG.Add() != nil {
		return
	}
	defer g.peerTG.Done()

	// Spin up a goroutine to listen for a shutdown signal from both the peer
	// and from the gateway. In the event of either, close the muxado session.
	connClosedChan := make(chan struct{})
	peerCloseChan := make(chan struct{})
	go func() {
		// Signal that the muxado session has been successfully closed, and
		// that this goroutine has terminated.
		defer close(connClosedChan)

		// Listen for a stop signal.
		select {
		case <-g.threads.StopChan():
		case <-peerCloseChan:
		}

		// Can't call Disconnect because it could return sync.ErrStopped.
		g.mu.Lock()
		delete(g.peers, p.NetAddress)
		g.mu.Unlock()
		if err := p.sess.Close(); err != nil {
			g.log.Debugf("WARN: error disconnecting from peer %q: %v", p.NetAddress, err)
		}
	}()

	for {
		conn, err := p.accept()
		if err != nil {
			g.log.Debugln("Peer connection closed:", p.NetAddress)
			break
		}

		// it is the handler's responsibility to close the connection
		go g.threadedHandleConn(conn)
	}
	// Signal that the goroutine can shutdown.
	close(peerCloseChan)
	// Wait for confirmation that the goroutine has shut down before returning
	// and releasing the threadgroup registration.
	<-connClosedChan
}

// threadedHandleConn reads header data from a connection, then routes it to the
// appropriate handler for further processing.
func (g *Gateway) threadedHandleConn(conn modules.PeerConn) {
	defer conn.Close()
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	var id rpcID
	if err := encoding.ReadObject(conn, &id, 8); err != nil {
		return
	}
	// call registered handler for this ID
	g.mu.RLock()
	fn, ok := g.handlers[id]
	g.mu.RUnlock()
	if !ok {
		g.log.Debugf("WARN: incoming conn %v requested unknown RPC \"%v\"", conn.RPCAddr(), id)
		return
	}
	g.log.Debugf("INFO: incoming conn %v requested RPC \"%v\"", conn.RPCAddr(), id)

	// call fn
	err := fn(conn)
	// don't log benign errors
	if err == modules.ErrDuplicateTransactionSet || err == modules.ErrBlockKnown {
		err = nil
	}
	if err != nil {
		g.log.Debugf("WARN: incoming RPC \"%v\" from conn %v failed: %v", id, conn.RPCAddr(), err)
	}
}

// Broadcast calls an RPC on all of the specified peers. The calls are run in
// parallel. Broadcasts are restricted to "one-way" RPCs, which simply write an
// object and disconnect. This is why Broadcast takes an interface{} instead of
// an RPCFunc.
func (g *Gateway) Broadcast(name string, obj interface{}, peers []modules.Peer) {
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	g.log.Debugf("INFO: broadcasting RPC %q to %v peers", name, len(peers))

	// only encode obj once, instead of using WriteObject
	enc := encoding.Marshal(obj)
	fn := func(conn modules.PeerConn) error {
		return encoding.WritePrefix(conn, enc)
	}

	var wg sync.WaitGroup
	for _, p := range peers {
		wg.Add(1)
		go func(addr modules.NetAddress) {
			defer wg.Done()
			err := g.managedRPC(addr, name, fn)
			if err != nil {
				g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed (attempting again in 10 seconds): %v", name, addr, err)
				// try one more time before giving up
				select {
				case <-time.After(10 * time.Second):
				case <-g.threads.StopChan():
					return
				}
				err := g.managedRPC(addr, name, fn)
				if err != nil {
					g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed twice: %v", name, addr, err)
				}
			}
		}(p.NetAddress)
	}
	wg.Wait()
}
<|endoftext|>


<commit_before>//Package memcached is a memcached store for onecache
package memcached

import (
	"time"

	"github.com/adelowo/onecache"
	"github.com/bradfitz/gomemcache/memcache"
)

type MemcachedStore struct {
	client *memcache.Client
	prefix string
}

//PREFIX prevents collision with other items stored in the db
const PREFIX = "onecache:"

//Returns a new instance of the memached store.
//If prefix is an empty string, it defaults to the package's prefix constant
func NewMemcachedStore(c *memcache.Client, prefix string) *MemcachedStore {

	var p string

	if prefix == "" {
		p = PREFIX
	} else {
		p = prefix
	}

	return &MemcachedStore{client: c, prefix: p}
}

func (m *MemcachedStore) key(k string) string {
	return m.prefix + k
}

func (m *MemcachedStore) Set(k string, data interface{}, expires time.Duration) error {

	i := &onecache.Item{Data: data}

	b, err := i.Bytes()

	if err != nil {
		return err
	}

	item := &memcache.Item{
		Key:        m.key(k),
		Value:      b,
		Expiration: int32(expires / time.Second),
	}

	return m.client.Set(item)
}
<commit_msg>Implemented Get for memcached store<commit_after>//Package memcached is a memcached store for onecache
package memcached

import (
	"time"

	"github.com/adelowo/onecache"
	"github.com/bradfitz/gomemcache/memcache"
)

type MemcachedStore struct {
	client *memcache.Client
	prefix string
}

//PREFIX prevents collision with other items stored in the db
const PREFIX = "onecache:"

//Returns a new instance of the memached store.
//If prefix is an empty string, it defaults to the package's prefix constant
func NewMemcachedStore(c *memcache.Client, prefix string) *MemcachedStore {

	var p string

	if prefix == "" {
		p = PREFIX
	} else {
		p = prefix
	}

	return &MemcachedStore{client: c, prefix: p}
}

func (m *MemcachedStore) key(k string) string {
	return m.prefix + k
}

func (m *MemcachedStore) Set(k string, data interface{}, expires time.Duration) error {

	i := &onecache.Item{Data: data}

	b, err := i.Bytes()

	if err != nil {
		return err
	}

	item := &memcache.Item{
		Key:        m.key(k),
		Value:      b,
		Expiration: int32(expires / time.Second),
	}

	return m.client.Set(item)
}

func (m *MemcachedStore) Get(k string) (interface{}, error) {

	i, err := m.client.Get(m.key(k))

	if err != nil {
		return nil, err
	}

	item, err := onecache.BytesToItem(i.Value)

	if err != nil {
		return nil, err
	}

	return item.Data, nil
}
<|endoftext|>


<commit_before>package host

import (
	"errors"
	"io"
	"net"
	"os"
	"path/filepath"

	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/types"
)

// rpcRetrieve is an RPC that uploads a specified file to a client.
//
// Mutexes are applied carefully to avoid locking during I/O. All necessary
// interaction with the host involves looking up the filepath of the file being
// requested. This is done all at once.
func (h *Host) rpcRetrieve(conn net.Conn) error {
	// Read the contract ID.
	var contractID types.FileContractID
	err := encoding.ReadObject(conn, &contractID, crypto.HashSize)
	if err != nil {
		return err
	}

	// Verify the file exists, using a mutex while reading the host.
	lockID := h.mu.RLock()
	contractObligation, exists := h.obligationsByID[contractID]
	if !exists {
		h.mu.RUnlock(lockID)
		return errors.New("no record of that file")
	}
	path := filepath.Join(h.saveDir, contractObligation.Path)
	h.mu.RUnlock(lockID)

	// Open the file.
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	// Transmit the file.
	_, err = io.CopyN(conn, file, int64(contractObligation.FileContract.FileSize))
	if err != nil {
		return err
	}

	return nil
}

// rpcDownload is an RPC that uploads requested segments of a file. After the
// RPC has been initiated, the host will read and process requests in a loop
// until the 'stop' signal is received or the connection times out.
func (h *Host) rpcDownload(conn net.Conn) error {
	// Read the contract ID.
	var contractID types.FileContractID
	err := encoding.ReadObject(conn, &contractID, crypto.HashSize)
	if err != nil {
		return err
	}

	// Verify the file exists, using a mutex while reading the host.
	lockID := h.mu.RLock()
	contractObligation, exists := h.obligationsByID[contractID]
	if !exists {
		h.mu.RUnlock(lockID)
		return errors.New("no record of that file")
	}
	path := filepath.Join(h.saveDir, contractObligation.Path)
	h.mu.RUnlock(lockID)

	// Open the file.
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	// Process requests until 'stop' signal is received.
	var request struct {
		Offset, Length uint64
	}
	for {
		if err := encoding.ReadObject(conn, &request, 16); err != nil {
			return err
		}
		// Check for termination signal.
		// TODO: perform other sanity checks on offset/length?
		if request.Length == 0 {
			break
		}
		// Write segment to conn.
		segment := io.NewSectionReader(file, int64(request.Offset), int64(request.Length))
		_, err := io.Copy(conn, segment)
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>host upload uses correct path<commit_after>package host

import (
	"errors"
	"io"
	"net"
	"os"
	"path/filepath"

	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

// rpcRetrieve is an RPC that uploads a specified file to a client.
//
// Mutexes are applied carefully to avoid locking during I/O. All necessary
// interaction with the host involves looking up the filepath of the file being
// requested. This is done all at once.
func (h *Host) rpcRetrieve(conn net.Conn) error {
	// Read the contract ID.
	var contractID types.FileContractID
	err := encoding.ReadObject(conn, &contractID, crypto.HashSize)
	if err != nil {
		return err
	}

	// Verify the file exists, using a mutex while reading the host.
	lockID := h.mu.RLock()
	contractObligation, exists := h.obligationsByID[contractID]
	if !exists {
		h.mu.RUnlock(lockID)
		return errors.New("no record of that file")
	}
	path := filepath.Join(h.saveDir, contractObligation.Path)
	h.mu.RUnlock(lockID)

	// Open the file.
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()

	// Transmit the file.
	_, err = io.CopyN(conn, file, int64(contractObligation.FileContract.FileSize))
	if err != nil {
		return err
	}

	return nil
}

// rpcDownload is an RPC that uploads requested segments of a file. After the
// RPC has been initiated, the host will read and process requests in a loop
// until the 'stop' signal is received or the connection times out.
func (h *Host) rpcDownload(conn net.Conn) error {
	// Read the contract ID.
	var contractID types.FileContractID
	err := encoding.ReadObject(conn, &contractID, crypto.HashSize)
	if err != nil {
		return err
	}

	// Verify the file exists, using a mutex while reading the host.
	lockID := h.mu.RLock()
	co, exists := h.obligationsByID[contractID]
	if !exists {
		h.mu.RUnlock(lockID)
		return errors.New("no record of that file")
	}
	h.mu.RUnlock(lockID)

	// Open the file.
	file, err := os.Open(co.Path)
	if err != nil {
		return err
	}
	defer file.Close()

	// Process requests until 'stop' signal is received.
	var request modules.DownloadRequest
	for {
		if err := encoding.ReadObject(conn, &request, 16); err != nil {
			return err
		}
		// Check for termination signal.
		// TODO: perform other sanity checks on offset/length?
		if request.Length == 0 {
			break
		}
		// Write segment to conn.
		segment := io.NewSectionReader(file, int64(request.Offset), int64(request.Length))
		_, err := io.Copy(conn, segment)
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>