<commit_before>package agent

import (
	"fmt"
	"github.com/golang/glog"
	"github.com/romana/core/common"
	"github.com/romana/core/pkg/util/firewall"
)

// addPolicy is a placeholder. TODO
func (a *Agent) addPolicy(input interface{}, ctx common.RestContext) (interface{}, error) {
	//	policy := input.(*common.Policy)
	return nil, nil
}

// deletePolicy is a placeholder. TODO
func (a *Agent) deletePolicy(input interface{}, ctx common.RestContext) (interface{}, error) {
	//	policyId := ctx.PathVariables["policyID"]
	return nil, nil
}

// listPolicies is a placeholder. TODO.
func (a *Agent) listPolicies(input interface{}, ctx common.RestContext) (interface{}, error) {
	return nil, nil
}

// statusHandler reports operational statistics.
func (a *Agent) statusHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	iptablesRules, err := fw.ListRules()
	if err != nil {
		return nil, err
	}
	return iptablesRules, nil
}

// podDownHandler cleans up after pod deleted.
func (a *Agent) podDownHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.V(1).Infoln("Agent: Entering podDownHandler()")
	netReq := input.(*NetworkRequest)
	netif := netReq.NetIf

	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	err = fw.Cleanup(netif)
	if err != nil {
		return nil, err
	}

	// Spawn new thread to process the request
	glog.Infof("Agent: Got request for pod teardown %v\n", netReq)

	return "OK", nil
}

// podUpHandler handles HTTP requests for endpoints provisioning.
func (a *Agent) podUpHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.Infof("Agent: Entering podUpHandler()")
	netReq := input.(*NetworkRequest)

	glog.Infof("Agent: Got request for network configuration: %v\n", netReq)
	// Spawn new thread to process the request

	// TODO don't know if fork-bombs are possible in go but if they are this
	// need to be refactored as buffered channel with fixed pool of workers
	go a.podUpHandlerAsync(*netReq)

	// TODO I wonder if this should actually return something like a
	// link to a status of this request which will later get updated
	// with success or failure -- Greg.
	return "OK", nil
}

// vmDownHandler handles HTTP requests for endpoints teardown.
func (a *Agent) vmDownHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	netif := input.(*NetIf)
	glog.V(1).Infof("In vmDownHandler() with Name %s, IP %s Mac %s\n", netif.Name, netif.IP, netif.Mac)

	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	err = fw.Cleanup(netif)
	if err != nil {
		return nil, err
	}

	return "OK", nil
}

// vmUpHandler handles HTTP requests for endpoints provisioning.
// Currently tested with Romana ML2 driver.
func (a *Agent) vmUpHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	// Parse out NetIf form the request
	netif := input.(*NetIf)

	glog.Infof("Got interface: Name %s, IP %s Mac %s\n", netif.Name, netif.IP, netif.Mac)
	// Spawn new thread to process the request

	// TODO don't know if fork-bombs are possible in go but if they are this
	// need to be refactored as buffered channel with fixed pool of workers
	go a.vmUpHandlerAsync(*netif)

	// TODO I wonder if this should actually return something like a
	// link to a status of this request which will later get updated
	// with success or failure -- Greg.
	return "OK", nil
}

// podUpHandlerAsync does a number of operations on given endpoint to ensure
// it's connected:
// 1. Ensures interface is ready
// 2. Creates ip route pointing new interface
// 3. Provisions firewall rules
func (a *Agent) podUpHandlerAsync(netReq NetworkRequest) error {
	glog.V(1).Info("Agent: Entering podUpHandlerAsync()")

	netif := netReq.NetIf
	if netif.Name == "" {
		return agentErrorString("Agent: Interface name required")
	}
	if !a.Helper.waitForIface(netif.Name) {
		// TODO should we resubmit failed interface in queue for later
		// retry ? ... considering openstack will give up as well after
		// timeout
		msg := fmt.Sprintf("Requested interface not available in time - %s", netif.Name)
		glog.Infoln("Agent: ", msg)
		return agentErrorString(msg)
	}

	glog.Info("Agent: creating endpoint routes")
	if err := a.Helper.ensureRouteToEndpoint(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: provisioning firewall")
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	if err1 := fw.Init(netif); err1 != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	metadata := fw.Metadata()
	chainNames := metadata["chains"].([]string)
	hostAddr := a.networkConfig.RomanaGW()
	hostMask, _ := a.networkConfig.RomanaGWMask().Size()

	// Default firewall rules for Kubernetes
	// Allow ICMP, and SSH between host and instances.
	var defaultRules []firewall.FirewallRule

	// ProvisionEndpoint applies default rules in reverse order
	// so DROP goes first
	inboundChain := chainNames[firewall.InputChainIndex]
	inboundRule := firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s -d %s/%d %s", inboundChain, hostAddr, hostMask, "-j DROP"))
	defaultRules = append(defaultRules, inboundRule)

	inboundRule = firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-p tcp --sport 22 -j ACCEPT"))
	defaultRules = append(defaultRules, inboundRule)

	inboundRule = firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-p icmp --icmp-type 0 -j ACCEPT"))
	defaultRules = append(defaultRules, inboundRule)

	outboundChain := chainNames[firewall.OutputChainIndex]
	outboundRule := firewall.NewFirewallRule()
	outboundRule.SetBody(fmt.Sprintf("%s %s", outboundChain, "-p tcp --dport 22 -j ACCEPT"))
	defaultRules = append(defaultRules, outboundRule)

	forwardInChain := chainNames[firewall.ForwardInChainIndex]
	forwardInRule := firewall.NewFirewallRule()
	forwardInRule.SetBody(fmt.Sprintf("%s %s", forwardInChain, "-m comment --comment Outgoing -j ACCEPT"))
	defaultRules = append(defaultRules, forwardInRule)

	forwardOutChain := chainNames[firewall.ForwardOutChainIndex]
	forwardOutRule := firewall.NewFirewallRule()
	forwardOutRule.SetBody(fmt.Sprintf("%s %s", forwardOutChain, "-m state --state RELATED,ESTABLISHED -j ACCEPT"))
	defaultRules = append(defaultRules, forwardOutRule)

	fw.SetDefaultRules(defaultRules)

	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: All good", netif)
	return nil
}

// vmUpHandlerAsync does a number of operations on given endpoint to ensure
// it's connected:
// 1. Ensures interface is ready
// 2. Checks if DHCP is running
// 3. Creates ip route pointing new interface
// 4. Provisions static DHCP lease for new interface
// 5. Provisions firewall rules
func (a *Agent) vmUpHandlerAsync(netif NetIf) error {
	glog.V(1).Info("Agent: Entering interfaceHandle()")
	if !a.Helper.waitForIface(netif.Name) {
		// TODO should we resubmit failed interface in queue for later
		// retry ? ... considering oenstack will give up as well after
		// timeout
		return agentErrorString(fmt.Sprintf("Requested interface not available in time - %s", netif.Name))
	}

	// dhcpPid is only needed here for fail fast check
	// will try to poll the pid again in provisionLease
	glog.Info("Agent: checking if DHCP is running")
	_, err := a.Helper.DhcpPid()
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	glog.Info("Agent: creating endpoint routes")
	if err := a.Helper.ensureRouteToEndpoint(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	glog.Info("Agent: provisioning DHCP")
	if err := a.leaseFile.provisionLease(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: provisioning firewall")
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	if err1 := fw.Init(netif); err1 != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	metadata := fw.Metadata()
	chainNames := metadata["chains"].([]string)
	u32filter := metadata["u32filter"]
	hostAddr := a.networkConfig.RomanaGW()
	hostMask, _ := a.networkConfig.RomanaGWMask().Size()

	// Default firewall rules for OpenStack
	// Allow ICMP, DHCP and SSH between host and instances.
	inboundChain := chainNames[firewall.InputChainIndex]
	var defaultRules []firewall.FirewallRule

	inboundRule := firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s -d %s/%d %s", inboundChain, hostAddr, hostMask, "-j DROP"))
	defaultRules = append(defaultRules, inboundRule)

	inboundRule = firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-d 255.255.255.255/32 -p udp -m udp --sport 68 --dport 67 -j ACCEPT"))
	defaultRules = append(defaultRules, inboundRule)

	inboundRule = firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-p tcp --sport 22 -j ACCEPT"))
	defaultRules = append(defaultRules, inboundRule)

	inboundRule = firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-p icmp --icmp-type 0 -j ACCEPT"))
	defaultRules = append(defaultRules, inboundRule)

	outboundChain := chainNames[firewall.OutputChainIndex]
	outboundRule := firewall.NewFirewallRule()
	outboundRule.SetBody(fmt.Sprintf("%s -s %s/32 -p udp -m udp --sport 67 --dport 68 -j ACCEPT", outboundChain, hostAddr))
	defaultRules = append(defaultRules, outboundRule)

	outboundRule = firewall.NewFirewallRule()
	outboundRule.SetBody(fmt.Sprintf("%s %s", outboundChain, "-p tcp --dport 22 -j ACCEPT"))
	defaultRules = append(defaultRules, outboundRule)

	forwardInChain := chainNames[firewall.ForwardInChainIndex]
	forwardInRule := firewall.NewFirewallRule()
	forwardInRule.SetBody(fmt.Sprintf("%s %s", forwardInChain, "-m comment --comment Outgoing -j ACCEPT"))
	defaultRules = append(defaultRules, forwardInRule)

	forwardOutChain := chainNames[firewall.ForwardOutChainIndex]
	forwardOutRule := firewall.NewFirewallRule()
	forwardOutRule.SetBody(fmt.Sprintf("%s %s", forwardOutChain, "-m state --state RELATED,ESTABLISHED -j ACCEPT"))
	defaultRules = append(defaultRules, forwardOutRule)

	forwardOutRule = firewall.NewFirewallRule()
	forwardOutRule.SetBody(fmt.Sprintf("%s -m u32 --u32 %s %s", forwardOutChain, u32filter, "-j ACCEPT"))
	defaultRules = append(defaultRules, forwardOutRule)

	fw.SetDefaultRules(defaultRules)

	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("All good", netif)
	return nil
}
<commit_msg>Remove ssh access rules for kubernetes<commit_after>package agent

import (
	"fmt"
	"github.com/golang/glog"
	"github.com/romana/core/common"
	"github.com/romana/core/pkg/util/firewall"
)

// addPolicy is a placeholder. TODO
func (a *Agent) addPolicy(input interface{}, ctx common.RestContext) (interface{}, error) {
	//	policy := input.(*common.Policy)
	return nil, nil
}

// deletePolicy is a placeholder. TODO
func (a *Agent) deletePolicy(input interface{}, ctx common.RestContext) (interface{}, error) {
	//	policyId := ctx.PathVariables["policyID"]
	return nil, nil
}

// listPolicies is a placeholder. TODO.
func (a *Agent) listPolicies(input interface{}, ctx common.RestContext) (interface{}, error) {
	return nil, nil
}

// statusHandler reports operational statistics.
func (a *Agent) statusHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	iptablesRules, err := fw.ListRules()
	if err != nil {
		return nil, err
	}
	return iptablesRules, nil
}

// podDownHandler cleans up after pod deleted.
func (a *Agent) podDownHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.V(1).Infoln("Agent: Entering podDownHandler()")
	netReq := input.(*NetworkRequest)
	netif := netReq.NetIf

	// We need new firewall instance here to use it's Cleanup()
	// to uninstall firewall rules related to the endpoint.
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	err = fw.Cleanup(netif)
	if err != nil {
		return nil, err
	}

	// Spawn new thread to process the request
	glog.Infof("Agent: Got request for pod teardown %v\n", netReq)

	return "OK", nil
}

// podUpHandler handles HTTP requests for endpoints provisioning.
func (a *Agent) podUpHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	glog.Infof("Agent: Entering podUpHandler()")
	netReq := input.(*NetworkRequest)

	glog.Infof("Agent: Got request for network configuration: %v\n", netReq)
	// Spawn new thread to process the request

	// TODO don't know if fork-bombs are possible in go but if they are this
	// need to be refactored as buffered channel with fixed pool of workers
	go a.podUpHandlerAsync(*netReq)

	// TODO I wonder if this should actually return something like a
	// link to a status of this request which will later get updated
	// with success or failure -- Greg.
	return "OK", nil
}

// vmDownHandler handles HTTP requests for endpoints teardown.
func (a *Agent) vmDownHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	netif := input.(*NetIf)
	glog.V(1).Infof("In vmDownHandler() with Name %s, IP %s Mac %s\n", netif.Name, netif.IP, netif.Mac)

	// We need new firewall instance here to use it's Cleanup()
	// to uninstall firewall rules related to the endpoint.
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		return nil, err
	}

	err = fw.Cleanup(netif)
	if err != nil {
		return nil, err
	}

	return "OK", nil
}

// vmUpHandler handles HTTP requests for endpoints provisioning.
// Currently tested with Romana ML2 driver.
func (a *Agent) vmUpHandler(input interface{}, ctx common.RestContext) (interface{}, error) {
	// Parse out NetIf form the request
	netif := input.(*NetIf)

	glog.Infof("Got interface: Name %s, IP %s Mac %s\n", netif.Name, netif.IP, netif.Mac)
	// Spawn new thread to process the request

	// TODO don't know if fork-bombs are possible in go but if they are this
	// need to be refactored as buffered channel with fixed pool of workers
	go a.vmUpHandlerAsync(*netif)

	// TODO I wonder if this should actually return something like a
	// link to a status of this request which will later get updated
	// with success or failure -- Greg.
	return "OK", nil
}

// podUpHandlerAsync does a number of operations on given endpoint to ensure
// it's connected:
// 1. Ensures interface is ready
// 2. Creates ip route pointing new interface
// 3. Provisions firewall rules
func (a *Agent) podUpHandlerAsync(netReq NetworkRequest) error {
	glog.V(1).Info("Agent: Entering podUpHandlerAsync()")

	netif := netReq.NetIf
	if netif.Name == "" {
		return agentErrorString("Agent: Interface name required")
	}
	if !a.Helper.waitForIface(netif.Name) {
		// TODO should we resubmit failed interface in queue for later
		// retry ? ... considering openstack will give up as well after
		// timeout
		msg := fmt.Sprintf("Requested interface not available in time - %s", netif.Name)
		glog.Infoln("Agent: ", msg)
		return agentErrorString(msg)
	}

	glog.Info("Agent: creating endpoint routes")
	if err := a.Helper.ensureRouteToEndpoint(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: provisioning firewall")
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	if err1 := fw.Init(netif); err1 != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	metadata := fw.Metadata()
	chainNames := metadata["chains"].([]string)
	hostAddr := a.networkConfig.RomanaGW()
	hostMask, _ := a.networkConfig.RomanaGWMask().Size()

	// Default firewall rules for Kubernetes
	// Allow ICMP, and SSH between host and instances.
	var defaultRules []firewall.FirewallRule

	// ProvisionEndpoint applies default rules in reverse order
	// so DROP goes first
	inboundChain := chainNames[firewall.InputChainIndex]
	inboundRule := firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s -d %s/%d %s", inboundChain, hostAddr, hostMask, "-j DROP"))
	defaultRules = append(defaultRules, inboundRule)

	inboundRule = firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-p icmp --icmp-type 0 -j ACCEPT"))
	defaultRules = append(defaultRules, inboundRule)

	forwardInChain := chainNames[firewall.ForwardInChainIndex]
	forwardInRule := firewall.NewFirewallRule()
	forwardInRule.SetBody(fmt.Sprintf("%s %s", forwardInChain, "-m comment --comment Outgoing -j ACCEPT"))
	defaultRules = append(defaultRules, forwardInRule)

	forwardOutChain := chainNames[firewall.ForwardOutChainIndex]
	forwardOutRule := firewall.NewFirewallRule()
	forwardOutRule.SetBody(fmt.Sprintf("%s %s", forwardOutChain, "-m state --state RELATED,ESTABLISHED -j ACCEPT"))
	defaultRules = append(defaultRules, forwardOutRule)

	fw.SetDefaultRules(defaultRules)

	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: All good", netif)
	return nil
}

// vmUpHandlerAsync does a number of operations on given endpoint to ensure
// it's connected:
// 1. Ensures interface is ready
// 2. Checks if DHCP is running
// 3. Creates ip route pointing new interface
// 4. Provisions static DHCP lease for new interface
// 5. Provisions firewall rules
func (a *Agent) vmUpHandlerAsync(netif NetIf) error {
	glog.V(1).Info("Agent: Entering interfaceHandle()")
	if !a.Helper.waitForIface(netif.Name) {
		// TODO should we resubmit failed interface in queue for later
		// retry ? ... considering oenstack will give up as well after
		// timeout
		return agentErrorString(fmt.Sprintf("Requested interface not available in time - %s", netif.Name))
	}

	// dhcpPid is only needed here for fail fast check
	// will try to poll the pid again in provisionLease
	glog.Info("Agent: checking if DHCP is running")
	_, err := a.Helper.DhcpPid()
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	glog.Info("Agent: creating endpoint routes")
	if err := a.Helper.ensureRouteToEndpoint(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	glog.Info("Agent: provisioning DHCP")
	if err := a.leaseFile.provisionLease(&netif); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("Agent: provisioning firewall")
	fw, err := firewall.NewFirewall(a.Helper.Executor, a.store, a.networkConfig)
	if err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	if err1 := fw.Init(netif); err1 != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	metadata := fw.Metadata()
	chainNames := metadata["chains"].([]string)
	u32filter := metadata["u32filter"]
	hostAddr := a.networkConfig.RomanaGW()
	hostMask, _ := a.networkConfig.RomanaGWMask().Size()

	// Default firewall rules for OpenStack
	// Allow ICMP, DHCP and SSH between host and instances.
	inboundChain := chainNames[firewall.InputChainIndex]
	var defaultRules []firewall.FirewallRule

	inboundRule := firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s -d %s/%d %s", inboundChain, hostAddr, hostMask, "-j DROP"))
	defaultRules = append(defaultRules, inboundRule)

	inboundRule = firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-d 255.255.255.255/32 -p udp -m udp --sport 68 --dport 67 -j ACCEPT"))
	defaultRules = append(defaultRules, inboundRule)

	inboundRule = firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-p tcp --sport 22 -j ACCEPT"))
	defaultRules = append(defaultRules, inboundRule)

	inboundRule = firewall.NewFirewallRule()
	inboundRule.SetBody(fmt.Sprintf("%s %s", inboundChain, "-p icmp --icmp-type 0 -j ACCEPT"))
	defaultRules = append(defaultRules, inboundRule)

	outboundChain := chainNames[firewall.OutputChainIndex]
	outboundRule := firewall.NewFirewallRule()
	outboundRule.SetBody(fmt.Sprintf("%s -s %s/32 -p udp -m udp --sport 67 --dport 68 -j ACCEPT", outboundChain, hostAddr))
	defaultRules = append(defaultRules, outboundRule)

	outboundRule = firewall.NewFirewallRule()
	outboundRule.SetBody(fmt.Sprintf("%s %s", outboundChain, "-p tcp --dport 22 -j ACCEPT"))
	defaultRules = append(defaultRules, outboundRule)

	forwardInChain := chainNames[firewall.ForwardInChainIndex]
	forwardInRule := firewall.NewFirewallRule()
	forwardInRule.SetBody(fmt.Sprintf("%s %s", forwardInChain, "-m comment --comment Outgoing -j ACCEPT"))
	defaultRules = append(defaultRules, forwardInRule)

	forwardOutChain := chainNames[firewall.ForwardOutChainIndex]
	forwardOutRule := firewall.NewFirewallRule()
	forwardOutRule.SetBody(fmt.Sprintf("%s %s", forwardOutChain, "-m state --state RELATED,ESTABLISHED -j ACCEPT"))
	defaultRules = append(defaultRules, forwardOutRule)

	forwardOutRule = firewall.NewFirewallRule()
	forwardOutRule.SetBody(fmt.Sprintf("%s -m u32 --u32 %s %s", forwardOutChain, u32filter, "-j ACCEPT"))
	defaultRules = append(defaultRules, forwardOutRule)

	fw.SetDefaultRules(defaultRules)

	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}
	if err := fw.ProvisionEndpoint(); err != nil {
		glog.Error(agentError(err))
		return agentError(err)
	}

	glog.Info("All good", netif)
	return nil
}
<|endoftext|>


<commit_before>package core

import (
	"crypto/rsa"
	"crypto/tls"
	"fmt"
	jwt "github.com/dgrijalva/jwt-go"
	"github.com/gorilla/sessions"
	"github.com/mendsley/gojwk"
	hydra "github.com/ory-am/hydra/sdk"
	"github.com/patrickmn/go-cache"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/clientcredentials"
	"io/ioutil"
	"net/http"
	"time"
)

const (
	VerifyPublicKey   = "VerifyPublic"
	ConsentPrivateKey = "ConsentPrivate"
	ClientInfo        = "ClientInfo"
)

var encryptionkey = "something-very-secret"

type IDPConfig struct {
	ClientID              string        `yaml:"client_id"`
	ClientSecret          string        `yaml:"client_secret"`
	ClusterURL            string        `yaml:"hydra_address"`
	KeyCacheExpiration    time.Duration `yaml:"key_cache_expiration"`
	ClientCacheExpiration time.Duration `yaml:"client_cache_expiration"`
	CacheCleanupInterval  time.Duration `yaml:"cache_cleanup_interval"`
	ChallengeStore        sessions.Store
}

type IDP struct {
	config *IDPConfig

	// Communication with Hydra
	hc *hydra.Client

	// Http client for communicating with Hydra
	client *http.Client

	// Cache for all private and public keys
	cache *cache.Cache
}

func NewIDP(config *IDPConfig) *IDP {
	var idp = new(IDP)
	idp.config = config

	// TODO: Pass TTL and refresh period from config
	idp.cache = cache.New(config.KeyCacheExpiration, config.CacheCleanupInterval)
	idp.cache.OnEvicted(func(key string, value interface{}) { idp.refreshCache(key) })

	return idp
}

// Called when any key expires
func (idp *IDP) refreshCache(key string) {
	switch key {
	case VerifyPublicKey:
		verifyKey, err := idp.getVerificationKey()
		if err != nil {
			return
		}
		idp.cache.Set(VerifyPublicKey, verifyKey, cache.DefaultExpiration)
		return

	case ConsentPrivateKey:
		consentKey, err := idp.getConsentKey()
		if err != nil {
			return
		}
		idp.cache.Set(ConsentPrivateKey, consentKey, cache.DefaultExpiration)
		return

	case ClientInfo:
		// consentKey, err := idp.getConsentKey()
		// if err != nil {
		// 	return
		// }
		// idp.cache.Set(ClientInfo, consentKey, idp.config.ClientCacheExpiration)
		return

	default:
		return
	}
}

// Gets the requested key from Hydra
func (idp *IDP) getKey(set string, kind string) (*gojwk.Key, error) {
	url := idp.config.ClusterURL + "/keys/" + set + "/" + kind

	resp, err := idp.client.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	key, err := gojwk.Unmarshal(body)
	if err != nil {
		return nil, err
	}

	return key.Keys[0], nil
}

// Downloads the hydra's public key
func (idp *IDP) getVerificationKey() (*rsa.PublicKey, error) {
	jwk, err := idp.getKey("consent.challenge", "public")
	if err != nil {
		return nil, err
	}

	key, err := jwk.DecodePublicKey()
	if err != nil {
		return nil, err
	}

	return key.(*rsa.PublicKey), err
}

// Downloads the private key used for signing the consent
func (idp *IDP) getConsentKey() (*rsa.PrivateKey, error) {
	jwk, err := idp.getKey("consent.endpoint", "private")
	if err != nil {
		return nil, err
	}

	key, err := jwk.DecodePrivateKey()
	if err != nil {
		return nil, err
	}

	return key.(*rsa.PrivateKey), err
}

func (idp *IDP) login() error {
	// Use the credentials to login to Hydra
	credentials := clientcredentials.Config{
		ClientID:     idp.config.ClientID,
		ClientSecret: idp.config.ClientSecret,
		TokenURL:     idp.config.ClusterURL + "/oauth2/token",
		Scopes:       []string{"core", "hydra.keys.get"},
	}

	// Skip verifying the certificate
	// TODO: Remove when Hydra implements passing key-cert pairs
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	c := &http.Client{Transport: tr}
	ctx := context.WithValue(oauth2.NoContext, oauth2.HTTPClient, c)

	// Prefetch the token - tests the connection``
	_, err := credentials.Token(ctx)
	if err != nil {
		return err
	}

	idp.client = credentials.Client(ctx)

	return nil
}

func (idp *IDP) Connect() error {
	var err error
	idp.hc, err = hydra.Connect(
		hydra.ClientID(idp.config.ClientID),
		hydra.ClientSecret(idp.config.ClientSecret),
		hydra.ClusterURL(idp.config.ClusterURL),
	)

	err = idp.login()
	if err != nil {
		return err
	}

	verifyKey, err := idp.getVerificationKey()
	if err != nil {
		return err
	}

	consentKey, err := idp.getConsentKey()
	if err != nil {
		return err
	}

	idp.cache.Set(VerifyPublicKey, verifyKey, cache.DefaultExpiration)
	idp.cache.Set(ConsentPrivateKey, consentKey, cache.DefaultExpiration)

	return err
}

// Parse and verify the challenge JWT
func (idp *IDP) getChallengeToken(challengeString string) (*jwt.Token, error) {
	token, err := jwt.Parse(challengeString, func(token *jwt.Token) (interface{}, error) {
		_, ok := token.Method.(*jwt.SigningMethodRSA)
		if !ok {
			return nil, fmt.Errorf("Unexpected signing method: %v", token.Header["alg"])
		}

		return idp.GetVerificationKey()
	})

	if err != nil {
		return nil, err
	}

	if !token.Valid {
		return nil, fmt.Errorf("Empty token")
	}

	return token, nil
}

func (idp *IDP) GetConsentKey() (*rsa.PrivateKey, error) {
	data, ok := idp.cache.Get(ConsentPrivateKey)
	if !ok {
		return nil, ErrorNoKey
	}

	key, ok := data.(*rsa.PrivateKey)
	if !ok {
		return nil, ErrorBadKey
	}

	return key, nil
}

func (idp *IDP) GetVerificationKey() (*rsa.PublicKey, error) {
	data, ok := idp.cache.Get(VerifyPublicKey)
	if !ok {
		return nil, ErrorNoKey
	}

	key, ok := data.(*rsa.PublicKey)
	if !ok {
		return nil, ErrorBadKey
	}

	return key, nil
}

func (idp *IDP) NewChallenge(r *http.Request, user string) (challenge *Challenge, err error) {
	tokenStr := r.FormValue("challenge")
	if tokenStr == "" {
		// No challenge token
		err = ErrorBadRequest
		return
	}

	token, err := idp.getChallengeToken(tokenStr)
	if err != nil {
		// Most probably, token can't be verified or parsed
		return
	}
	claims := token.Claims.(jwt.MapClaims)

	challenge = new(Challenge)
	challenge.Expires = time.Unix(int64(claims["exp"].(float64)), 0)
	if challenge.Expires.Before(time.Now()) {
		challenge = nil
		err = ErrorChallengeExpired
		return
	}

	// Get data from the challenge jwt
	challenge.Client = claims["aud"].(string)
	challenge.Redirect = claims["redir"].(string)

	challenge.User = user
	challenge.idp = idp

	scopes := claims["scp"].([]interface{})
	challenge.Scopes = make([]string, len(scopes), len(scopes))
	for i, scope := range scopes {
		challenge.Scopes[i] = scope.(string)
	}

	return
}

func (idp *IDP) GetChallenge(r *http.Request) (*Challenge, error) {
	session, err := idp.config.ChallengeStore.Get(r, SessionCookieName)
	if err != nil {
		return nil, err
	}

	challenge, ok := session.Values[SessionCookieName].(*Challenge)
	if !ok {
		return nil, ErrorBadChallengeCookie
	}

	if challenge.Expires.Before(time.Now()) {
		return nil, ErrorChallengeExpired
	}

	challenge.idp = idp

	return challenge, nil
}

func (idp *IDP) Close() {
	fmt.Println("IDP closed")
	idp.client = nil

	// Removes all keys from the cache
	idp.cache.Flush()
}
<commit_msg>core: use ory-am/hydra SDK<commit_after>package core

import (
	"crypto/rsa"
	"errors"
	"fmt"
	"net/http"
	"time"

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/gorilla/sessions"
	hjwk "github.com/ory-am/hydra/jwk"
	hoauth2 "github.com/ory-am/hydra/oauth2"
	hydra "github.com/ory-am/hydra/sdk"
	"github.com/patrickmn/go-cache"
)

const (
	VerifyPublicKey   = "VerifyPublic"
	ConsentPrivateKey = "ConsentPrivate"
	ClientInfo        = "ClientInfo"
)

var encryptionkey = "something-very-secret"

type IDPConfig struct {
	ClientID              string        `yaml:"client_id"`
	ClientSecret          string        `yaml:"client_secret"`
	ClusterURL            string        `yaml:"hydra_address"`
	KeyCacheExpiration    time.Duration `yaml:"key_cache_expiration"`
	ClientCacheExpiration time.Duration `yaml:"client_cache_expiration"`
	CacheCleanupInterval  time.Duration `yaml:"cache_cleanup_interval"`
	ChallengeStore        sessions.Store
}

type IDP struct {
	config *IDPConfig

	// Communication with Hydra
	hc *hydra.Client

	// Http client for communicating with Hydra
	client *http.Client

	// Cache for all private and public keys
	cache *cache.Cache
}

func NewIDP(config *IDPConfig) *IDP {
	var idp = new(IDP)
	idp.config = config

	// TODO: Pass TTL and refresh period from config
	idp.cache = cache.New(config.KeyCacheExpiration, config.CacheCleanupInterval)
	idp.cache.OnEvicted(func(key string, value interface{}) { idp.refreshCache(key) })

	return idp
}

// Called when any key expires
func (idp *IDP) refreshCache(key string) {
	switch key {
	case VerifyPublicKey:
		verifyKey, err := idp.getVerificationKey()
		if err != nil {
			return
		}
		idp.cache.Set(VerifyPublicKey, verifyKey, cache.DefaultExpiration)
		return

	case ConsentPrivateKey:
		consentKey, err := idp.getConsentKey()
		if err != nil {
			return
		}
		idp.cache.Set(ConsentPrivateKey, consentKey, cache.DefaultExpiration)
		return

	case ClientInfo:
		// consentKey, err := idp.getConsentKey()
		// if err != nil {
		// 	return
		// }
		// idp.cache.Set(ClientInfo, consentKey, idp.config.ClientCacheExpiration)
		return

	default:
		return
	}
}

// Downloads the hydra's public key
func (idp *IDP) getVerificationKey() (*rsa.PublicKey, error) {

	jwk, err := idp.hc.JWK.GetKey(hoauth2.ConsentChallengeKey, "public")
	if err != nil {
		return nil, err
	}

	rsaKey, ok := hjwk.First(jwk.Keys).Key.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("Could not convert to RSA Public Key")
	}

	return rsaKey, nil
}

// Downloads the private key used for signing the consent
func (idp *IDP) getConsentKey() (*rsa.PrivateKey, error) {
	jwk, err := idp.hc.JWK.GetKey(hoauth2.ConsentEndpointKey, "private")
	if err != nil {
		return nil, err
	}

	rsaKey, ok := hjwk.First(jwk.Keys).Key.(*rsa.PrivateKey)
	if !ok {
		return nil, errors.New("Could not convert to RSA Private Key")
	}

	return rsaKey, nil
}

func (idp *IDP) Connect() error {
	var err error
	idp.hc, err = hydra.Connect(
		hydra.ClientID(idp.config.ClientID),
		hydra.ClientSecret(idp.config.ClientSecret),
		hydra.ClusterURL(idp.config.ClusterURL),
		hydra.SkipTLSVerify(),
	)

	if err != nil {
		return err
	}

	verifyKey, err := idp.getVerificationKey()
	if err != nil {
		return err
	}

	consentKey, err := idp.getConsentKey()
	if err != nil {
		return err
	}

	idp.cache.Set(VerifyPublicKey, verifyKey, cache.DefaultExpiration)
	idp.cache.Set(ConsentPrivateKey, consentKey, cache.DefaultExpiration)

	return err
}

// Parse and verify the challenge JWT
func (idp *IDP) getChallengeToken(challengeString string) (*jwt.Token, error) {
	token, err := jwt.Parse(challengeString, func(token *jwt.Token) (interface{}, error) {
		_, ok := token.Method.(*jwt.SigningMethodRSA)
		if !ok {
			return nil, fmt.Errorf("Unexpected signing method: %v", token.Header["alg"])
		}

		return idp.GetVerificationKey()
	})

	if err != nil {
		return nil, err
	}

	if !token.Valid {
		return nil, fmt.Errorf("Empty token")
	}

	return token, nil
}

func (idp *IDP) GetConsentKey() (*rsa.PrivateKey, error) {
	data, ok := idp.cache.Get(ConsentPrivateKey)
	if !ok {
		return nil, ErrorNoKey
	}

	key, ok := data.(*rsa.PrivateKey)
	if !ok {
		return nil, ErrorBadKey
	}

	return key, nil
}

func (idp *IDP) GetVerificationKey() (*rsa.PublicKey, error) {
	data, ok := idp.cache.Get(VerifyPublicKey)
	if !ok {
		return nil, ErrorNoKey
	}

	key, ok := data.(*rsa.PublicKey)
	if !ok {
		return nil, ErrorBadKey
	}

	return key, nil
}

func (idp *IDP) NewChallenge(r *http.Request, user string) (challenge *Challenge, err error) {
	tokenStr := r.FormValue("challenge")
	if tokenStr == "" {
		// No challenge token
		err = ErrorBadRequest
		return
	}

	token, err := idp.getChallengeToken(tokenStr)
	if err != nil {
		// Most probably, token can't be verified or parsed
		return
	}
	claims := token.Claims.(jwt.MapClaims)

	challenge = new(Challenge)
	challenge.Expires = time.Unix(int64(claims["exp"].(float64)), 0)
	if challenge.Expires.Before(time.Now()) {
		challenge = nil
		err = ErrorChallengeExpired
		return
	}

	// Get data from the challenge jwt
	challenge.Client = claims["aud"].(string)
	challenge.Redirect = claims["redir"].(string)

	challenge.User = user
	challenge.idp = idp

	scopes := claims["scp"].([]interface{})
	challenge.Scopes = make([]string, len(scopes), len(scopes))
	for i, scope := range scopes {
		challenge.Scopes[i] = scope.(string)
	}

	return
}

func (idp *IDP) GetChallenge(r *http.Request) (*Challenge, error) {
	session, err := idp.config.ChallengeStore.Get(r, SessionCookieName)
	if err != nil {
		return nil, err
	}

	challenge, ok := session.Values[SessionCookieName].(*Challenge)
	if !ok {
		return nil, ErrorBadChallengeCookie
	}

	if challenge.Expires.Before(time.Now()) {
		return nil, ErrorChallengeExpired
	}

	challenge.idp = idp

	return challenge, nil
}

func (idp *IDP) Close() {
	fmt.Println("IDP closed")
	idp.client = nil

	// Removes all keys from the cache
	idp.cache.Flush()
}
<|endoftext|>


<commit_before>package handler

import (
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"strings"

	"github.com/dchest/safefile"
	"github.com/go-chi/chi"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/webhippie/terrastate/pkg/config"
)

// Update is used to update a specific state.
func Update(logger log.Logger) http.HandlerFunc {
	logger = log.WithPrefix(logger, "handler", "update")

	return func(w http.ResponseWriter, req *http.Request) {
		dir := strings.Replace(
			path.Join(
				config.Server.Storage,
				chi.URLParam(req, "*"),
			),
			"../", "", -1,
		)

		full := path.Join(
			dir,
			"terraform.tfstate",
		)

		// TODO: ID param can be a lock id
		level.Info(logger).Log(
			"msg", "debugging",
			"id", req.URL.Query().Get("ID"),
		)

		content, err := ioutil.ReadAll(req.Body)

		if err != nil {
			level.Info(logger).Log(
				"msg", "failed to load request body",
				"err", err,
			)

			http.Error(
				w,
				http.StatusText(http.StatusInternalServerError),
				http.StatusInternalServerError,
			)

			return
		}

		if err := os.MkdirAll(dir, 0755); err != nil {
			level.Info(logger).Log(
				"msg", "failed to create state dir",
				"dir", dir,
				"err", err,
			)

			http.Error(
				w,
				http.StatusText(http.StatusInternalServerError),
				http.StatusInternalServerError,
			)

			return
		}

		if _, err := os.Stat(full); os.IsNotExist(err) {
			if err := safefile.WriteFile(full, content, 0644); err != nil {
				level.Info(logger).Log(
					"msg", "failed to create state file",
					"err", err,
				)

				http.Error(
					w,
					http.StatusText(http.StatusInternalServerError),
					http.StatusInternalServerError,
				)

				return
			}

			level.Info(logger).Log(
				"msg", "successfully created state file",
				"file", full,
			)
		} else {
			if err := safefile.WriteFile(full, content, 0644); err != nil {
				level.Info(logger).Log(
					"msg", "failed to update state file",
					"err", err,
				)

				http.Error(
					w,
					http.StatusText(http.StatusInternalServerError),
					http.StatusInternalServerError,
				)

				return
			}

			level.Info(logger).Log(
				"msg", "successfully updated state file",
				"file", full,
			)
		}

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
	}
}
<commit_msg>Ignore lock id within update for now<commit_after>package handler

import (
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"strings"

	"github.com/dchest/safefile"
	"github.com/go-chi/chi"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/webhippie/terrastate/pkg/config"
)

// Update is used to update a specific state.
func Update(logger log.Logger) http.HandlerFunc {
	logger = log.WithPrefix(logger, "handler", "update")

	return func(w http.ResponseWriter, req *http.Request) {
		dir := strings.Replace(
			path.Join(
				config.Server.Storage,
				chi.URLParam(req, "*"),
			),
			"../", "", -1,
		)

		full := path.Join(
			dir,
			"terraform.tfstate",
		)

		content, err := ioutil.ReadAll(req.Body)

		if err != nil {
			level.Info(logger).Log(
				"msg", "failed to load request body",
				"err", err,
			)

			http.Error(
				w,
				http.StatusText(http.StatusInternalServerError),
				http.StatusInternalServerError,
			)

			return
		}

		if err := os.MkdirAll(dir, 0755); err != nil {
			level.Info(logger).Log(
				"msg", "failed to create state dir",
				"dir", dir,
				"err", err,
			)

			http.Error(
				w,
				http.StatusText(http.StatusInternalServerError),
				http.StatusInternalServerError,
			)

			return
		}

		if _, err := os.Stat(full); os.IsNotExist(err) {
			if err := safefile.WriteFile(full, content, 0644); err != nil {
				level.Info(logger).Log(
					"msg", "failed to create state file",
					"err", err,
				)

				http.Error(
					w,
					http.StatusText(http.StatusInternalServerError),
					http.StatusInternalServerError,
				)

				return
			}

			level.Info(logger).Log(
				"msg", "successfully created state file",
				"file", full,
			)
		} else {
			if err := safefile.WriteFile(full, content, 0644); err != nil {
				level.Info(logger).Log(
					"msg", "failed to update state file",
					"err", err,
				)

				http.Error(
					w,
					http.StatusText(http.StatusInternalServerError),
					http.StatusInternalServerError,
				)

				return
			}

			level.Info(logger).Log(
				"msg", "successfully updated state file",
				"file", full,
			)
		}

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
	}
}
<|endoftext|>


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

package frontend

import (
	"context"
	"fmt"
	"math"
	"net/http"
	"path"
	"strings"

	"golang.org/x/discovery/internal"
	"golang.org/x/discovery/internal/derrors"
	"golang.org/x/discovery/internal/log"
	"golang.org/x/discovery/internal/postgres"
	"golang.org/x/xerrors"
)

const defaultSearchLimit = 10

// SearchPage contains all of the data that the search template needs to
// populate.
type SearchPage struct {
	basePage
	Pagination pagination
	Results    []*SearchResult
}

// SearchResult contains data needed to display a single search result.
type SearchResult struct {
	Name           string
	PackagePath    string
	ModulePath     string
	Synopsis       string
	DisplayVersion string
	Licenses       []string
	CommitTime     string
	NumImportedBy  uint64
	Approximate    bool
}

// fetchSearchPage fetches data matching the search query from the database and
// returns a SearchPage.
func fetchSearchPage(ctx context.Context, ds DataSource, query, method string, pageParams paginationParams) (*SearchPage, error) {
	var (
		dbresults []*postgres.SearchResult
		err       error
	)
	switch method {
	case "slow":
		dbresults, err = ds.Search(ctx, query, pageParams.limit, pageParams.offset())
	case "deep":
		dbresults, err = ds.DeepSearch(ctx, query, pageParams.limit, pageParams.offset())
	case "partial-fast":
		dbresults, err = ds.PartialFastSearch(ctx, query, pageParams.limit, pageParams.offset())
	case "popular":
		dbresults, err = ds.PopularSearch(ctx, query, pageParams.limit, pageParams.offset())
	default:
		dbresults, err = ds.FastSearch(ctx, query, pageParams.limit, pageParams.offset())
	}
	if err != nil {
		return nil, err
	}

	var results []*SearchResult
	for _, r := range dbresults {
		results = append(results, &SearchResult{
			Name:           r.Name,
			PackagePath:    r.PackagePath,
			ModulePath:     r.ModulePath,
			Synopsis:       r.Synopsis,
			DisplayVersion: displayVersion(r.Version, r.ModulePath),
			Licenses:       r.Licenses,
			CommitTime:     elapsedTime(r.CommitTime),
			NumImportedBy:  r.NumImportedBy,
		})
	}

	var (
		numResults  int
		approximate bool
	)
	if len(dbresults) > 0 {
		numResults = int(dbresults[0].NumResults)
		if dbresults[0].Approximate {
			// 128 buckets corresponds to a standard error of 10%.
			// http://algo.inria.fr/flajolet/Publications/FlFuGaMe07.pdf
			numResults = approximateNumber(numResults, 0.1)
			approximate = true
		}
	}

	pgs := newPagination(pageParams, len(results), numResults)
	pgs.Approximate = approximate
	return &SearchPage{
		Results:    results,
		Pagination: pgs,
	}, nil
}

// approximateNumber returns an approximation of the estimate, calibrated by
// the statistical estimate of standard error.
// i.e., a number that isn't misleading when we say '1-10 of approximately N
// results', but that is still close to our estimate.
func approximateNumber(estimate int, sigma float64) int {
	expectedErr := sigma * float64(estimate)
	// Compute the unit by rounding the error the logarithmically closest power
	// of 10, so that 300->100, but 400->1000.
	unit := math.Pow(10, math.Round(math.Log10(expectedErr)))
	// Now round the estimate to the nearest unit.
	return int(unit * math.Round(float64(estimate)/unit))
}

// handleSearch applies database data to the search template. Handles endpoint
// /search?q=<query>. If <query> is an exact match for a package path, the user
// will be redirected to the details page.
func (s *Server) handleSearch(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	query := searchQuery(r)
	if query == "" {
		http.Redirect(w, r, "/", http.StatusFound)
		return
	}

	if strings.Contains(query, "/") {
		pkg, err := s.ds.GetPackage(ctx, path.Clean(query), internal.UnknownModulePath, internal.LatestVersion)
		if err == nil {
			http.Redirect(w, r, fmt.Sprintf("/%s?tab=overview", pkg.Path), http.StatusFound)
			return
		} else if !xerrors.Is(err, derrors.NotFound) {
			log.Errorf("error getting package for %s: %v", path.Clean(query), err)
		}
	}

	searchMethod := r.FormValue("method")
	page, err := fetchSearchPage(ctx, s.ds, query, searchMethod, newPaginationParams(r, defaultSearchLimit))
	if err != nil {
		log.Errorf("fetchSearchDetails(ctx, db, %q): %v", query, err)
		s.serveErrorPage(w, r, http.StatusInternalServerError, nil)
		return
	}
	page.basePage = newBasePage(r, query)
	s.servePage(w, "search.tmpl", page)
}

// searchQuery extracts a search query from the request.
func searchQuery(r *http.Request) string {
	return strings.TrimSpace(r.FormValue("q"))
}
<commit_msg>Revert "internal/frontend: link search results to overview tab"<commit_after>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package frontend

import (
	"context"
	"fmt"
	"math"
	"net/http"
	"path"
	"strings"

	"golang.org/x/discovery/internal"
	"golang.org/x/discovery/internal/derrors"
	"golang.org/x/discovery/internal/log"
	"golang.org/x/discovery/internal/postgres"
	"golang.org/x/xerrors"
)

const defaultSearchLimit = 10

// SearchPage contains all of the data that the search template needs to
// populate.
type SearchPage struct {
	basePage
	Pagination pagination
	Results    []*SearchResult
}

// SearchResult contains data needed to display a single search result.
type SearchResult struct {
	Name           string
	PackagePath    string
	ModulePath     string
	Synopsis       string
	DisplayVersion string
	Licenses       []string
	CommitTime     string
	NumImportedBy  uint64
	Approximate    bool
}

// fetchSearchPage fetches data matching the search query from the database and
// returns a SearchPage.
func fetchSearchPage(ctx context.Context, ds DataSource, query, method string, pageParams paginationParams) (*SearchPage, error) {
	var (
		dbresults []*postgres.SearchResult
		err       error
	)
	switch method {
	case "slow":
		dbresults, err = ds.Search(ctx, query, pageParams.limit, pageParams.offset())
	case "deep":
		dbresults, err = ds.DeepSearch(ctx, query, pageParams.limit, pageParams.offset())
	case "partial-fast":
		dbresults, err = ds.PartialFastSearch(ctx, query, pageParams.limit, pageParams.offset())
	case "popular":
		dbresults, err = ds.PopularSearch(ctx, query, pageParams.limit, pageParams.offset())
	default:
		dbresults, err = ds.FastSearch(ctx, query, pageParams.limit, pageParams.offset())
	}
	if err != nil {
		return nil, err
	}

	var results []*SearchResult
	for _, r := range dbresults {
		results = append(results, &SearchResult{
			Name:           r.Name,
			PackagePath:    r.PackagePath,
			ModulePath:     r.ModulePath,
			Synopsis:       r.Synopsis,
			DisplayVersion: displayVersion(r.Version, r.ModulePath),
			Licenses:       r.Licenses,
			CommitTime:     elapsedTime(r.CommitTime),
			NumImportedBy:  r.NumImportedBy,
		})
	}

	var (
		numResults  int
		approximate bool
	)
	if len(dbresults) > 0 {
		numResults = int(dbresults[0].NumResults)
		if dbresults[0].Approximate {
			// 128 buckets corresponds to a standard error of 10%.
			// http://algo.inria.fr/flajolet/Publications/FlFuGaMe07.pdf
			numResults = approximateNumber(numResults, 0.1)
			approximate = true
		}
	}

	pgs := newPagination(pageParams, len(results), numResults)
	pgs.Approximate = approximate
	return &SearchPage{
		Results:    results,
		Pagination: pgs,
	}, nil
}

// approximateNumber returns an approximation of the estimate, calibrated by
// the statistical estimate of standard error.
// i.e., a number that isn't misleading when we say '1-10 of approximately N
// results', but that is still close to our estimate.
func approximateNumber(estimate int, sigma float64) int {
	expectedErr := sigma * float64(estimate)
	// Compute the unit by rounding the error the logarithmically closest power
	// of 10, so that 300->100, but 400->1000.
	unit := math.Pow(10, math.Round(math.Log10(expectedErr)))
	// Now round the estimate to the nearest unit.
	return int(unit * math.Round(float64(estimate)/unit))
}

// handleSearch applies database data to the search template. Handles endpoint
// /search?q=<query>. If <query> is an exact match for a package path, the user
// will be redirected to the details page.
func (s *Server) handleSearch(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	query := searchQuery(r)
	if query == "" {
		http.Redirect(w, r, "/", http.StatusFound)
		return
	}

	if strings.Contains(query, "/") {
		pkg, err := s.ds.GetPackage(ctx, path.Clean(query), internal.UnknownModulePath, internal.LatestVersion)
		if err == nil {
			http.Redirect(w, r, fmt.Sprintf("/%s", pkg.Path), http.StatusFound)
			return
		} else if !xerrors.Is(err, derrors.NotFound) {
			log.Errorf("error getting package for %s: %v", path.Clean(query), err)
		}
	}

	searchMethod := r.FormValue("method")
	page, err := fetchSearchPage(ctx, s.ds, query, searchMethod, newPaginationParams(r, defaultSearchLimit))
	if err != nil {
		log.Errorf("fetchSearchDetails(ctx, db, %q): %v", query, err)
		s.serveErrorPage(w, r, http.StatusInternalServerError, nil)
		return
	}
	page.basePage = newBasePage(r, query)
	s.servePage(w, "search.tmpl", page)
}

// searchQuery extracts a search query from the request.
func searchQuery(r *http.Request) string {
	return strings.TrimSpace(r.FormValue("q"))
}
<|endoftext|>


<commit_before>package config

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"syscall"
)

type Broker struct {
	Name               string
	ServiceGenericName string
	IP                 string
	Port               int
	CertFile           string
	KeyFile            string
	AuthExchange       string
	AuthAllExchange    string
	WebProtocol        string
}

type Config struct {
	Aws struct {
		Key    string
		Secret string
	}
	BuildNumber int
	Environment string
	Regions     struct {
		Vagrant string
		SJ      string
		AWS     string
		Premium string
	}
	ProjectRoot     string
	UserSitesDomain string
	ContainerSubnet string
	VmPool          string
	Version         string
	Client          struct {
		StaticFilesBaseUrl string
		RuntimeOptions     RuntimeOptions
	}
	Mongo string
	Mq    struct {
		Host     string
		Port     int
		Login    string
		Password string
		Vhost    string
		LogLevel string
	}
	Broker            Broker
	PremiumBroker     Broker
	BrokerKite        Broker
	PremiumBrokerKite Broker
	ProxyKite         struct {
		Domain   string
		CertFile string
		KeyFile  string
	}
	Kontrold struct {
		Vhost    string
		Overview struct {
			ApiPort    int
			ApiHost    string
			Port       int
			KodingHost string
			SocialHost string
		}
		Api struct {
			Port int
			URL  string
		}
		Proxy struct {
			Port    int
			PortSSL int
			FTPIP   string
		}
	}
	Slack struct {
		Token   string
		Channel string
	}
	LogLevel map[string]string
	Redis    struct {
		Url string
	}
	SubscriptionEndpoint string
	Gowebserver          struct {
		Port int
	}
	Rerouting struct {
		Port int
	}
	SocialApi struct {
		ProxyUrl     string
		CustomDomain struct {
			Public string
			Local  string
		}
	}
	Vmwatcher struct {
		Port           string
		AwsKey         string
		AwsSecret      string
		KloudSecretKey string
		KloudAddr      string
	}
	Segment        string
	GatherIngestor struct {
		Port int
	}
	Mailgun struct {
		Domain     string
		PrivateKey string
		PublicKey  string
	}
}

type RuntimeOptions struct {
	Kites struct {
		DisableWebSocketByDefault bool `json:"disableWebSocketByDefault"`
		Kontrol                   struct {
			Username string `json:"username"`
		} `json:"kontrol"`
		Os struct {
			Version string `json:"version"`
		} `json:"os"`
		Terminal struct {
			Version string `json:"version"`
		} `json:"terminal"`
		Klient struct {
			Version string `json:"version"`
		} `json:"klient"`
		Kloud struct {
			Version string `json:"version"`
		} `json:"kloud"`
	} `json:"kites"`
	Algolia struct {
		AppId       string `json:"appId"`
		ApiKey      string `json:"apiKey"`
		IndexSuffix string `json:"indexSuffix"`
	} `json:"algolia"`
	SuppressLogs    bool   `json:"suppressLogs"`
	AuthExchange    string `json:"authExchange"`
	Environment     string `json:"environment"`
	Version         string `json:"version"`
	ResourceName    string `json:"resourceName"`
	UserSitesDomain string `json:"userSitesDomain"`
	LogResourceName string `json:"logResourceName"`
	SocialApiUri    string `json:"socialApiUri"`
	ApiUri          string `json:"apiUri"`
	MainUri         string `json:"mainUri"`
	SourceMapsUri   string `json:"sourceMapsUri"`
	Broker          struct {
		Uri string `json:"uri"`
	} `json:"broker"`
	AppsUri            string `json:"appsUri"`
	UploadsUri         string `json:"uploadsUri"`
	UploadsUriForGroup string `json:"uploadsUriForGroup"`
	FileFetchTimeout   int    `json:"fileFetchTimeout"`
	UserIdleMs         int    `json:"userIdleMs"`
	Embedly            struct {
		ApiKey string `json:"apiKey"`
	} `json:"embedly"`
	Github struct {
		ClientId string `json:"clientId"`
	} `json:"github"`
	Newkontrol struct {
		Url string `json:"url"`
	} `json:"newkontrol"`
	SessionCookie struct {
		MaxAge int  `json:"maxAge"`
		Secure bool `json:"secure"`
	} `json:"sessionCookie"`
	Stripe struct {
		Token string `json:"token"`
	} `json:"stripe"`
	ExternalProfiles struct {
		Google struct {
			Nicename string `json:"nicename"`
		} `json:"google"`
		Linkedin struct {
			Nicename string `json:"nicename"`
		} `json:"linkedin"`
		Twitter struct {
			Nicename string `json:"nicename"`
		} `json:"twitter"`
		Facebook struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"facebook"`
		Github struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"github"`
	} `json:"externalProfiles"`
	EntryPoint struct {
		Slug string `json:"slug"`
		Type string `json:"type"`
	} `json:"entryPoint"`
	Roles       []string      `json:"roles"`
	Permissions []interface{} `json:"permissions"`
	SiftScience string        `json:"siftScience"`
	Paypal      struct {
		FormUrl string `json:"formUrl"`
	} `json:"paypal"`
	Pubnub struct {
		SubscribeKey string `json:"subscribekey"`
		Enabled      bool   `json:"enabled"`
		SSL          bool   `json:"ssl"`
	} `json:"pubnub"`
	Collaboration struct {
		Timeout int `json:"timeout"`
	} `json:"collaboration"`
	PaymentBlockDuration float64 `json:"paymentBlockDuration"`
	DisabledFeatures     struct {
		Moderation bool `json:"moderation"`
		Teams      bool `json:"teams"`
		BotChannel bool `json:"botchannel"`
	} `json:"disabledFeatures"`
	ContentRotatorUrl string `json:"contentRotatorUrl"`
	Integration       struct {
		Url string `json:"url"`
	} `json:"integration"`
	WebhookMiddleware struct {
		Url string `json:"url"`
	} `json:"webhookMiddleware"`
	Google struct {
		ApiKey string `json:"apiKey"`
	} `json:"google"`
	Recaptcha struct {
		Key     string `json:"key"`
		Enabled bool   `json:"enabled"`
	} `json:"recaptcha"`
	Domains struct {
		Base string `json:"base"`
		Mail string `json:"mail"`
		Main string `json:"main"`
		Port string `json:"port"`
	} `json:"domains"`
}

// TODO: THIS IS ADDED SO ALL GO PACKAGES CLEANLY EXIT EVEN WHEN
// RUN WITH RERUN

func init() {

	go func() {
		signals := make(chan os.Signal, 1)
		signal.Notify(signals)
		for {
			signal := <-signals
			switch signal {
			case syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGSTOP:
				os.Exit(0)
			}
		}
	}()
}

func MustConfig(profile string) *Config {
	conf, err := readConfig("", profile)
	if err != nil {
		panic(err)
	}

	return conf
}

// MustEnv is like Env, but panics if the Config cannot be read successfully.
func MustEnv() *Config {
	conf, err := Env()
	if err != nil {
		panic(err)
	}

	return conf
}

// Env reads from the KONFIG_JSON environment variable and intitializes the
// Config struct
func Env() (*Config, error) {
	return readConfig("", "")
}

func readConfig(configDir, profile string) (*Config, error) {
	jsonData := os.Getenv("KONFIG_JSON")
	if jsonData == "" {
		return nil, errors.New("KONFIG_JSON is not set")
	}

	conf := new(Config)
	err := json.Unmarshal([]byte(jsonData), &conf)
	if err != nil {
		return nil, fmt.Errorf("Configuration error, make sure KONFIG_JSON is set: %s\nConfiguration source output:\n%s\n",
			err.Error(), string(jsonData))
	}

	return conf, nil
}
<commit_msg>go/config: remove unused configuration fields<commit_after>package config

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"syscall"
)

type Broker struct {
	Name               string
	ServiceGenericName string
	IP                 string
	Port               int
	CertFile           string
	KeyFile            string
	AuthExchange       string
	AuthAllExchange    string
	WebProtocol        string
}

type Config struct {
	BuildNumber     int
	Environment     string
	ProjectRoot     string
	UserSitesDomain string
	Version         string
	Client          struct {
		StaticFilesBaseUrl string
		RuntimeOptions     RuntimeOptions
	}
	Mongo string
	Mq    struct {
		Host     string
		Port     int
		Login    string
		Password string
		Vhost    string
		LogLevel string
	}
	Broker            Broker
	PremiumBroker     Broker
	BrokerKite        Broker
	PremiumBrokerKite Broker
	Slack             struct {
		Token   string
		Channel string
	}
	LogLevel map[string]string
	Redis    struct {
		Url string
	}
	SubscriptionEndpoint string
	Gowebserver          struct {
		Port int
	}
	Rerouting struct {
		Port int
	}
	SocialApi struct {
		ProxyUrl     string
		CustomDomain struct {
			Public string
			Local  string
		}
	}
	Vmwatcher struct {
		Port           string
		AwsKey         string
		AwsSecret      string
		KloudSecretKey string
		KloudAddr      string
	}
	Segment        string
	GatherIngestor struct {
		Port int
	}
	Mailgun struct {
		Domain     string
		PrivateKey string
		PublicKey  string
	}
}

type RuntimeOptions struct {
	Kites struct {
		DisableWebSocketByDefault bool `json:"disableWebSocketByDefault"`
		Kontrol                   struct {
			Username string `json:"username"`
		} `json:"kontrol"`
		Os struct {
			Version string `json:"version"`
		} `json:"os"`
		Terminal struct {
			Version string `json:"version"`
		} `json:"terminal"`
		Klient struct {
			Version string `json:"version"`
		} `json:"klient"`
		Kloud struct {
			Version string `json:"version"`
		} `json:"kloud"`
	} `json:"kites"`
	Algolia struct {
		AppId       string `json:"appId"`
		ApiKey      string `json:"apiKey"`
		IndexSuffix string `json:"indexSuffix"`
	} `json:"algolia"`
	SuppressLogs    bool   `json:"suppressLogs"`
	AuthExchange    string `json:"authExchange"`
	Environment     string `json:"environment"`
	Version         string `json:"version"`
	ResourceName    string `json:"resourceName"`
	UserSitesDomain string `json:"userSitesDomain"`
	LogResourceName string `json:"logResourceName"`
	SocialApiUri    string `json:"socialApiUri"`
	ApiUri          string `json:"apiUri"`
	MainUri         string `json:"mainUri"`
	SourceMapsUri   string `json:"sourceMapsUri"`
	Broker          struct {
		Uri string `json:"uri"`
	} `json:"broker"`
	AppsUri            string `json:"appsUri"`
	UploadsUri         string `json:"uploadsUri"`
	UploadsUriForGroup string `json:"uploadsUriForGroup"`
	FileFetchTimeout   int    `json:"fileFetchTimeout"`
	UserIdleMs         int    `json:"userIdleMs"`
	Embedly            struct {
		ApiKey string `json:"apiKey"`
	} `json:"embedly"`
	Github struct {
		ClientId string `json:"clientId"`
	} `json:"github"`
	Newkontrol struct {
		Url string `json:"url"`
	} `json:"newkontrol"`
	SessionCookie struct {
		MaxAge int  `json:"maxAge"`
		Secure bool `json:"secure"`
	} `json:"sessionCookie"`
	Stripe struct {
		Token string `json:"token"`
	} `json:"stripe"`
	ExternalProfiles struct {
		Google struct {
			Nicename string `json:"nicename"`
		} `json:"google"`
		Linkedin struct {
			Nicename string `json:"nicename"`
		} `json:"linkedin"`
		Twitter struct {
			Nicename string `json:"nicename"`
		} `json:"twitter"`
		Facebook struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"facebook"`
		Github struct {
			Nicename    string `json:"nicename"`
			UrlLocation string `json:"urlLocation"`
		} `json:"github"`
	} `json:"externalProfiles"`
	EntryPoint struct {
		Slug string `json:"slug"`
		Type string `json:"type"`
	} `json:"entryPoint"`
	Roles       []string      `json:"roles"`
	Permissions []interface{} `json:"permissions"`
	SiftScience string        `json:"siftScience"`
	Paypal      struct {
		FormUrl string `json:"formUrl"`
	} `json:"paypal"`
	Pubnub struct {
		SubscribeKey string `json:"subscribekey"`
		Enabled      bool   `json:"enabled"`
		SSL          bool   `json:"ssl"`
	} `json:"pubnub"`
	Collaboration struct {
		Timeout int `json:"timeout"`
	} `json:"collaboration"`
	PaymentBlockDuration float64 `json:"paymentBlockDuration"`
	DisabledFeatures     struct {
		Moderation bool `json:"moderation"`
		Teams      bool `json:"teams"`
		BotChannel bool `json:"botchannel"`
	} `json:"disabledFeatures"`
	ContentRotatorUrl string `json:"contentRotatorUrl"`
	Integration       struct {
		Url string `json:"url"`
	} `json:"integration"`
	WebhookMiddleware struct {
		Url string `json:"url"`
	} `json:"webhookMiddleware"`
	Google struct {
		ApiKey string `json:"apiKey"`
	} `json:"google"`
	Recaptcha struct {
		Key     string `json:"key"`
		Enabled bool   `json:"enabled"`
	} `json:"recaptcha"`
	Domains struct {
		Base string `json:"base"`
		Mail string `json:"mail"`
		Main string `json:"main"`
		Port string `json:"port"`
	} `json:"domains"`
}

// TODO: THIS IS ADDED SO ALL GO PACKAGES CLEANLY EXIT EVEN WHEN
// RUN WITH RERUN

func init() {

	go func() {
		signals := make(chan os.Signal, 1)
		signal.Notify(signals)
		for {
			signal := <-signals
			switch signal {
			case syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGSTOP:
				os.Exit(0)
			}
		}
	}()
}

func MustConfig(profile string) *Config {
	conf, err := readConfig("", profile)
	if err != nil {
		panic(err)
	}

	return conf
}

// MustEnv is like Env, but panics if the Config cannot be read successfully.
func MustEnv() *Config {
	conf, err := Env()
	if err != nil {
		panic(err)
	}

	return conf
}

// Env reads from the KONFIG_JSON environment variable and intitializes the
// Config struct
func Env() (*Config, error) {
	return readConfig("", "")
}

func readConfig(configDir, profile string) (*Config, error) {
	jsonData := os.Getenv("KONFIG_JSON")
	if jsonData == "" {
		return nil, errors.New("KONFIG_JSON is not set")
	}

	conf := new(Config)
	err := json.Unmarshal([]byte(jsonData), &conf)
	if err != nil {
		return nil, fmt.Errorf("Configuration error, make sure KONFIG_JSON is set: %s\nConfiguration source output:\n%s\n",
			err.Error(), string(jsonData))
	}

	return conf, nil
}
<|endoftext|>


<commit_before>package dynamodb

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"log"
	"math"
	"net/http"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	godynamodb "github.com/aws/aws-sdk-go/service/dynamodb"
	godynamodbiface "github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
	"github.com/pkg/errors"

	"github.com/yuuki/diamondb/pkg/config"
	"github.com/yuuki/diamondb/pkg/model"
	"github.com/yuuki/diamondb/pkg/storage/util"
	"github.com/yuuki/diamondb/pkg/timeparser"
)

//go:generate mockgen -source ../../../vendor/github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface/interface.go -destination dynamodb_mock.go -package dynamodb

// ReadWriter defines the interface for DynamoDB reader and writer.
type ReadWriter interface {
	Ping() error
	Client() godynamodbiface.DynamoDBAPI
	CreateTable(*CreateTableParam) error
	Fetch(string, time.Time, time.Time) (model.SeriesMap, error)
	batchGet(q *query) (model.SeriesMap, error)
	Put(string, string, string, int64, map[int64]float64) error
}

// DynamoDB provides a dynamodb client.
type DynamoDB struct {
	svc godynamodbiface.DynamoDBAPI
}

type timeSlot struct {
	itemEpoch int64
	step      int
}

type query struct {
	names []string
	start time.Time
	end   time.Time
	slot  *timeSlot
	// context
}

const (
	pingTimeout = time.Duration(5) * time.Second

	oneYear time.Duration = time.Duration(24*360) * time.Hour
	oneWeek time.Duration = time.Duration(24*7) * time.Hour
	oneDay  time.Duration = time.Duration(24*1) * time.Hour
)

var (
	dynamodbBatchLimit = 100

	oneYearSeconds = int(oneYear.Seconds())
	oneWeekSeconds = int(oneWeek.Seconds())
	oneDaySeconds  = int(oneDay.Seconds())
)

var _ ReadWriter = &DynamoDB{}

// New creates a new DynamoDB.
func New() (*DynamoDB, error) {
	awsConf := aws.NewConfig().WithRegion(config.Config.DynamoDBRegion)
	if config.Config.DynamoDBEndpoint != "" {
		// For dynamodb-local configuration
		awsConf.WithEndpoint(config.Config.DynamoDBEndpoint)
		awsConf.WithCredentials(credentials.NewStaticCredentials("dummy", "dummy", "dummy"))
	}
	awsConf.WithHTTPClient(&http.Client{
		Timeout:   10 * time.Second,
		Transport: http.DefaultTransport,
	})
	sess, err := session.NewSession(awsConf)
	if err != nil {
		return nil, errors.Wrapf(err,
			"failed to create session for dynamodb (%s,%s)",
			config.Config.DynamoDBRegion,
			config.Config.DynamoDBEndpoint,
		)
	}
	return &DynamoDB{
		svc: godynamodb.New(sess),
	}, nil
}

// Client returns the DynamoDB client.
func (d *DynamoDB) Client() godynamodbiface.DynamoDBAPI {
	return d.svc
}

// Ping pings DynamoDB endpoint.
func (d *DynamoDB) Ping() error {
	ctx, cancel := context.WithTimeout(context.TODO(), pingTimeout)
	defer cancel()
	var opt request.Option = func(r *request.Request) {}
	_, err := d.svc.DescribeTableWithContext(ctx, &godynamodb.DescribeTableInput{
		TableName: aws.String(config.Config.DynamoDBTableName),
	}, opt)
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok && aerr.Code() == request.CanceledErrorCode {
			// If the SDK can determine the request or retry delay was canceled
			// by a context the CanceledErrorCode error code will be returned.
			return errors.Wrap(err, "failed to ping dynamodb due to timeout")
		}
		return errors.Wrapf(err, "failed to ping dynamodb")
	}
	return nil
}

// CreateTableParam is parameter set of CreateTable.
type CreateTableParam struct {
	Name string
	RCU  int64 // ReadCapacityUnits
	WCU  int64 // WriteCapacityUnits
}

// CreateTable creates a dynamodb table to store time series data.
// Skip creating table if the table already exists.
func (d *DynamoDB) CreateTable(param *CreateTableParam) error {
	_, err := d.svc.CreateTable(&godynamodb.CreateTableInput{
		TableName: aws.String(param.Name),
		AttributeDefinitions: []*godynamodb.AttributeDefinition{
			{
				AttributeName: aws.String("Name"),
				AttributeType: aws.String(godynamodb.ScalarAttributeTypeS),
			},
			{
				AttributeName: aws.String("Timestamp"),
				AttributeType: aws.String(godynamodb.ScalarAttributeTypeS),
			},
		},
		KeySchema: []*godynamodb.KeySchemaElement{
			{
				AttributeName: aws.String("Name"),
				KeyType:       aws.String("HASH"),
			},
			{
				AttributeName: aws.String("Timestamp"),
				KeyType:       aws.String("RANGE"),
			},
		},
		ProvisionedThroughput: &godynamodb.ProvisionedThroughput{
			ReadCapacityUnits:  aws.Int64(param.RCU),
			WriteCapacityUnits: aws.Int64(param.WCU),
		},
		// TODO StreamSpecification to export to s3
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceInUseException" {
				// Skip if the table already exists
				log.Printf("Skip creating DynamoDB table because %s already exists\n", param.Name)
				return nil
			}
		}
		return errors.Wrapf(err, "failed to create dynamodb table (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}

	log.Printf("Creating DynamoDB table (name:%s, rcu:%d, wcu:%d) ...\n",
		param.Name, param.RCU, param.WCU)

	err = d.svc.WaitUntilTableExists(&godynamodb.DescribeTableInput{
		TableName: aws.String(param.Name),
	})
	if err != nil {
		return errors.Wrapf(err, "failed to wait until table exists (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}

	if config.Config.DynamoDBTTL {
		_, err = d.svc.UpdateTimeToLive(&godynamodb.UpdateTimeToLiveInput{
			TableName: aws.String(param.Name),
			TimeToLiveSpecification: &godynamodb.TimeToLiveSpecification{
				AttributeName: aws.String("TTL"),
				Enabled:       aws.Bool(true),
			},
		})
		if err != nil {
			return errors.Wrapf(err, "failed to set TTL to (%s,%d,%d)",
				param.Name, param.RCU, param.WCU)
		}
	}
	return nil
}

// Fetch fetches datapoints by name from start until end.
func (d *DynamoDB) Fetch(name string, start, end time.Time) (model.SeriesMap, error) {
	slots := selectTimeSlots(start, end)
	nameGroups := util.GroupNames(util.SplitName(name), dynamodbBatchLimit)
	numQueries := len(slots) * len(nameGroups)

	type result struct {
		value model.SeriesMap
		err   error
	}
	c := make(chan *result, numQueries)
	for _, slot := range slots {
		for _, names := range nameGroups {
			q := &query{
				names: names,
				start: start,
				end:   end,
				slot:  slot,
			}
			go func(q *query) {
				sm, err := d.batchGet(q)
				c <- &result{value: sm, err: err}
			}(q)
		}
	}
	sm := make(model.SeriesMap, len(nameGroups))
	for i := 0; i < numQueries; i++ {
		ret := <-c
		if ret.err != nil {
			return nil, ret.err
		}
		sm.MergePointsToMap(ret.value)
	}
	return sm, nil
}

func batchGetResultToMap(resp *godynamodb.BatchGetItemOutput, q *query) model.SeriesMap {
	sm := make(model.SeriesMap, len(resp.Responses))
	for _, xs := range resp.Responses {
		for _, x := range xs {
			name := (*x["Name"].S)
			points := make(model.DataPoints, 0, len(x["Values"].BS))
			for _, y := range x["Values"].BS {
				t := int64(binary.BigEndian.Uint64(y[0:8]))
				v := math.Float64frombits(binary.BigEndian.Uint64(y[8:]))
				// Trim datapoints out of [start, end]
				if t < q.start.Unix() || q.end.Unix() < t {
					continue
				}
				points = append(points, model.NewDataPoint(t, v))
			}
			sm[name] = model.NewSeriesPoint(name, points, q.slot.step)
		}
	}
	return sm
}

func (d *DynamoDB) batchGet(q *query) (model.SeriesMap, error) {
	var keys []map[string]*godynamodb.AttributeValue
	for _, name := range q.names {
		keys = append(keys, map[string]*godynamodb.AttributeValue{
			"Name":      {S: aws.String(name)},
			"Timestamp": {S: aws.String(fmt.Sprintf("%d:%d", q.slot.itemEpoch, q.slot.step))},
		})
	}
	items := make(map[string]*godynamodb.KeysAndAttributes)
	items[config.Config.DynamoDBTableName] = &godynamodb.KeysAndAttributes{Keys: keys}
	params := &godynamodb.BatchGetItemInput{
		RequestItems:           items,
		ReturnConsumedCapacity: aws.String("NONE"),
	}
	resp, err := d.svc.BatchGetItem(params)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				// Don't handle ResourceNotFoundException as error
				// bacause diamondb web return length 0 series as 200.
				return model.SeriesMap{}, nil
			}
		}
		return nil, errors.Wrapf(err,
			"failed to call dynamodb API batchGetItem (%s,%d,%d)",
			config.Config.DynamoDBTableName, q.slot.itemEpoch, q.slot.step,
		)
	}
	return batchGetResultToMap(resp, q), nil
}

// Put writes the datapoints into DynamoDB. It creates item
// if item doesn't exist and updates item if it exists.
func (d *DynamoDB) Put(name, slot, history string, itemEpoch int64, tv map[int64]float64) error {
	stepDuration, err := timeparser.ParseTimeOffset(slot)
	if err != nil {
		return err
	}
	historyDuration, err := timeparser.ParseTimeOffset(history)
	if err != nil {
		return err
	}
	ttl := itemEpoch + int64(historyDuration.Seconds())

	vals := make([][]byte, 0, len(tv))
	for timestamp, value := range tv {
		buf := new(bytes.Buffer)
		binary.Write(buf, binary.BigEndian, timestamp)
		binary.Write(buf, binary.BigEndian, math.Float64bits(value))
		vals = append(vals, buf.Bytes())
	}

	params := &godynamodb.UpdateItemInput{
		TableName: aws.String(config.Config.DynamoDBTableName),
		Key: map[string]*godynamodb.AttributeValue{
			"Name":      {S: aws.String(name)},
			"Timestamp": {S: aws.String(fmt.Sprintf("%d:%d", itemEpoch, int64(stepDuration.Seconds())))},
		},
		UpdateExpression: aws.String(`
			SET #ttl = :new_ttl
			ADD #values_set :new_values
		`),
		ExpressionAttributeNames: map[string]*string{
			"#ttl":        aws.String("TTL"),
			"#values_set": aws.String("Values"),
		},
		ExpressionAttributeValues: map[string]*godynamodb.AttributeValue{
			":new_ttl":    {N: aws.String(fmt.Sprintf("%d", ttl))},
			":new_values": {BS: vals},
		},
		ReturnValues: aws.String("NONE"),
	}
	if _, err := d.svc.UpdateItem(params); err != nil {
		return errors.Wrapf(err, "failed to call dynamodb API putItem (%s,%s,%d)",
			config.Config.DynamoDBTableName, name, itemEpoch)
	}
	return nil
}

func selectTimeSlots(startTime, endTime time.Time) []*timeSlot {
	var (
		step          int
		itemEpochStep int
	)
	diff := endTime.Sub(startTime)
	switch {
	case oneYear <= diff:
		itemEpochStep = oneYearSeconds
		step = 60 * 60 * 24
	case oneWeek <= diff:
		itemEpochStep = oneWeekSeconds
		step = 60 * 60
	case oneDay <= diff:
		itemEpochStep = oneDaySeconds
		step = 5 * 60
	default:
		itemEpochStep = 60 * 60
		step = 60
	}

	var slots []*timeSlot
	startItemEpoch := startTime.Unix() - startTime.Unix()%int64(itemEpochStep)
	endItemEpoch := endTime.Unix()
	for epoch := startItemEpoch; epoch < endItemEpoch; epoch += int64(itemEpochStep) {
		slots = append(slots, &timeSlot{itemEpoch: epoch, step: step})
	}

	return slots
}
<commit_msg>Add updateItem timeout with context<commit_after>package dynamodb

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"log"
	"math"
	"net/http"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	godynamodb "github.com/aws/aws-sdk-go/service/dynamodb"
	godynamodbiface "github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
	"github.com/pkg/errors"

	"github.com/yuuki/diamondb/pkg/config"
	"github.com/yuuki/diamondb/pkg/model"
	"github.com/yuuki/diamondb/pkg/storage/util"
	"github.com/yuuki/diamondb/pkg/timeparser"
)

//go:generate mockgen -source ../../../vendor/github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface/interface.go -destination dynamodb_mock.go -package dynamodb

// ReadWriter defines the interface for DynamoDB reader and writer.
type ReadWriter interface {
	Ping() error
	Client() godynamodbiface.DynamoDBAPI
	CreateTable(*CreateTableParam) error
	Fetch(string, time.Time, time.Time) (model.SeriesMap, error)
	batchGet(q *query) (model.SeriesMap, error)
	Put(string, string, string, int64, map[int64]float64) error
}

// DynamoDB provides a dynamodb client.
type DynamoDB struct {
	svc godynamodbiface.DynamoDBAPI
}

type timeSlot struct {
	itemEpoch int64
	step      int
}

type query struct {
	names []string
	start time.Time
	end   time.Time
	slot  *timeSlot
	// context
}

const (
	pingTimeout   = time.Duration(5) * time.Second
	updateTimeout = time.Duration(5) * time.Second

	oneYear time.Duration = time.Duration(24*360) * time.Hour
	oneWeek time.Duration = time.Duration(24*7) * time.Hour
	oneDay  time.Duration = time.Duration(24*1) * time.Hour
)

var (
	dynamodbBatchLimit = 100

	oneYearSeconds = int(oneYear.Seconds())
	oneWeekSeconds = int(oneWeek.Seconds())
	oneDaySeconds  = int(oneDay.Seconds())
)

var _ ReadWriter = &DynamoDB{}

// New creates a new DynamoDB.
func New() (*DynamoDB, error) {
	awsConf := aws.NewConfig().WithRegion(config.Config.DynamoDBRegion)
	if config.Config.DynamoDBEndpoint != "" {
		// For dynamodb-local configuration
		awsConf.WithEndpoint(config.Config.DynamoDBEndpoint)
		awsConf.WithCredentials(credentials.NewStaticCredentials("dummy", "dummy", "dummy"))
	}
	awsConf.WithHTTPClient(&http.Client{
		Timeout:   10 * time.Second,
		Transport: http.DefaultTransport,
	})
	sess, err := session.NewSession(awsConf)
	if err != nil {
		return nil, errors.Wrapf(err,
			"failed to create session for dynamodb (%s,%s)",
			config.Config.DynamoDBRegion,
			config.Config.DynamoDBEndpoint,
		)
	}
	return &DynamoDB{
		svc: godynamodb.New(sess),
	}, nil
}

// Client returns the DynamoDB client.
func (d *DynamoDB) Client() godynamodbiface.DynamoDBAPI {
	return d.svc
}

// Ping pings DynamoDB endpoint.
func (d *DynamoDB) Ping() error {
	ctx, cancel := context.WithTimeout(context.TODO(), pingTimeout)
	defer cancel()
	var opt request.Option = func(r *request.Request) {}
	_, err := d.svc.DescribeTableWithContext(ctx, &godynamodb.DescribeTableInput{
		TableName: aws.String(config.Config.DynamoDBTableName),
	}, opt)
	if err != nil {
		if aerr, ok := err.(awserr.Error); ok && aerr.Code() == request.CanceledErrorCode {
			// If the SDK can determine the request or retry delay was canceled
			// by a context the CanceledErrorCode error code will be returned.
			return errors.Wrap(err, "failed to ping dynamodb due to timeout")
		}
		return errors.Wrapf(err, "failed to ping dynamodb")
	}
	return nil
}

// CreateTableParam is parameter set of CreateTable.
type CreateTableParam struct {
	Name string
	RCU  int64 // ReadCapacityUnits
	WCU  int64 // WriteCapacityUnits
}

// CreateTable creates a dynamodb table to store time series data.
// Skip creating table if the table already exists.
func (d *DynamoDB) CreateTable(param *CreateTableParam) error {
	_, err := d.svc.CreateTable(&godynamodb.CreateTableInput{
		TableName: aws.String(param.Name),
		AttributeDefinitions: []*godynamodb.AttributeDefinition{
			{
				AttributeName: aws.String("Name"),
				AttributeType: aws.String(godynamodb.ScalarAttributeTypeS),
			},
			{
				AttributeName: aws.String("Timestamp"),
				AttributeType: aws.String(godynamodb.ScalarAttributeTypeS),
			},
		},
		KeySchema: []*godynamodb.KeySchemaElement{
			{
				AttributeName: aws.String("Name"),
				KeyType:       aws.String("HASH"),
			},
			{
				AttributeName: aws.String("Timestamp"),
				KeyType:       aws.String("RANGE"),
			},
		},
		ProvisionedThroughput: &godynamodb.ProvisionedThroughput{
			ReadCapacityUnits:  aws.Int64(param.RCU),
			WriteCapacityUnits: aws.Int64(param.WCU),
		},
		// TODO StreamSpecification to export to s3
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceInUseException" {
				// Skip if the table already exists
				log.Printf("Skip creating DynamoDB table because %s already exists\n", param.Name)
				return nil
			}
		}
		return errors.Wrapf(err, "failed to create dynamodb table (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}

	log.Printf("Creating DynamoDB table (name:%s, rcu:%d, wcu:%d) ...\n",
		param.Name, param.RCU, param.WCU)

	err = d.svc.WaitUntilTableExists(&godynamodb.DescribeTableInput{
		TableName: aws.String(param.Name),
	})
	if err != nil {
		return errors.Wrapf(err, "failed to wait until table exists (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}

	if config.Config.DynamoDBTTL {
		_, err = d.svc.UpdateTimeToLive(&godynamodb.UpdateTimeToLiveInput{
			TableName: aws.String(param.Name),
			TimeToLiveSpecification: &godynamodb.TimeToLiveSpecification{
				AttributeName: aws.String("TTL"),
				Enabled:       aws.Bool(true),
			},
		})
		if err != nil {
			return errors.Wrapf(err, "failed to set TTL to (%s,%d,%d)",
				param.Name, param.RCU, param.WCU)
		}
	}
	return nil
}

// Fetch fetches datapoints by name from start until end.
func (d *DynamoDB) Fetch(name string, start, end time.Time) (model.SeriesMap, error) {
	slots := selectTimeSlots(start, end)
	nameGroups := util.GroupNames(util.SplitName(name), dynamodbBatchLimit)
	numQueries := len(slots) * len(nameGroups)

	type result struct {
		value model.SeriesMap
		err   error
	}
	c := make(chan *result, numQueries)
	for _, slot := range slots {
		for _, names := range nameGroups {
			q := &query{
				names: names,
				start: start,
				end:   end,
				slot:  slot,
			}
			go func(q *query) {
				sm, err := d.batchGet(q)
				c <- &result{value: sm, err: err}
			}(q)
		}
	}
	sm := make(model.SeriesMap, len(nameGroups))
	for i := 0; i < numQueries; i++ {
		ret := <-c
		if ret.err != nil {
			return nil, ret.err
		}
		sm.MergePointsToMap(ret.value)
	}
	return sm, nil
}

func batchGetResultToMap(resp *godynamodb.BatchGetItemOutput, q *query) model.SeriesMap {
	sm := make(model.SeriesMap, len(resp.Responses))
	for _, xs := range resp.Responses {
		for _, x := range xs {
			name := (*x["Name"].S)
			points := make(model.DataPoints, 0, len(x["Values"].BS))
			for _, y := range x["Values"].BS {
				t := int64(binary.BigEndian.Uint64(y[0:8]))
				v := math.Float64frombits(binary.BigEndian.Uint64(y[8:]))
				// Trim datapoints out of [start, end]
				if t < q.start.Unix() || q.end.Unix() < t {
					continue
				}
				points = append(points, model.NewDataPoint(t, v))
			}
			sm[name] = model.NewSeriesPoint(name, points, q.slot.step)
		}
	}
	return sm
}

func (d *DynamoDB) batchGet(q *query) (model.SeriesMap, error) {
	var keys []map[string]*godynamodb.AttributeValue
	for _, name := range q.names {
		keys = append(keys, map[string]*godynamodb.AttributeValue{
			"Name":      {S: aws.String(name)},
			"Timestamp": {S: aws.String(fmt.Sprintf("%d:%d", q.slot.itemEpoch, q.slot.step))},
		})
	}
	items := make(map[string]*godynamodb.KeysAndAttributes)
	items[config.Config.DynamoDBTableName] = &godynamodb.KeysAndAttributes{Keys: keys}
	params := &godynamodb.BatchGetItemInput{
		RequestItems:           items,
		ReturnConsumedCapacity: aws.String("NONE"),
	}
	resp, err := d.svc.BatchGetItem(params)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				// Don't handle ResourceNotFoundException as error
				// bacause diamondb web return length 0 series as 200.
				return model.SeriesMap{}, nil
			}
		}
		return nil, errors.Wrapf(err,
			"failed to call dynamodb API batchGetItem (%s,%d,%d)",
			config.Config.DynamoDBTableName, q.slot.itemEpoch, q.slot.step,
		)
	}
	return batchGetResultToMap(resp, q), nil
}

// Put writes the datapoints into DynamoDB. It creates item
// if item doesn't exist and updates item if it exists.
func (d *DynamoDB) Put(name, slot, history string, itemEpoch int64, tv map[int64]float64) error {
	stepDuration, err := timeparser.ParseTimeOffset(slot)
	if err != nil {
		return err
	}
	historyDuration, err := timeparser.ParseTimeOffset(history)
	if err != nil {
		return err
	}
	ttl := itemEpoch + int64(historyDuration.Seconds())

	vals := make([][]byte, 0, len(tv))
	for timestamp, value := range tv {
		buf := new(bytes.Buffer)
		binary.Write(buf, binary.BigEndian, timestamp)
		binary.Write(buf, binary.BigEndian, math.Float64bits(value))
		vals = append(vals, buf.Bytes())
	}

	params := &godynamodb.UpdateItemInput{
		TableName: aws.String(config.Config.DynamoDBTableName),
		Key: map[string]*godynamodb.AttributeValue{
			"Name":      {S: aws.String(name)},
			"Timestamp": {S: aws.String(fmt.Sprintf("%d:%d", itemEpoch, int64(stepDuration.Seconds())))},
		},
		UpdateExpression: aws.String(`
			SET #ttl = :new_ttl
			ADD #values_set :new_values
		`),
		ExpressionAttributeNames: map[string]*string{
			"#ttl":        aws.String("TTL"),
			"#values_set": aws.String("Values"),
		},
		ExpressionAttributeValues: map[string]*godynamodb.AttributeValue{
			":new_ttl":    {N: aws.String(fmt.Sprintf("%d", ttl))},
			":new_values": {BS: vals},
		},
		ReturnValues: aws.String("NONE"),
	}

	ctx, cancel := context.WithTimeout(context.TODO(), updateTimeout)
	defer cancel()
	var opt request.Option = func(r *request.Request) {}
	if _, err := d.svc.UpdateItemWithContext(ctx, params, opt); err != nil {
		if aerr, ok := err.(awserr.Error); ok && aerr.Code() == request.CanceledErrorCode {
			// If the SDK can determine the request or retry delay was canceled
			// by a context the CanceledErrorCode error code will be returned.
			return errors.Wrap(err, "failed to updateItem dynamodb due to timeout")
		}
		return errors.Wrapf(err, "failed to call dynamodb API putItem (%s,%s,%d)",
			config.Config.DynamoDBTableName, name, itemEpoch)
	}
	return nil
}

func selectTimeSlots(startTime, endTime time.Time) []*timeSlot {
	var (
		step          int
		itemEpochStep int
	)
	diff := endTime.Sub(startTime)
	switch {
	case oneYear <= diff:
		itemEpochStep = oneYearSeconds
		step = 60 * 60 * 24
	case oneWeek <= diff:
		itemEpochStep = oneWeekSeconds
		step = 60 * 60
	case oneDay <= diff:
		itemEpochStep = oneDaySeconds
		step = 5 * 60
	default:
		itemEpochStep = 60 * 60
		step = 60
	}

	var slots []*timeSlot
	startItemEpoch := startTime.Unix() - startTime.Unix()%int64(itemEpochStep)
	endItemEpoch := endTime.Unix()
	for epoch := startItemEpoch; epoch < endItemEpoch; epoch += int64(itemEpochStep) {
		slots = append(slots, &timeSlot{itemEpoch: epoch, step: step})
	}

	return slots
}
<|endoftext|>


<commit_before>/*
Copyright 2017 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 azure_dd

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	libstrings "strings"

	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-10-01/compute"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/sets"
	api "k8s.io/kubernetes/pkg/apis/core"
	"k8s.io/kubernetes/pkg/cloudprovider/providers/azure"
	"k8s.io/kubernetes/pkg/util/mount"
	"k8s.io/kubernetes/pkg/util/strings"
	"k8s.io/kubernetes/pkg/volume"
)

const (
	defaultStorageAccountType       = compute.StandardLRS
	defaultAzureDiskKind            = v1.AzureManagedDisk
	defaultAzureDataDiskCachingMode = v1.AzureDataDiskCachingNone
)

type dataDisk struct {
	volume.MetricsProvider
	volumeName string
	diskName   string
	podUID     types.UID
	plugin     *azureDataDiskPlugin
}

var (
	supportedCachingModes = sets.NewString(
		string(api.AzureDataDiskCachingNone),
		string(api.AzureDataDiskCachingReadOnly),
		string(api.AzureDataDiskCachingReadWrite))

	supportedDiskKinds = sets.NewString(
		string(api.AzureSharedBlobDisk),
		string(api.AzureDedicatedBlobDisk),
		string(api.AzureManagedDisk))

	lunPathRE = regexp.MustCompile(`/dev/disk/azure/scsi(?:.*)/lun(.+)`)
)

func getPath(uid types.UID, volName string, host volume.VolumeHost) string {
	return host.GetPodVolumeDir(uid, strings.EscapeQualifiedNameForDisk(azureDataDiskPluginName), volName)
}

// creates a unique path for disks (even if they share the same *.vhd name)
func makeGlobalPDPath(host volume.VolumeHost, diskUri string, isManaged bool) (string, error) {
	diskUri = libstrings.ToLower(diskUri) // always lower uri because users may enter it in caps.
	uniqueDiskNameTemplate := "%s%s"
	hashedDiskUri := azure.MakeCRC32(diskUri)
	prefix := "b"
	if isManaged {
		prefix = "m"
	}
	// "{m for managed b for blob}{hashed diskUri or DiskId depending on disk kind }"
	diskName := fmt.Sprintf(uniqueDiskNameTemplate, prefix, hashedDiskUri)
	pdPath := filepath.Join(host.GetPluginDir(azureDataDiskPluginName), mount.MountsInGlobalPDPath, diskName)

	return pdPath, nil
}

func makeDataDisk(volumeName string, podUID types.UID, diskName string, host volume.VolumeHost, plugin *azureDataDiskPlugin) *dataDisk {
	var metricProvider volume.MetricsProvider
	if podUID != "" {
		metricProvider = volume.NewMetricsStatFS(getPath(podUID, volumeName, host))
	}

	return &dataDisk{
		MetricsProvider: metricProvider,
		volumeName:      volumeName,
		diskName:        diskName,
		podUID:          podUID,
		plugin:          plugin,
	}
}

func getVolumeSource(spec *volume.Spec) (volumeSource *v1.AzureDiskVolumeSource, readOnly bool, err error) {
	if spec.Volume != nil && spec.Volume.AzureDisk != nil {
		return spec.Volume.AzureDisk, spec.Volume.AzureDisk.ReadOnly != nil && *spec.Volume.AzureDisk.ReadOnly, nil
	}

	if spec.PersistentVolume != nil && spec.PersistentVolume.Spec.AzureDisk != nil {
		return spec.PersistentVolume.Spec.AzureDisk, spec.ReadOnly, nil
	}

	return nil, false, fmt.Errorf("azureDisk - Spec does not reference an Azure disk volume type")
}

func normalizeKind(kind string) (v1.AzureDataDiskKind, error) {
	if kind == "" {
		return defaultAzureDiskKind, nil
	}

	if !supportedDiskKinds.Has(kind) {
		return "", fmt.Errorf("azureDisk - %s is not supported disk kind. Supported values are %s", kind, supportedDiskKinds.List())
	}

	return v1.AzureDataDiskKind(kind), nil
}

func normalizeStorageAccountType(storageAccountType string) (compute.DiskStorageAccountTypes, error) {
	if storageAccountType == "" {
		return defaultStorageAccountType, nil
	}

	sku := compute.DiskStorageAccountTypes(storageAccountType)
	supportedSkuNames := compute.PossibleDiskStorageAccountTypesValues()
	for _, s := range supportedSkuNames {
		if sku == s {
			return sku, nil
		}
	}

	return "", fmt.Errorf("azureDisk - %s is not supported sku/storageaccounttype. Supported values are %s", storageAccountType, supportedSkuNames)
}

func normalizeCachingMode(cachingMode v1.AzureDataDiskCachingMode) (v1.AzureDataDiskCachingMode, error) {
	if cachingMode == "" {
		return defaultAzureDataDiskCachingMode, nil
	}

	if !supportedCachingModes.Has(string(cachingMode)) {
		return "", fmt.Errorf("azureDisk - %s is not supported cachingmode. Supported values are %s", cachingMode, supportedCachingModes.List())
	}

	return cachingMode, nil
}

type ioHandler interface {
	ReadDir(dirname string) ([]os.FileInfo, error)
	WriteFile(filename string, data []byte, perm os.FileMode) error
	Readlink(name string) (string, error)
	ReadFile(filename string) ([]byte, error)
}

//TODO: check if priming the iscsi interface is actually needed

type osIOHandler struct{}

func (handler *osIOHandler) ReadDir(dirname string) ([]os.FileInfo, error) {
	return ioutil.ReadDir(dirname)
}

func (handler *osIOHandler) WriteFile(filename string, data []byte, perm os.FileMode) error {
	return ioutil.WriteFile(filename, data, perm)
}

func (handler *osIOHandler) Readlink(name string) (string, error) {
	return os.Readlink(name)
}

func (handler *osIOHandler) ReadFile(filename string) ([]byte, error) {
	return ioutil.ReadFile(filename)
}

func getDiskController(host volume.VolumeHost) (DiskController, error) {
	cloudProvider := host.GetCloudProvider()
	az, ok := cloudProvider.(*azure.Cloud)

	if !ok || az == nil {
		return nil, fmt.Errorf("AzureDisk -  failed to get Azure Cloud Provider. GetCloudProvider returned %v instead", cloudProvider)
	}
	return az, nil
}

func getCloud(host volume.VolumeHost) (*azure.Cloud, error) {
	cloudProvider := host.GetCloudProvider()
	az, ok := cloudProvider.(*azure.Cloud)

	if !ok || az == nil {
		return nil, fmt.Errorf("AzureDisk -  failed to get Azure Cloud Provider. GetCloudProvider returned %v instead", cloudProvider)
	}
	return az, nil
}

func strFirstLetterToUpper(str string) string {
	if len(str) < 2 {
		return str
	}
	return libstrings.ToUpper(string(str[0])) + str[1:]
}

// getDiskLUN : deviceInfo could be a LUN number or a device path, e.g. /dev/disk/azure/scsi1/lun2
func getDiskLUN(deviceInfo string) (int32, error) {
	var diskLUN string
	if len(deviceInfo) <= 2 {
		diskLUN = deviceInfo
	} else {
		// extract the LUN num from a device path
		matches := lunPathRE.FindStringSubmatch(deviceInfo)
		if len(matches) == 2 {
			diskLUN = matches[1]
		} else {
			return -1, fmt.Errorf("cannot parse deviceInfo: %s", deviceInfo)
		}
	}

	lun, err := strconv.Atoi(diskLUN)
	if err != nil {
		return -1, err
	}
	return int32(lun), nil
}
<commit_msg>change azure disk host cache to ReadOnly by default<commit_after>/*
Copyright 2017 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 azure_dd

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	libstrings "strings"

	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-10-01/compute"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/sets"
	api "k8s.io/kubernetes/pkg/apis/core"
	"k8s.io/kubernetes/pkg/cloudprovider/providers/azure"
	"k8s.io/kubernetes/pkg/util/mount"
	"k8s.io/kubernetes/pkg/util/strings"
	"k8s.io/kubernetes/pkg/volume"
)

const (
	defaultStorageAccountType       = compute.StandardLRS
	defaultAzureDiskKind            = v1.AzureManagedDisk
	defaultAzureDataDiskCachingMode = v1.AzureDataDiskCachingReadOnly
)

type dataDisk struct {
	volume.MetricsProvider
	volumeName string
	diskName   string
	podUID     types.UID
	plugin     *azureDataDiskPlugin
}

var (
	supportedCachingModes = sets.NewString(
		string(api.AzureDataDiskCachingNone),
		string(api.AzureDataDiskCachingReadOnly),
		string(api.AzureDataDiskCachingReadWrite))

	supportedDiskKinds = sets.NewString(
		string(api.AzureSharedBlobDisk),
		string(api.AzureDedicatedBlobDisk),
		string(api.AzureManagedDisk))

	lunPathRE = regexp.MustCompile(`/dev/disk/azure/scsi(?:.*)/lun(.+)`)
)

func getPath(uid types.UID, volName string, host volume.VolumeHost) string {
	return host.GetPodVolumeDir(uid, strings.EscapeQualifiedNameForDisk(azureDataDiskPluginName), volName)
}

// creates a unique path for disks (even if they share the same *.vhd name)
func makeGlobalPDPath(host volume.VolumeHost, diskUri string, isManaged bool) (string, error) {
	diskUri = libstrings.ToLower(diskUri) // always lower uri because users may enter it in caps.
	uniqueDiskNameTemplate := "%s%s"
	hashedDiskUri := azure.MakeCRC32(diskUri)
	prefix := "b"
	if isManaged {
		prefix = "m"
	}
	// "{m for managed b for blob}{hashed diskUri or DiskId depending on disk kind }"
	diskName := fmt.Sprintf(uniqueDiskNameTemplate, prefix, hashedDiskUri)
	pdPath := filepath.Join(host.GetPluginDir(azureDataDiskPluginName), mount.MountsInGlobalPDPath, diskName)

	return pdPath, nil
}

func makeDataDisk(volumeName string, podUID types.UID, diskName string, host volume.VolumeHost, plugin *azureDataDiskPlugin) *dataDisk {
	var metricProvider volume.MetricsProvider
	if podUID != "" {
		metricProvider = volume.NewMetricsStatFS(getPath(podUID, volumeName, host))
	}

	return &dataDisk{
		MetricsProvider: metricProvider,
		volumeName:      volumeName,
		diskName:        diskName,
		podUID:          podUID,
		plugin:          plugin,
	}
}

func getVolumeSource(spec *volume.Spec) (volumeSource *v1.AzureDiskVolumeSource, readOnly bool, err error) {
	if spec.Volume != nil && spec.Volume.AzureDisk != nil {
		return spec.Volume.AzureDisk, spec.Volume.AzureDisk.ReadOnly != nil && *spec.Volume.AzureDisk.ReadOnly, nil
	}

	if spec.PersistentVolume != nil && spec.PersistentVolume.Spec.AzureDisk != nil {
		return spec.PersistentVolume.Spec.AzureDisk, spec.ReadOnly, nil
	}

	return nil, false, fmt.Errorf("azureDisk - Spec does not reference an Azure disk volume type")
}

func normalizeKind(kind string) (v1.AzureDataDiskKind, error) {
	if kind == "" {
		return defaultAzureDiskKind, nil
	}

	if !supportedDiskKinds.Has(kind) {
		return "", fmt.Errorf("azureDisk - %s is not supported disk kind. Supported values are %s", kind, supportedDiskKinds.List())
	}

	return v1.AzureDataDiskKind(kind), nil
}

func normalizeStorageAccountType(storageAccountType string) (compute.DiskStorageAccountTypes, error) {
	if storageAccountType == "" {
		return defaultStorageAccountType, nil
	}

	sku := compute.DiskStorageAccountTypes(storageAccountType)
	supportedSkuNames := compute.PossibleDiskStorageAccountTypesValues()
	for _, s := range supportedSkuNames {
		if sku == s {
			return sku, nil
		}
	}

	return "", fmt.Errorf("azureDisk - %s is not supported sku/storageaccounttype. Supported values are %s", storageAccountType, supportedSkuNames)
}

func normalizeCachingMode(cachingMode v1.AzureDataDiskCachingMode) (v1.AzureDataDiskCachingMode, error) {
	if cachingMode == "" {
		return defaultAzureDataDiskCachingMode, nil
	}

	if !supportedCachingModes.Has(string(cachingMode)) {
		return "", fmt.Errorf("azureDisk - %s is not supported cachingmode. Supported values are %s", cachingMode, supportedCachingModes.List())
	}

	return cachingMode, nil
}

type ioHandler interface {
	ReadDir(dirname string) ([]os.FileInfo, error)
	WriteFile(filename string, data []byte, perm os.FileMode) error
	Readlink(name string) (string, error)
	ReadFile(filename string) ([]byte, error)
}

//TODO: check if priming the iscsi interface is actually needed

type osIOHandler struct{}

func (handler *osIOHandler) ReadDir(dirname string) ([]os.FileInfo, error) {
	return ioutil.ReadDir(dirname)
}

func (handler *osIOHandler) WriteFile(filename string, data []byte, perm os.FileMode) error {
	return ioutil.WriteFile(filename, data, perm)
}

func (handler *osIOHandler) Readlink(name string) (string, error) {
	return os.Readlink(name)
}

func (handler *osIOHandler) ReadFile(filename string) ([]byte, error) {
	return ioutil.ReadFile(filename)
}

func getDiskController(host volume.VolumeHost) (DiskController, error) {
	cloudProvider := host.GetCloudProvider()
	az, ok := cloudProvider.(*azure.Cloud)

	if !ok || az == nil {
		return nil, fmt.Errorf("AzureDisk -  failed to get Azure Cloud Provider. GetCloudProvider returned %v instead", cloudProvider)
	}
	return az, nil
}

func getCloud(host volume.VolumeHost) (*azure.Cloud, error) {
	cloudProvider := host.GetCloudProvider()
	az, ok := cloudProvider.(*azure.Cloud)

	if !ok || az == nil {
		return nil, fmt.Errorf("AzureDisk -  failed to get Azure Cloud Provider. GetCloudProvider returned %v instead", cloudProvider)
	}
	return az, nil
}

func strFirstLetterToUpper(str string) string {
	if len(str) < 2 {
		return str
	}
	return libstrings.ToUpper(string(str[0])) + str[1:]
}

// getDiskLUN : deviceInfo could be a LUN number or a device path, e.g. /dev/disk/azure/scsi1/lun2
func getDiskLUN(deviceInfo string) (int32, error) {
	var diskLUN string
	if len(deviceInfo) <= 2 {
		diskLUN = deviceInfo
	} else {
		// extract the LUN num from a device path
		matches := lunPathRE.FindStringSubmatch(deviceInfo)
		if len(matches) == 2 {
			diskLUN = matches[1]
		} else {
			return -1, fmt.Errorf("cannot parse deviceInfo: %s", deviceInfo)
		}
	}

	lun, err := strconv.Atoi(diskLUN)
	if err != nil {
		return -1, err
	}
	return int32(lun), nil
}
<|endoftext|>