<commit_before>package osdn

import (
	"fmt"
	"net"
	"strings"
	"time"

	log "github.com/golang/glog"

	"github.com/openshift/openshift-sdn/pkg/netutils"
	"github.com/openshift/openshift-sdn/plugins/osdn/api"

	utildbus "k8s.io/kubernetes/pkg/util/dbus"
	kerrors "k8s.io/kubernetes/pkg/util/errors"
	kexec "k8s.io/kubernetes/pkg/util/exec"
	"k8s.io/kubernetes/pkg/util/iptables"
)

type PluginHooks interface {
	PluginStartMaster(clusterNetworkCIDR string, clusterBitsPerSubnet uint, serviceNetworkCIDR string) error
	PluginStartNode(mtu uint) error
}

type OvsController struct {
	pluginHooks     PluginHooks
	Registry        *Registry
	localIP         string
	localSubnet     *api.Subnet
	hostName        string
	subnetAllocator *netutils.SubnetAllocator
	sig             chan struct{}
	podNetworkReady chan struct{}
	flowController  FlowController
	VNIDMap         map[string]uint
	netIDManager    *netutils.NetIDAllocator
	adminNamespaces []string
	services        map[string]api.Service
}

type FlowController interface {
	Setup(localSubnetCIDR, clusterNetworkCIDR, serviceNetworkCIDR string, mtu uint) error

	AddOFRules(nodeIP, nodeSubnetCIDR, localIP string) error
	DelOFRules(nodeIP, localIP string) error

	AddServiceOFRules(netID uint, IP string, protocol api.ServiceProtocol, port uint) error
	DelServiceOFRules(netID uint, IP string, protocol api.ServiceProtocol, port uint) error

	UpdatePod(namespace, podName, containerID string, netID uint) error
}

// Called by plug factory functions to initialize the generic plugin instance
func (oc *OvsController) BaseInit(registry *Registry, flowController FlowController, pluginHooks PluginHooks, hostname string, selfIP string) error {
	if hostname == "" {
		output, err := kexec.New().Command("uname", "-n").CombinedOutput()
		if err != nil {
			return err
		}
		hostname = strings.TrimSpace(string(output))
	}

	if selfIP == "" {
		var err error
		selfIP, err = netutils.GetNodeIP(hostname)
		if err != nil {
			return err
		}
	}
	log.Infof("Self IP: %s.", selfIP)

	oc.pluginHooks = pluginHooks
	oc.Registry = registry
	oc.flowController = flowController
	oc.localIP = selfIP
	oc.hostName = hostname
	oc.VNIDMap = make(map[string]uint)
	oc.sig = make(chan struct{})
	oc.podNetworkReady = make(chan struct{})
	oc.adminNamespaces = make([]string, 0)
	oc.services = make(map[string]api.Service)

	return nil
}

func (oc *OvsController) validateClusterNetwork(networkCIDR string, subnetsInUse []string, hostIPNets []*net.IPNet) error {
	clusterIP, clusterIPNet, err := net.ParseCIDR(networkCIDR)
	if err != nil {
		return fmt.Errorf("Failed to parse network address: %s", networkCIDR)
	}

	errList := []error{}
	for _, ipNet := range hostIPNets {
		if ipNet.Contains(clusterIP) {
			errList = append(errList, fmt.Errorf("Error: Cluster IP: %s conflicts with host network: %s", clusterIP.String(), ipNet.String()))
		}
		if clusterIPNet.Contains(ipNet.IP) {
			errList = append(errList, fmt.Errorf("Error: Host network with IP: %s conflicts with cluster network: %s", ipNet.IP.String(), networkCIDR))
		}
	}

	for _, netStr := range subnetsInUse {
		subnetIP, _, err := net.ParseCIDR(netStr)
		if err != nil {
			errList = append(errList, fmt.Errorf("Failed to parse network address: %s", netStr))
			continue
		}
		if !clusterIPNet.Contains(subnetIP) {
			errList = append(errList, fmt.Errorf("Error: Existing node subnet: %s is not part of cluster network: %s", netStr, networkCIDR))
		}
	}
	return kerrors.NewAggregate(errList)
}

func (oc *OvsController) validateServiceNetwork(networkCIDR string, hostIPNets []*net.IPNet) error {
	serviceIP, serviceIPNet, err := net.ParseCIDR(networkCIDR)
	if err != nil {
		return fmt.Errorf("Failed to parse network address: %s", networkCIDR)
	}

	errList := []error{}
	for _, ipNet := range hostIPNets {
		if ipNet.Contains(serviceIP) {
			errList = append(errList, fmt.Errorf("Error: Service IP: %s conflicts with host network: %s", ipNet.String(), networkCIDR))
		}
		if serviceIPNet.Contains(ipNet.IP) {
			errList = append(errList, fmt.Errorf("Error: Host network with IP: %s conflicts with service network: %s", ipNet.IP.String(), networkCIDR))
		}
	}

	services, _, err := oc.Registry.GetServices()
	if err != nil {
		return err
	}
	for _, svc := range services {
		if !serviceIPNet.Contains(net.ParseIP(svc.IP)) {
			errList = append(errList, fmt.Errorf("Error: Existing service with IP: %s is not part of service network: %s", svc.IP, networkCIDR))
		}
	}
	return kerrors.NewAggregate(errList)
}

func (oc *OvsController) validateNetworkConfig(clusterNetworkCIDR, serviceNetworkCIDR string, subnetsInUse []string) error {
	// TODO: Instead of hardcoding 'tun0' and 'lbr0', get it from common place.
	// This will ensure both the kube/multitenant scripts and master validations use the same name.
	hostIPNets, err := netutils.GetHostIPNetworks([]string{"tun0", "lbr0"})
	if err != nil {
		return err
	}

	errList := []error{}
	if err := oc.validateClusterNetwork(clusterNetworkCIDR, subnetsInUse, hostIPNets); err != nil {
		errList = append(errList, err)
	}
	if err := oc.validateServiceNetwork(serviceNetworkCIDR, hostIPNets); err != nil {
		errList = append(errList, err)
	}
	return kerrors.NewAggregate(errList)
}

func (oc *OvsController) StartMaster(clusterNetworkCIDR string, clusterBitsPerSubnet uint, serviceNetworkCIDR string) error {
	// Any mismatch in cluster/service network is handled by WriteNetworkConfig
	// For any new cluster/service network, ensure existing node subnets belong
	// to the given cluster network and service IPs belong to the given service network
	if _, err := oc.Registry.GetClusterNetworkCIDR(); err != nil {
		subrange := make([]string, 0)
		subnets, _, err := oc.Registry.GetSubnets()
		if err != nil {
			log.Errorf("Error in initializing/fetching subnets: %v", err)
			return err
		}
		for _, sub := range subnets {
			subrange = append(subrange, sub.SubnetCIDR)
		}

		err = oc.validateNetworkConfig(clusterNetworkCIDR, serviceNetworkCIDR, subrange)
		if err != nil {
			return err
		}
	}

	if err := oc.Registry.WriteNetworkConfig(clusterNetworkCIDR, clusterBitsPerSubnet, serviceNetworkCIDR); err != nil {
		return err
	}

	if err := oc.pluginHooks.PluginStartMaster(clusterNetworkCIDR, clusterBitsPerSubnet, serviceNetworkCIDR); err != nil {
		return fmt.Errorf("Failed to start plugin: %v", err)
	}

	return nil
}

func (oc *OvsController) StartNode(mtu uint) error {
	// Assume we are working with IPv4
	clusterNetworkCIDR, err := oc.Registry.GetClusterNetworkCIDR()
	if err != nil {
		log.Errorf("Failed to obtain ClusterNetwork: %v", err)
		return err
	}

	ipt := iptables.New(kexec.New(), utildbus.New(), iptables.ProtocolIpv4)
	if err := SetupIptables(ipt, clusterNetworkCIDR); err != nil {
		return fmt.Errorf("Failed to set up iptables: %v", err)
	}

	ipt.AddReloadFunc(func() {
		err := SetupIptables(ipt, clusterNetworkCIDR)
		if err != nil {
			log.Errorf("Error reloading iptables: %v\n", err)
		}
	})

	if err := oc.pluginHooks.PluginStartNode(mtu); err != nil {
		return fmt.Errorf("Failed to start plugin: %v", err)
	}

	oc.markPodNetworkReady()

	return nil
}

func (oc *OvsController) markPodNetworkReady() {
	close(oc.podNetworkReady)
}

func (oc *OvsController) WaitForPodNetworkReady() error {
	timeout := 2 * time.Minute
	select {
	case <-oc.podNetworkReady:
		// Wait for StartNode() to finish SDN setup
	case <-time.After(timeout):
		return fmt.Errorf("SDN pod network is not ready(timeout: %v)", timeout)
	}
	return nil
}

func (oc *OvsController) Stop() {
	close(oc.sig)
}

// Wait for ready signal from Watch interface for the given resource
// Closes the ready channel as we don't need it anymore after this point
func waitForWatchReadiness(ready chan bool, resourceName string) {
	timeout := time.Minute
	select {
	case <-ready:
		close(ready)
	case <-time.After(timeout):
		log.Fatalf("Watch for resource %s is not ready(timeout: %v)", resourceName, timeout)
	}
	return
}

type watchWatcher func(oc *OvsController, ready chan<- bool, start <-chan string)
type watchGetter func(registry *Registry) (interface{}, string, error)

// watchAndGetResource will fetch current items in etcd and watch for any new
// changes for the given resource.
// Supported resources: nodes, subnets, namespaces, services, netnamespaces, and pods.
//
// To avoid any potential race conditions during this process, these steps are followed:
// 1. Initiator(master/node): Watch for a resource as an async op, lets say WatchProcess
// 2. WatchProcess: When ready for watching, send ready signal to initiator
// 3. Initiator: Wait for watch resource to be ready
//    This is needed as step-1 is an asynchronous operation
// 4. WatchProcess: Collect new changes in the queue but wait for initiator
//    to indicate which version to start from
// 5. Initiator: Get existing items with their latest version for the resource
// 6. Initiator: Send version from step-5 to WatchProcess
// 7. WatchProcess: Ignore any items with version <= start version got from initiator on step-6
// 8. WatchProcess: Handle new changes
func (oc *OvsController) watchAndGetResource(resourceName string, watcher watchWatcher, getter watchGetter) (interface{}, error) {
	ready := make(chan bool)
	start := make(chan string)

	go watcher(oc, ready, start)
	waitForWatchReadiness(ready, strings.ToLower(resourceName))
	getOutput, version, err := getter(oc.Registry)
	if err != nil {
		return nil, err
	}

	start <- version

	return getOutput, nil
}

type FirewallRule struct {
	table string
	chain string
	args  []string
}

func SetupIptables(ipt iptables.Interface, clusterNetworkCIDR string) error {
	rules := []FirewallRule{
		{"nat", "POSTROUTING", []string{"-s", clusterNetworkCIDR, "!", "-d", clusterNetworkCIDR, "-j", "MASQUERADE"}},
		{"filter", "INPUT", []string{"-p", "udp", "-m", "multiport", "--dports", "4789", "-m", "comment", "--comment", "001 vxlan incoming", "-j", "ACCEPT"}},
		{"filter", "INPUT", []string{"-i", "tun0", "-m", "comment", "--comment", "traffic from docker for internet", "-j", "ACCEPT"}},
		{"filter", "FORWARD", []string{"-d", clusterNetworkCIDR, "-j", "ACCEPT"}},
		{"filter", "FORWARD", []string{"-s", clusterNetworkCIDR, "-j", "ACCEPT"}},
	}

	for _, rule := range rules {
		_, err := ipt.EnsureRule(iptables.Prepend, iptables.Table(rule.table), iptables.Chain(rule.chain), rule.args...)
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Log message when pod is waiting for sdn initialization<commit_after>package osdn

import (
	"fmt"
	"net"
	"strings"
	"time"

	log "github.com/golang/glog"

	"github.com/openshift/openshift-sdn/pkg/netutils"
	"github.com/openshift/openshift-sdn/plugins/osdn/api"

	utildbus "k8s.io/kubernetes/pkg/util/dbus"
	kerrors "k8s.io/kubernetes/pkg/util/errors"
	kexec "k8s.io/kubernetes/pkg/util/exec"
	"k8s.io/kubernetes/pkg/util/iptables"
)

type PluginHooks interface {
	PluginStartMaster(clusterNetworkCIDR string, clusterBitsPerSubnet uint, serviceNetworkCIDR string) error
	PluginStartNode(mtu uint) error
}

type OvsController struct {
	pluginHooks     PluginHooks
	Registry        *Registry
	localIP         string
	localSubnet     *api.Subnet
	hostName        string
	subnetAllocator *netutils.SubnetAllocator
	sig             chan struct{}
	podNetworkReady chan struct{}
	flowController  FlowController
	VNIDMap         map[string]uint
	netIDManager    *netutils.NetIDAllocator
	adminNamespaces []string
	services        map[string]api.Service
}

type FlowController interface {
	Setup(localSubnetCIDR, clusterNetworkCIDR, serviceNetworkCIDR string, mtu uint) error

	AddOFRules(nodeIP, nodeSubnetCIDR, localIP string) error
	DelOFRules(nodeIP, localIP string) error

	AddServiceOFRules(netID uint, IP string, protocol api.ServiceProtocol, port uint) error
	DelServiceOFRules(netID uint, IP string, protocol api.ServiceProtocol, port uint) error

	UpdatePod(namespace, podName, containerID string, netID uint) error
}

// Called by plug factory functions to initialize the generic plugin instance
func (oc *OvsController) BaseInit(registry *Registry, flowController FlowController, pluginHooks PluginHooks, hostname string, selfIP string) error {
	if hostname == "" {
		output, err := kexec.New().Command("uname", "-n").CombinedOutput()
		if err != nil {
			return err
		}
		hostname = strings.TrimSpace(string(output))
	}

	if selfIP == "" {
		var err error
		selfIP, err = netutils.GetNodeIP(hostname)
		if err != nil {
			return err
		}
	}
	log.Infof("Self IP: %s.", selfIP)

	oc.pluginHooks = pluginHooks
	oc.Registry = registry
	oc.flowController = flowController
	oc.localIP = selfIP
	oc.hostName = hostname
	oc.VNIDMap = make(map[string]uint)
	oc.sig = make(chan struct{})
	oc.podNetworkReady = make(chan struct{})
	oc.adminNamespaces = make([]string, 0)
	oc.services = make(map[string]api.Service)

	return nil
}

func (oc *OvsController) validateClusterNetwork(networkCIDR string, subnetsInUse []string, hostIPNets []*net.IPNet) error {
	clusterIP, clusterIPNet, err := net.ParseCIDR(networkCIDR)
	if err != nil {
		return fmt.Errorf("Failed to parse network address: %s", networkCIDR)
	}

	errList := []error{}
	for _, ipNet := range hostIPNets {
		if ipNet.Contains(clusterIP) {
			errList = append(errList, fmt.Errorf("Error: Cluster IP: %s conflicts with host network: %s", clusterIP.String(), ipNet.String()))
		}
		if clusterIPNet.Contains(ipNet.IP) {
			errList = append(errList, fmt.Errorf("Error: Host network with IP: %s conflicts with cluster network: %s", ipNet.IP.String(), networkCIDR))
		}
	}

	for _, netStr := range subnetsInUse {
		subnetIP, _, err := net.ParseCIDR(netStr)
		if err != nil {
			errList = append(errList, fmt.Errorf("Failed to parse network address: %s", netStr))
			continue
		}
		if !clusterIPNet.Contains(subnetIP) {
			errList = append(errList, fmt.Errorf("Error: Existing node subnet: %s is not part of cluster network: %s", netStr, networkCIDR))
		}
	}
	return kerrors.NewAggregate(errList)
}

func (oc *OvsController) validateServiceNetwork(networkCIDR string, hostIPNets []*net.IPNet) error {
	serviceIP, serviceIPNet, err := net.ParseCIDR(networkCIDR)
	if err != nil {
		return fmt.Errorf("Failed to parse network address: %s", networkCIDR)
	}

	errList := []error{}
	for _, ipNet := range hostIPNets {
		if ipNet.Contains(serviceIP) {
			errList = append(errList, fmt.Errorf("Error: Service IP: %s conflicts with host network: %s", ipNet.String(), networkCIDR))
		}
		if serviceIPNet.Contains(ipNet.IP) {
			errList = append(errList, fmt.Errorf("Error: Host network with IP: %s conflicts with service network: %s", ipNet.IP.String(), networkCIDR))
		}
	}

	services, _, err := oc.Registry.GetServices()
	if err != nil {
		return err
	}
	for _, svc := range services {
		if !serviceIPNet.Contains(net.ParseIP(svc.IP)) {
			errList = append(errList, fmt.Errorf("Error: Existing service with IP: %s is not part of service network: %s", svc.IP, networkCIDR))
		}
	}
	return kerrors.NewAggregate(errList)
}

func (oc *OvsController) validateNetworkConfig(clusterNetworkCIDR, serviceNetworkCIDR string, subnetsInUse []string) error {
	// TODO: Instead of hardcoding 'tun0' and 'lbr0', get it from common place.
	// This will ensure both the kube/multitenant scripts and master validations use the same name.
	hostIPNets, err := netutils.GetHostIPNetworks([]string{"tun0", "lbr0"})
	if err != nil {
		return err
	}

	errList := []error{}
	if err := oc.validateClusterNetwork(clusterNetworkCIDR, subnetsInUse, hostIPNets); err != nil {
		errList = append(errList, err)
	}
	if err := oc.validateServiceNetwork(serviceNetworkCIDR, hostIPNets); err != nil {
		errList = append(errList, err)
	}
	return kerrors.NewAggregate(errList)
}

func (oc *OvsController) StartMaster(clusterNetworkCIDR string, clusterBitsPerSubnet uint, serviceNetworkCIDR string) error {
	// Any mismatch in cluster/service network is handled by WriteNetworkConfig
	// For any new cluster/service network, ensure existing node subnets belong
	// to the given cluster network and service IPs belong to the given service network
	if _, err := oc.Registry.GetClusterNetworkCIDR(); err != nil {
		subrange := make([]string, 0)
		subnets, _, err := oc.Registry.GetSubnets()
		if err != nil {
			log.Errorf("Error in initializing/fetching subnets: %v", err)
			return err
		}
		for _, sub := range subnets {
			subrange = append(subrange, sub.SubnetCIDR)
		}

		err = oc.validateNetworkConfig(clusterNetworkCIDR, serviceNetworkCIDR, subrange)
		if err != nil {
			return err
		}
	}

	if err := oc.Registry.WriteNetworkConfig(clusterNetworkCIDR, clusterBitsPerSubnet, serviceNetworkCIDR); err != nil {
		return err
	}

	if err := oc.pluginHooks.PluginStartMaster(clusterNetworkCIDR, clusterBitsPerSubnet, serviceNetworkCIDR); err != nil {
		return fmt.Errorf("Failed to start plugin: %v", err)
	}

	return nil
}

func (oc *OvsController) StartNode(mtu uint) error {
	// Assume we are working with IPv4
	clusterNetworkCIDR, err := oc.Registry.GetClusterNetworkCIDR()
	if err != nil {
		log.Errorf("Failed to obtain ClusterNetwork: %v", err)
		return err
	}

	ipt := iptables.New(kexec.New(), utildbus.New(), iptables.ProtocolIpv4)
	if err := SetupIptables(ipt, clusterNetworkCIDR); err != nil {
		return fmt.Errorf("Failed to set up iptables: %v", err)
	}

	ipt.AddReloadFunc(func() {
		err := SetupIptables(ipt, clusterNetworkCIDR)
		if err != nil {
			log.Errorf("Error reloading iptables: %v\n", err)
		}
	})

	if err := oc.pluginHooks.PluginStartNode(mtu); err != nil {
		return fmt.Errorf("Failed to start plugin: %v", err)
	}

	oc.markPodNetworkReady()

	return nil
}

func (oc *OvsController) markPodNetworkReady() {
	close(oc.podNetworkReady)
}

func (oc *OvsController) WaitForPodNetworkReady() error {
	logInterval := 10 * time.Second
	numIntervals := 12 // timeout: 2 mins

	for i := 0; i < numIntervals; i++ {
		select {
		// Wait for StartNode() to finish SDN setup
		case <-oc.podNetworkReady:
			return nil
		case <-time.After(logInterval):
			log.Infof("Waiting for SDN pod network to be ready...")
		}
	}
	return fmt.Errorf("SDN pod network is not ready(timeout: 2 mins)")
}

func (oc *OvsController) Stop() {
	close(oc.sig)
}

// Wait for ready signal from Watch interface for the given resource
// Closes the ready channel as we don't need it anymore after this point
func waitForWatchReadiness(ready chan bool, resourceName string) {
	timeout := time.Minute
	select {
	case <-ready:
		close(ready)
	case <-time.After(timeout):
		log.Fatalf("Watch for resource %s is not ready(timeout: %v)", resourceName, timeout)
	}
	return
}

type watchWatcher func(oc *OvsController, ready chan<- bool, start <-chan string)
type watchGetter func(registry *Registry) (interface{}, string, error)

// watchAndGetResource will fetch current items in etcd and watch for any new
// changes for the given resource.
// Supported resources: nodes, subnets, namespaces, services, netnamespaces, and pods.
//
// To avoid any potential race conditions during this process, these steps are followed:
// 1. Initiator(master/node): Watch for a resource as an async op, lets say WatchProcess
// 2. WatchProcess: When ready for watching, send ready signal to initiator
// 3. Initiator: Wait for watch resource to be ready
//    This is needed as step-1 is an asynchronous operation
// 4. WatchProcess: Collect new changes in the queue but wait for initiator
//    to indicate which version to start from
// 5. Initiator: Get existing items with their latest version for the resource
// 6. Initiator: Send version from step-5 to WatchProcess
// 7. WatchProcess: Ignore any items with version <= start version got from initiator on step-6
// 8. WatchProcess: Handle new changes
func (oc *OvsController) watchAndGetResource(resourceName string, watcher watchWatcher, getter watchGetter) (interface{}, error) {
	ready := make(chan bool)
	start := make(chan string)

	go watcher(oc, ready, start)
	waitForWatchReadiness(ready, strings.ToLower(resourceName))
	getOutput, version, err := getter(oc.Registry)
	if err != nil {
		return nil, err
	}

	start <- version

	return getOutput, nil
}

type FirewallRule struct {
	table string
	chain string
	args  []string
}

func SetupIptables(ipt iptables.Interface, clusterNetworkCIDR string) error {
	rules := []FirewallRule{
		{"nat", "POSTROUTING", []string{"-s", clusterNetworkCIDR, "!", "-d", clusterNetworkCIDR, "-j", "MASQUERADE"}},
		{"filter", "INPUT", []string{"-p", "udp", "-m", "multiport", "--dports", "4789", "-m", "comment", "--comment", "001 vxlan incoming", "-j", "ACCEPT"}},
		{"filter", "INPUT", []string{"-i", "tun0", "-m", "comment", "--comment", "traffic from docker for internet", "-j", "ACCEPT"}},
		{"filter", "FORWARD", []string{"-d", clusterNetworkCIDR, "-j", "ACCEPT"}},
		{"filter", "FORWARD", []string{"-s", clusterNetworkCIDR, "-j", "ACCEPT"}},
	}

	for _, rule := range rules {
		_, err := ipt.EnsureRule(iptables.Prepend, iptables.Table(rule.table), iptables.Chain(rule.chain), rule.args...)
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package index

import (
	. "github.com/balzaczyy/golucene/core/codec/spi"
	. "github.com/balzaczyy/golucene/core/index/model"
	"github.com/balzaczyy/golucene/core/store"
	"github.com/balzaczyy/golucene/core/util"
)

/* Default general purpose indexing chain, which handles indexing all types of fields */
type DefaultIndexingChain struct {
	bytesUsed  util.Counter
	docState   *docState
	docWriter  *DocumentsWriterPerThread
	fieldInfos *FieldInfosBuilder

	// Writes postings and term vectors:
	termsHash TermsHash

	storedFieldsWriter StoredFieldsWriter // lazy init
	lastStoredDocId    int

	fieldHash []*PerField

	nextFieldGen int64

	// Holds fields seen in each document
	fields []*PerField
}

func newDefaultIndexingChain(docWriter *DocumentsWriterPerThread) *DefaultIndexingChain {
	termVectorsWriter := newTermVectorsConsumer(docWriter)
	return &DefaultIndexingChain{
		docWriter:  docWriter,
		fieldInfos: docWriter.fieldInfos,
		docState:   docWriter.docState,
		bytesUsed:  docWriter._bytesUsed,
		termsHash:  newFreqProxTermsWriter(docWriter, termVectorsWriter),
		fields:     make([]*PerField, 1),
	}
}

// TODO: can we remove this lazy-init / make cleaner / do it another way...?
func (c *DefaultIndexingChain) initStoredFieldsWriter() (err error) {
	if c.storedFieldsWriter == nil {
		c.storedFieldsWriter, err = c.docWriter.codec.StoredFieldsFormat().FieldsWriter(
			c.docWriter.directory, c.docWriter.segmentInfo, store.IO_CONTEXT_DEFAULT)
	}
	return
}

func (c *DefaultIndexingChain) flush(state *SegmentWriteState) error {
	panic("not implemented yet")
}

/*
Catch up for all docs before us that had no stored fields, or hit
non-aborting errors before writing stored fields.
*/
func (c *DefaultIndexingChain) fillStoredFields(docId int) (err error) {
	for err == nil && c.lastStoredDocId < docId {
		err = c.startStoredFields()
		if err == nil {
			err = c.finishStoredFields()
		}
	}
	return
}

func (c *DefaultIndexingChain) abort() {
	// E.g. close any open files in the stored fields writer:
	if c.storedFieldsWriter != nil {
		c.storedFieldsWriter.Abort() // ignore error
	}

	// E.g. close any open files in the term vectors writer:
	c.termsHash.abort()

	for i, _ := range c.fieldHash {
		c.fieldHash[i] = nil
	}
}

/* Calls StoredFieldsWriter.startDocument, aborting the segment if it hits any error. */
func (c *DefaultIndexingChain) startStoredFields() (err error) {
	var success = false
	defer func() {
		if !success {
			c.docWriter.setAborting()
		}
	}()

	if err = c.initStoredFieldsWriter(); err != nil {
		return
	}
	if err = c.storedFieldsWriter.StartDocument(); err != nil {
		return
	}
	success = true

	c.lastStoredDocId++
	return nil
}

/* Calls StoredFieldsWriter.finishDocument(), aborting the segment if it hits any error. */
func (c *DefaultIndexingChain) finishStoredFields() error {
	panic("not implemented yet")
}

func (c *DefaultIndexingChain) processDocument() (err error) {
	// How many indexed field names we've seen (collapses multiple
	// field instances by the same name):
	fieldCount := 0

	fieldGen := c.nextFieldGen
	c.nextFieldGen++

	// NOTE: we need to passes here, in case there are multi-valued
	// fields, because we must process all instances of a given field
	// at once, since the anlayzer is free to reuse TOkenStream across
	// fields (i.e., we cannot have more than one TokenStream running
	// "at once"):

	c.termsHash.startDocument()

	if err = c.fillStoredFields(c.docState.docID); err != nil {
		return
	}
	if err = c.startStoredFields(); err != nil {
		return
	}

	if err = func() error {
		defer func() {
			if !c.docWriter.aborting {
				// Finish each indexed field name seen in the document:
				for _, field := range c.fields[:fieldCount] {
					err = mergeError(err, field.finish())
				}
				err = mergeError(err, c.finishStoredFields())
			}
		}()

		for _, field := range c.docState.doc {
			if fieldCount, err = c.processField(field, fieldGen, fieldCount); err != nil {
				return err
			}
		}
		return nil
	}(); err != nil {
		return
	}

	var success = false
	defer func() {
		if !success {
			// Must abort, on the possibility that on-disk term vectors are now corrupt:
			c.docWriter.setAborting()
		}
	}()

	if err = c.termsHash.finishDocument(); err != nil {
		return
	}
	success = true
	return nil
}

func (c *DefaultIndexingChain) processField(field IndexableField,
	fieldGen int64, fieldCount int) (int, error) {
	panic("not implemented yet")
}

type PerField struct {
	*DefaultIndexingChain // acess at least docState, termsHash.
}

func (f *PerField) finish() error {
	panic("not implemented yet")
}
<commit_msg>implement DefaultIndexingChain.finishStoredFields()<commit_after>package index

import (
	. "github.com/balzaczyy/golucene/core/codec/spi"
	. "github.com/balzaczyy/golucene/core/index/model"
	"github.com/balzaczyy/golucene/core/store"
	"github.com/balzaczyy/golucene/core/util"
)

/* Default general purpose indexing chain, which handles indexing all types of fields */
type DefaultIndexingChain struct {
	bytesUsed  util.Counter
	docState   *docState
	docWriter  *DocumentsWriterPerThread
	fieldInfos *FieldInfosBuilder

	// Writes postings and term vectors:
	termsHash TermsHash

	storedFieldsWriter StoredFieldsWriter // lazy init
	lastStoredDocId    int

	fieldHash []*PerField

	nextFieldGen int64

	// Holds fields seen in each document
	fields []*PerField
}

func newDefaultIndexingChain(docWriter *DocumentsWriterPerThread) *DefaultIndexingChain {
	termVectorsWriter := newTermVectorsConsumer(docWriter)
	return &DefaultIndexingChain{
		docWriter:  docWriter,
		fieldInfos: docWriter.fieldInfos,
		docState:   docWriter.docState,
		bytesUsed:  docWriter._bytesUsed,
		termsHash:  newFreqProxTermsWriter(docWriter, termVectorsWriter),
		fields:     make([]*PerField, 1),
	}
}

// TODO: can we remove this lazy-init / make cleaner / do it another way...?
func (c *DefaultIndexingChain) initStoredFieldsWriter() (err error) {
	if c.storedFieldsWriter == nil {
		c.storedFieldsWriter, err = c.docWriter.codec.StoredFieldsFormat().FieldsWriter(
			c.docWriter.directory, c.docWriter.segmentInfo, store.IO_CONTEXT_DEFAULT)
	}
	return
}

func (c *DefaultIndexingChain) flush(state *SegmentWriteState) error {
	panic("not implemented yet")
}

/*
Catch up for all docs before us that had no stored fields, or hit
non-aborting errors before writing stored fields.
*/
func (c *DefaultIndexingChain) fillStoredFields(docId int) (err error) {
	for err == nil && c.lastStoredDocId < docId {
		err = c.startStoredFields()
		if err == nil {
			err = c.finishStoredFields()
		}
	}
	return
}

func (c *DefaultIndexingChain) abort() {
	// E.g. close any open files in the stored fields writer:
	if c.storedFieldsWriter != nil {
		c.storedFieldsWriter.Abort() // ignore error
	}

	// E.g. close any open files in the term vectors writer:
	c.termsHash.abort()

	for i, _ := range c.fieldHash {
		c.fieldHash[i] = nil
	}
}

/* Calls StoredFieldsWriter.startDocument, aborting the segment if it hits any error. */
func (c *DefaultIndexingChain) startStoredFields() (err error) {
	var success = false
	defer func() {
		if !success {
			c.docWriter.setAborting()
		}
	}()

	if err = c.initStoredFieldsWriter(); err != nil {
		return
	}
	if err = c.storedFieldsWriter.StartDocument(); err != nil {
		return
	}
	success = true

	c.lastStoredDocId++
	return nil
}

/* Calls StoredFieldsWriter.finishDocument(), aborting the segment if it hits any error. */
func (c *DefaultIndexingChain) finishStoredFields() error {
	var success = false
	defer func() {
		if !success {
			c.docWriter.setAborting()
		}
	}()
	if err := c.storedFieldsWriter.FinishDocument(); err != nil {
		return err
	}
	success = true
	return nil
}

func (c *DefaultIndexingChain) processDocument() (err error) {
	// How many indexed field names we've seen (collapses multiple
	// field instances by the same name):
	fieldCount := 0

	fieldGen := c.nextFieldGen
	c.nextFieldGen++

	// NOTE: we need to passes here, in case there are multi-valued
	// fields, because we must process all instances of a given field
	// at once, since the anlayzer is free to reuse TOkenStream across
	// fields (i.e., we cannot have more than one TokenStream running
	// "at once"):

	c.termsHash.startDocument()

	if err = c.fillStoredFields(c.docState.docID); err != nil {
		return
	}
	if err = c.startStoredFields(); err != nil {
		return
	}

	if err = func() error {
		defer func() {
			if !c.docWriter.aborting {
				// Finish each indexed field name seen in the document:
				for _, field := range c.fields[:fieldCount] {
					err = mergeError(err, field.finish())
				}
				err = mergeError(err, c.finishStoredFields())
			}
		}()

		for _, field := range c.docState.doc {
			if fieldCount, err = c.processField(field, fieldGen, fieldCount); err != nil {
				return err
			}
		}
		return nil
	}(); err != nil {
		return
	}

	var success = false
	defer func() {
		if !success {
			// Must abort, on the possibility that on-disk term vectors are now corrupt:
			c.docWriter.setAborting()
		}
	}()

	if err = c.termsHash.finishDocument(); err != nil {
		return
	}
	success = true
	return nil
}

func (c *DefaultIndexingChain) processField(field IndexableField,
	fieldGen int64, fieldCount int) (int, error) {
	panic("not implemented yet")
}

type PerField struct {
	*DefaultIndexingChain // acess at least docState, termsHash.
}

func (f *PerField) finish() error {
	panic("not implemented yet")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
	"time"
)

// start when there is no config should report that fact right away
func TestStartWithoutConfig(t *testing.T) {
	done := start("/not/notaconfig.conf")
	message := <- done
	if message != "Could not read file at /not/notaconfig.conf" {
		t.Errorf("did not report the file that was missing")
	}
}

// start when there is a bad config should report that fact right away
func TestStartWithBadConfig(t *testing.T) {
	done := start("./test_bad.conf")
	message := <- done
	if message != "Could not decode config" {
		t.Errorf("incorrect bad config mesage\n\"%s\"", message)
	}
}

// start with a valid file should not be done right away
func TestStartWithGoodConfig(t *testing.T) {
	done := start("./test_good.conf")
	select {
	case message := <-done:
		t.Errorf("done with message\n \"%s\"", message)
	case <-time.After(time.Millisecond * 50):
		fmt.Print("Stayed up with good config\n")
	}
}

// Stabilizer.ServeHTTP should return the first good response
func TestStabilizerReturnsFirstResponse(t *testing.T) {
	// mock handler that first errors, then takes a long time then returns a
	// a good response
	reqCount := 0
	mockHandler := func(w http.ResponseWriter, r *http.Request) {
		reqCount++
		if reqCount % 3 == 1 {
			http.Error(w, "test error", 1234567890)
		}
		if reqCount % 3 == 2 {
			time.Sleep(100 * time.Millisecond)
			fmt.Fprintf(w, "slow response")
		}
		if reqCount % 3 == 0 {
			fmt.Fprintf(w, "fast response")
		}
	}

	mockUnstableBackend := httptest.NewServer(http.HandlerFunc(mockHandler))
	defer mockUnstableBackend.Close()
	u, err := url.Parse(mockUnstableBackend.URL)
	if err != nil { t.Errorf("error parsing backend url test broken") }
	testStabilizer := &Stabilizer{u, 4}
	testStableServer := httptest.NewServer(
		http.TimeoutHandler(testStabilizer, 5 * time.Second, "timeout"),
	)
	defer testStableServer.Close()

	// make many requests and make sure they are all the fast response
	for i := 0; i < 3; i++ {
		res, err := http.Get(testStableServer.URL)
		if err != nil { t.Errorf("error response from stable server") }

		message, err := ioutil.ReadAll(res.Body)
		if err != nil { t.Errorf("error reading response body from stable server") }
		res.Body.Close()

		// ensure that all responses are the fast response
		if string(message) != "fast response" {
			t.Errorf(string(message))
		}

		time.Sleep(100 * time.Millisecond)
	}
}

// TestCanPass proves that tests are running
func TestCanPass(t *testing.T) {
	if true != true {
		t.Errorf("true is not true,\ncheck your premises,\n consider clojure?")
	}
}
<commit_msg>tests are way gentler on the system now<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
	"time"
)

// start when there is no config should report that fact right away
func TestStartWithoutConfig(t *testing.T) {
	done := start("/not/notaconfig.conf")
	message := <- done
	if message != "Could not read file at /not/notaconfig.conf" {
		t.Errorf("did not report the file that was missing")
	}
}

// start when there is a bad config should report that fact right away
func TestStartWithBadConfig(t *testing.T) {
	done := start("./test_bad.conf")
	message := <- done
	if message != "Could not decode config" {
		t.Errorf("incorrect bad config mesage\n\"%s\"", message)
	}
}

// start with a valid file should not be done right away
func TestStartWithGoodConfig(t *testing.T) {
	done := start("./test_good.conf")
	select {
	case message := <-done:
		t.Errorf("done with message\n \"%s\"", message)
	case <-time.After(time.Millisecond * 50):
		fmt.Print("Stayed up with good config\n")
	}
}

// Stabilizer.ServeHTTP should return the first good response
func TestStabilizerReturnsFirstResponse(t *testing.T) {
	// mock handler that first errors, then takes a long time then returns a
	// a good response
	reqCount := 0
	mockHandler := func(w http.ResponseWriter, r *http.Request) {
		reqCount++
		if reqCount % 3 == 1 {
			http.Error(w, "test error", 1234567890)
		}
		if reqCount % 3 == 2 {
			time.Sleep(100 * time.Millisecond)
			fmt.Fprintf(w, "slow response")
		}
		if reqCount % 3 == 0 {
			fmt.Fprintf(w, "fast response")
		}
	}

	mockUnstableBackend := httptest.NewServer(http.HandlerFunc(mockHandler))
	defer mockUnstableBackend.Close()
	u, err := url.Parse(mockUnstableBackend.URL)
	if err != nil { t.Errorf("error parsing backend url test broken") }
	testStabilizer := &Stabilizer{u, 4}
	testStableServer := httptest.NewServer(
		http.TimeoutHandler(testStabilizer, 5 * time.Second, "timeout"),
	)
	defer testStableServer.Close()

	res, err := http.Get(testStableServer.URL)
	if err != nil { t.Errorf("error getting response from stable server") }

	message, err := ioutil.ReadAll(res.Body)
	if err != nil { t.Errorf("error reading response body from stable server") }
	res.Body.Close()

	// ensure that all responses are the fast response
	if string(message) != "fast response" {
		t.Errorf(string(message))
	}

}

// TestCanPass proves that tests are running
func TestCanPass(t *testing.T) {
	if true != true {
		t.Errorf("true is not true,\ncheck your premises,\n consider clojure?")
	}
}
<|endoftext|>


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

// Package proxy provides the image proxy.
package proxy

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"reflect"
	"strings"
	"time"

	"github.com/golang/glog"
	"github.com/gregjones/httpcache"
)

// Proxy serves image requests.
type Proxy struct {
	Client *http.Client // client used to fetch remote URLs
	Cache  Cache

	// Whitelist specifies a list of remote hosts that images can be proxied from.  An empty list means all hosts are allowed.
	Whitelist []string

	StorageAddr string

	MaxWidth  int
	MaxHeight int
}

// NewProxy constructs a new proxy.  The provided http Client will be used to
// fetch remote URLs.  If nil is provided, http.DefaultClient will be used.
func NewProxy(transport http.RoundTripper, cache Cache) *Proxy {
	if transport == nil {
		transport = http.DefaultTransport
	}
	if cache == nil {
		cache = NopCache
	}

	client := new(http.Client)
	client.Transport = &httpcache.Transport{
		Transport:           &TransformingTransport{transport, client},
		Cache:               cache,
		MarkCachedResponses: true,
	}

	return &Proxy{
		Client: client,
		Cache:  cache,
	}
}

// ServeHTTP handles image requests.
func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {

	req, err := NewRequest(r, p.StorageAddr)

	if err != nil {
		glog.Errorf("invalid request URL: %v", err)
		http.Error(w, fmt.Sprintf("invalid request URL: %v", err), http.StatusBadRequest)
		return
	}
	req.StorageAddr = p.StorageAddr
	u := req.URL.String()

	if p.MaxWidth > 0 && int(req.Options.Width) > p.MaxWidth {
		req.Options.Width = float64(p.MaxWidth)
	}
	if p.MaxHeight > 0 && int(req.Options.Height) > p.MaxHeight {
		req.Options.Height = float64(p.MaxHeight)
	}

	/*
	if !p.allowed(req.URL) {
		glog.Errorf("remote URL is not for an allowed host: %v", req.URL)
		http.Error(w, fmt.Sprintf("remote URL is not for an allowed host: %v", req.URL), http.StatusBadRequest)
		return
	}
	*/

	//u := req.URL.String()
	if req.Options != nil && !reflect.DeepEqual(req.Options, emptyOptions) {
		u += "#" + req.Options.String()
	}

	resp, err := p.Client.Get(u)
	if err != nil {

		http.Error(w, fmt.Printf("%v", http.StatusInternalServerError), 500)
		return

		glog.Errorf("error fetching remote image: %v", err)
		http.Error(w, fmt.Sprintf("Error fetching remote image: %v", err), http.StatusInternalServerError)
		return
	}

	if resp.StatusCode != http.StatusOK {
		http.Error(w, fmt.Sprintf("Remote URL %q returned status: %v", req.URL, resp.Status), http.StatusInternalServerError)
		return
	}

	w.Header().Add("Last-Modified", resp.Header.Get("Last-Modified"))
	w.Header().Add("Expires", resp.Header.Get("Expires"))
	w.Header().Add("Etag", resp.Header.Get("Etag"))

	if is304 := check304(w, r, resp); is304 {
		w.WriteHeader(http.StatusNotModified)
		return
	}

	w.Header().Add("Content-Length", resp.Header.Get("Content-Length"))
	defer resp.Body.Close()
	io.Copy(w, resp.Body)
}

// allowed returns whether the specified URL is on the whitelist of remote hosts.
func (p *Proxy) allowed(u *url.URL) bool {
	if len(p.Whitelist) == 0 {
		return true
	}

	for _, host := range p.Whitelist {
		if u.Host == host {
			return true
		}
		if strings.HasPrefix(host, "*.") && strings.HasSuffix(u.Host, host[2:]) {
			return true
		}
	}

	return false
}

func check304(w http.ResponseWriter, req *http.Request, resp *http.Response) bool {
	etag := resp.Header.Get("Etag")
	if etag != "" && etag == req.Header.Get("If-None-Match") {
		return true
	}

	lastModified, err := time.Parse(time.RFC1123, resp.Header.Get("Last-Modified"))
	if err != nil {
		return false
	}
	ifModSince, err := time.Parse(time.RFC1123, req.Header.Get("If-Modified-Since"))
	if err != nil {
		return false
	}
	if lastModified.Before(ifModSince) {
		return true
	}

	return false
}

// TransformingTransport is an implementation of http.RoundTripper that
// optionally transforms images using the options specified in the request URL
// fragment.
type TransformingTransport struct {
	// Transport is used to satisfy non-transform requests (those that do not include a URL fragment)
	Transport http.RoundTripper

	// Client is used to fetch images to be resized.
	Client *http.Client
}

// RoundTrip implements http.RoundTripper.
func (t *TransformingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if req.URL.Fragment == "" {
		// normal requests pass through
		glog.Infof("fetching remote URL: %v", req.URL)
		return t.Transport.RoundTrip(req)
	}

	u := *req.URL
	u.Fragment = ""
	resp, err := t.Client.Get(u.String())

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

	opt := ParseOptions(req.URL.Fragment)
	img, err := Transform(b, opt)
	if err != nil {
		img = b
	}

	// replay response with transformed image and updated content length
	buf := new(bytes.Buffer)
	fmt.Fprintf(buf, "%s %s\n", resp.Proto, resp.Status)
	resp.Header.WriteSubset(buf, map[string]bool{"Content-Length": true})
	fmt.Fprintf(buf, "Content-Length: %d\n\n", len(img))
	buf.Write(img)

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

// Package proxy provides the image proxy.
package proxy

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"reflect"
	"strings"
	"time"

	"github.com/golang/glog"
	"github.com/gregjones/httpcache"
)

// Proxy serves image requests.
type Proxy struct {
	Client *http.Client // client used to fetch remote URLs
	Cache  Cache

	// Whitelist specifies a list of remote hosts that images can be proxied from.  An empty list means all hosts are allowed.
	Whitelist []string

	StorageAddr string

	MaxWidth  int
	MaxHeight int
}

// NewProxy constructs a new proxy.  The provided http Client will be used to
// fetch remote URLs.  If nil is provided, http.DefaultClient will be used.
func NewProxy(transport http.RoundTripper, cache Cache) *Proxy {
	if transport == nil {
		transport = http.DefaultTransport
	}
	if cache == nil {
		cache = NopCache
	}

	client := new(http.Client)
	client.Transport = &httpcache.Transport{
		Transport:           &TransformingTransport{transport, client},
		Cache:               cache,
		MarkCachedResponses: true,
	}

	return &Proxy{
		Client: client,
		Cache:  cache,
	}
}

// ServeHTTP handles image requests.
func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {

	req, err := NewRequest(r, p.StorageAddr)

	if err != nil {
		glog.Errorf("invalid request URL: %v", err)
		http.Error(w, fmt.Sprintf("invalid request URL: %v", err), http.StatusBadRequest)
		return
	}
	req.StorageAddr = p.StorageAddr
	u := req.URL.String()

	if p.MaxWidth > 0 && int(req.Options.Width) > p.MaxWidth {
		req.Options.Width = float64(p.MaxWidth)
	}
	if p.MaxHeight > 0 && int(req.Options.Height) > p.MaxHeight {
		req.Options.Height = float64(p.MaxHeight)
	}

	/*
	if !p.allowed(req.URL) {
		glog.Errorf("remote URL is not for an allowed host: %v", req.URL)
		http.Error(w, fmt.Sprintf("remote URL is not for an allowed host: %v", req.URL), http.StatusBadRequest)
		return
	}
	*/

	//u := req.URL.String()
	if req.Options != nil && !reflect.DeepEqual(req.Options, emptyOptions) {
		u += "#" + req.Options.String()
	}

	resp, err := p.Client.Get(u)
	if err != nil {

		http.Error(w, fmt.Printf("%v", err), 500)
		return

		glog.Errorf("error fetching remote image: %v", err)
		http.Error(w, fmt.Sprintf("Error fetching remote image: %v", err), http.StatusInternalServerError)
		return
	}

	if resp.StatusCode != http.StatusOK {
		http.Error(w, fmt.Sprintf("Remote URL %q returned status: %v", req.URL, resp.Status), http.StatusInternalServerError)
		return
	}

	w.Header().Add("Last-Modified", resp.Header.Get("Last-Modified"))
	w.Header().Add("Expires", resp.Header.Get("Expires"))
	w.Header().Add("Etag", resp.Header.Get("Etag"))

	if is304 := check304(w, r, resp); is304 {
		w.WriteHeader(http.StatusNotModified)
		return
	}

	w.Header().Add("Content-Length", resp.Header.Get("Content-Length"))
	defer resp.Body.Close()
	io.Copy(w, resp.Body)
}

// allowed returns whether the specified URL is on the whitelist of remote hosts.
func (p *Proxy) allowed(u *url.URL) bool {
	if len(p.Whitelist) == 0 {
		return true
	}

	for _, host := range p.Whitelist {
		if u.Host == host {
			return true
		}
		if strings.HasPrefix(host, "*.") && strings.HasSuffix(u.Host, host[2:]) {
			return true
		}
	}

	return false
}

func check304(w http.ResponseWriter, req *http.Request, resp *http.Response) bool {
	etag := resp.Header.Get("Etag")
	if etag != "" && etag == req.Header.Get("If-None-Match") {
		return true
	}

	lastModified, err := time.Parse(time.RFC1123, resp.Header.Get("Last-Modified"))
	if err != nil {
		return false
	}
	ifModSince, err := time.Parse(time.RFC1123, req.Header.Get("If-Modified-Since"))
	if err != nil {
		return false
	}
	if lastModified.Before(ifModSince) {
		return true
	}

	return false
}

// TransformingTransport is an implementation of http.RoundTripper that
// optionally transforms images using the options specified in the request URL
// fragment.
type TransformingTransport struct {
	// Transport is used to satisfy non-transform requests (those that do not include a URL fragment)
	Transport http.RoundTripper

	// Client is used to fetch images to be resized.
	Client *http.Client
}

// RoundTrip implements http.RoundTripper.
func (t *TransformingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if req.URL.Fragment == "" {
		// normal requests pass through
		glog.Infof("fetching remote URL: %v", req.URL)
		return t.Transport.RoundTrip(req)
	}

	u := *req.URL
	u.Fragment = ""
	resp, err := t.Client.Get(u.String())

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

	opt := ParseOptions(req.URL.Fragment)
	img, err := Transform(b, opt)
	if err != nil {
		img = b
	}

	// replay response with transformed image and updated content length
	buf := new(bytes.Buffer)
	fmt.Fprintf(buf, "%s %s\n", resp.Proto, resp.Status)
	resp.Header.WriteSubset(buf, map[string]bool{"Content-Length": true})
	fmt.Fprintf(buf, "Content-Length: %d\n\n", len(img))
	buf.Write(img)

	return http.ReadResponse(bufio.NewReader(buf), req)
}
<|endoftext|>


<commit_before>package proxy

import (
	"fmt"
	"log"
	"sort"
	"time"

	"github.com/SoCloz/goprismic"
)

type Config struct {
	// Cache size
	CacheSize int
	// Documents are cached for a maximum time of ttl.
	TTL time.Duration
	// API Master ref refresh frequency
	MasterRefresh time.Duration
	// Base refresh chance after master reference chance
	BaselineRefreshChance float32
	// inherited from api
	debug bool
}

type Proxy struct {
	cache *Cache
	api   *goprismic.Api

	Config Config

	lastRefresh           time.Time
}

// Creates a proxy
//
// All documents are cached in a LRU cache of cacheSize elements.
func New(u, accessToken string, apiCfg goprismic.Config, cfg Config) (*Proxy, error) {
	a, err := goprismic.Get(u, accessToken, apiCfg)
	if err != nil {
		return nil, err
	}
	if cfg.MasterRefresh == 0 {
		cfg.MasterRefresh = time.Minute
	}
	cfg.debug = apiCfg.Debug
	if cfg.BaselineRefreshChance == 0 {
		cfg.BaselineRefreshChance = 1.0
	}
	c := NewCache(cfg.CacheSize, cfg.TTL, 1.0)
	c.revision = a.GetMasterRef()
	p := &Proxy{
		cache:                 c,
		api:                   a,
		Config:                cfg,
		lastRefresh:           time.Now(),
	}
	go p.loopRefresh()
	return p, nil
}

// Returns the cache stats
func (p *Proxy) GetStats() Stats {
	return p.cache.Stats
}

// Queries the API
//
//   proxy.Direct().Master().Form("everything").Submit()
func (p *Proxy) Direct() *goprismic.Api {
	return p.api
}

// Fetches a document by id
func (p *Proxy) GetDocument(id string) (*goprismic.Document, error) {
	sr, err := p.Search().Form("everything").Query("[[:d = at(document.id, \"" + id + "\")]]").Submit()
	if err != nil || sr.TotalResults == 0 {
		return nil, err
	}
	return &sr.Results[0], nil

}

// Fetches a document of a specific type by a field value
func (p *Proxy) GetDocumentBy(docType, field string, value interface{}) (*goprismic.Document, error) {
	query := fmt.Sprintf("[[:d = at(my.%s.%s, \"%v\")][:d = any(document.type, [\"%s\"])]]", docType, field, value, docType)
	sr, err := p.Search().Form("everything").Query(query).Submit()
	if err != nil || sr.TotalResults == 0 {
		return nil, err
	}
	return &sr.Results[0], nil
}

// Search documents
func (p *Proxy) Search() *SearchForm {
	f := &SearchForm{sf: p.api.Master(), p: p}
	f.sig = sort.StringSlice{}
	return f
}

// Refreshes the master ref, returns true if master ref has changed
func (p *Proxy) Refresh() bool {
	p.api.Refresh()
	oldRev := p.cache.revision
	p.cache.revision = p.api.GetMasterRef()
	if oldRev != p.cache.revision {
		p.lastRefresh = time.Now()
		// refresh : 100% cache miss expected => we switch to baseline
		p.cache.refreshChance = p.Config.BaselineRefreshChance
		if p.Config.debug {
			log.Printf("Prismic - refreshing master ref and lowering refresh chance to %.1f%%", p.cache.refreshChance*100)
		}
		return true
	}
	return false
}

// Refreshes the master ref, returns true if master ref has changed
func (p *Proxy) loopRefresh() {
	prevRefreshError, prevRefresh := 0, 0
	tick := time.Tick(p.Config.MasterRefresh)
	for {
		select {
		case <-tick:
			refreshed := p.Refresh()
			deltaRefreshError := p.cache.Stats.RefreshError - prevRefreshError
			deltaRefresh := p.cache.Stats.Refresh - prevRefresh
			prevRefreshError, prevRefresh = p.cache.Stats.RefreshError, p.cache.Stats.Refresh
			if !refreshed && deltaRefresh+deltaRefreshError > 0 {
				refreshChance := float32(deltaRefresh) / float32(deltaRefresh+deltaRefreshError) * p.cache.refreshChance
				if refreshChance < p.cache.refreshChance {
					p.cache.refreshChance = refreshChance
					if p.Config.debug {
						log.Printf("Prismic - lowering refresh chance to %.2f%%", p.cache.refreshChance*100.0)
					}
				}
				if refreshChance > p.cache.refreshChance {
					if refreshChance < p.cache.refreshChance*1.1 {
						p.cache.refreshChance = refreshChance
					} else {
						p.cache.refreshChance = p.cache.refreshChance * 1.1
					}
					if p.cache.refreshChance >= 1.0 {
						p.cache.refreshChance = 1.0
					}
					if p.Config.debug {
						log.Printf("Prismic - raising refresh chance to %.2f%%", p.cache.refreshChance*100.0)
					}
				}
			}
		}
	}
}

// Get something from the proxy
func (p *Proxy) Get(key string, refresh RefreshFn) (interface{}, error) {
	return p.cache.Get(key, refresh)
}

// Clears the cache
func (p *Proxy) Clear() {
	p.cache.Clear()
}
<commit_msg>proxy: refresh chance was not raised properly<commit_after>package proxy

import (
	"fmt"
	"log"
	"sort"
	"time"

	"github.com/SoCloz/goprismic"
)

type Config struct {
	// Cache size
	CacheSize int
	// Documents are cached for a maximum time of ttl.
	TTL time.Duration
	// API Master ref refresh frequency
	MasterRefresh time.Duration
	// Base refresh chance after master reference chance
	BaselineRefreshChance float32
	// inherited from api
	debug bool
}

type Proxy struct {
	cache *Cache
	api   *goprismic.Api

	Config Config

	lastRefresh           time.Time
}

// Creates a proxy
//
// All documents are cached in a LRU cache of cacheSize elements.
func New(u, accessToken string, apiCfg goprismic.Config, cfg Config) (*Proxy, error) {
	a, err := goprismic.Get(u, accessToken, apiCfg)
	if err != nil {
		return nil, err
	}
	if cfg.MasterRefresh == 0 {
		cfg.MasterRefresh = time.Minute
	}
	cfg.debug = apiCfg.Debug
	if cfg.BaselineRefreshChance == 0 {
		cfg.BaselineRefreshChance = 1.0
	}
	c := NewCache(cfg.CacheSize, cfg.TTL, 1.0)
	c.revision = a.GetMasterRef()
	p := &Proxy{
		cache:                 c,
		api:                   a,
		Config:                cfg,
		lastRefresh:           time.Now(),
	}
	go p.loopRefresh()
	return p, nil
}

// Returns the cache stats
func (p *Proxy) GetStats() Stats {
	return p.cache.Stats
}

// Queries the API
//
//   proxy.Direct().Master().Form("everything").Submit()
func (p *Proxy) Direct() *goprismic.Api {
	return p.api
}

// Fetches a document by id
func (p *Proxy) GetDocument(id string) (*goprismic.Document, error) {
	sr, err := p.Search().Form("everything").Query("[[:d = at(document.id, \"" + id + "\")]]").Submit()
	if err != nil || sr.TotalResults == 0 {
		return nil, err
	}
	return &sr.Results[0], nil

}

// Fetches a document of a specific type by a field value
func (p *Proxy) GetDocumentBy(docType, field string, value interface{}) (*goprismic.Document, error) {
	query := fmt.Sprintf("[[:d = at(my.%s.%s, \"%v\")][:d = any(document.type, [\"%s\"])]]", docType, field, value, docType)
	sr, err := p.Search().Form("everything").Query(query).Submit()
	if err != nil || sr.TotalResults == 0 {
		return nil, err
	}
	return &sr.Results[0], nil
}

// Search documents
func (p *Proxy) Search() *SearchForm {
	f := &SearchForm{sf: p.api.Master(), p: p}
	f.sig = sort.StringSlice{}
	return f
}

// Refreshes the master ref, returns true if master ref has changed
func (p *Proxy) Refresh() bool {
	p.api.Refresh()
	oldRev := p.cache.revision
	p.cache.revision = p.api.GetMasterRef()
	if oldRev != p.cache.revision {
		p.lastRefresh = time.Now()
		// refresh : 100% cache miss expected => we switch to baseline
		p.cache.refreshChance = p.Config.BaselineRefreshChance
		if p.Config.debug {
			log.Printf("Prismic - refreshing master ref and lowering refresh chance to %.1f%%", p.cache.refreshChance*100)
		}
		return true
	}
	return false
}

// Refreshes the master ref, returns true if master ref has changed
func (p *Proxy) loopRefresh() {
	prevRefreshError, prevRefresh, lastNoError := 0, 0, 0
	tick := time.Tick(p.Config.MasterRefresh)
	for {
		select {
		case <-tick:
			refreshed := p.Refresh()
			if refreshed {
				lastNoError = 0
			} else {
				deltaRefreshError := p.cache.Stats.RefreshError - prevRefreshError
				deltaRefresh := p.cache.Stats.Refresh - prevRefresh
				prevRefreshError, prevRefresh = p.cache.Stats.RefreshError, p.cache.Stats.Refresh
				var refreshChance float32
				if deltaRefreshError > 0 {
					lastNoError = 0
					refreshChance = float32(deltaRefresh) / float32(deltaRefresh+deltaRefreshError) * p.cache.refreshChance
				} else {
					lastNoError++
					refreshChance = p.Config.BaselineRefreshChance*(1.0+float32(lastNoError*(lastNoError-1))/3)
					if refreshChance > 1.0 {
						refreshChance = 1.0
					}
				}
				if refreshChance < p.cache.refreshChance {
					p.cache.refreshChance = refreshChance
					if p.Config.debug {
						log.Printf("Prismic - lowering refresh chance to %.2f%%", p.cache.refreshChance*100.0)
					}
				}
				if refreshChance > p.cache.refreshChance {
					if refreshChance >= 1.0 && p.cache.refreshChance >= 0.99 {
						p.cache.refreshChance = 1.0
					} else {
						p.cache.refreshChance = (p.cache.refreshChance + refreshChance) / 2
					}
					if p.Config.debug {
						log.Printf("Prismic - raising refresh chance to %.2f%%", p.cache.refreshChance*100.0)
					}
				}
			}
		}
	}
}

// Get something from the proxy
func (p *Proxy) Get(key string, refresh RefreshFn) (interface{}, error) {
	return p.cache.Get(key, refresh)
}

// Clears the cache
func (p *Proxy) Clear() {
	p.cache.Clear()
}
<|endoftext|>


<commit_before>package proxy

import (
	"bytes"
	"fmt"
	"github.com/robryk/zproxy/cache"
	"github.com/robryk/zproxy/hasher"
	"io"
	"io/ioutil"
	"log"
	"net/http"
)

const SizeCutoff = 10 // for testing

type Proxy struct {
	Cr    hasher.ChunkedRetriever
	Cache cache.Cache
}

func getSize(url string) (int64, error) {
	resp, err := http.Head(url)
	if err != nil {
		return 0, err
	}
	resp.Body.Close()
	return resp.ContentLength, nil
}

type byteRange struct {
	begin int64
	end   int64
}

func getContents(url string, r byteRange) ([]byte, http.Header, int, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, nil, 0, err
	}
	if r.begin != 0 || r.end != 0 {
		req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", r.begin, r.end))
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, nil, 0, err
	}
	defer resp.Body.Close()

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

	return contents, resp.Header, resp.StatusCode, nil
}

func copyHeaders(dst http.Header, src http.Header) {
	for k := range dst {
		delete(dst, k)
	}
	for k, v := range src {
		dst[k] = v
	}
}

func directProxy(rw http.ResponseWriter, url string) {
	resp, err := http.Get(url)
	if err != nil {
		rw.WriteHeader(http.StatusInternalServerError)
		return
	}
	defer resp.Body.Close()

	copyHeaders(rw.Header(), resp.Header)
	rw.WriteHeader(resp.StatusCode)
	io.Copy(rw, resp.Body)
}

func (p *Proxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if req.Method != "GET" {
		rw.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	url := req.URL.String()
	log.Printf("Serving %s", url)

	if size, err := getSize(url); size < SizeCutoff && size != -1 && err == nil {
		log.Printf("Size %d, serving directly", size)
		directProxy(rw, url)
		return
	}
	chunked, err := p.Cr.GetChunked(url)
	if err != nil {
		directProxy(rw, url)
		return
	}

	// TODO: deal with status and headers
	// TODO: emit content-length!

	for _, chunk := range chunked.Chunks {
		r, err := p.Cache.Read(string(chunk.Digest))
		if err == nil {
			io.Copy(rw, r)
		} else {
			// TODO: handle cache write failures
			contents, _, _, err := getContents(string(url), byteRange{int64(chunk.Offset), int64(chunk.Offset + chunk.Length)})
			if err != nil {
				log.Printf("Error retrieving missing chunk: %v", err)
				return
			}
			// errors fixme
			go p.Cache.Write(string(chunk.Digest), bytes.NewBuffer(contents))
			rw.Write(contents)
		}
	}
}
<commit_msg>Proxy now sends a correct initial HEAD probe.<commit_after>package proxy

import (
	"bytes"
	"fmt"
	"github.com/robryk/zproxy/cache"
	"github.com/robryk/zproxy/hasher"
	"io"
	"io/ioutil"
	"log"
	"net/http"
)

const SizeCutoff = 10 // for testing

type Proxy struct {
	Cr    hasher.ChunkedRetriever
	Cache cache.Cache
}

func sanitizeRequest(req *http.Request) *http.Request {
	r := &http.Request{}
	// FIXME: What if old proto?

	r.Method = req.Method
	r.URL = req.URL // FIXME: Should this be deeper?
	// Proto* are ignored on outgoing requests
	r.Header = make(http.Header)
	for k, vv := range req.Header {
		if k == "Connection" {
			continue
		}
		for _, v := range vv {
			r.Header.Add(k, v)
		}
		// FIXME: Remove other headers
	}
	r.Body = req.Body // We might wish to warn about this
	r.ContentLength = req.ContentLength
	if r.ContentLength == -1 {
		r.ContentLength = 0 // necessary?
	}
	// FIXME: Deal better with Transfer-Encoding
	r.TransferEncoding = nil
	r.Close = false
	r.Host = req.Host // necessary?
	// *Form is ignored on outgoing requests
	// We set nil Trailer, because there is no way for us to set it early enough (or so I think)
	// RemoteAddr, RequestURI and TLS make no sense on an outgoing request
	return r
}

type byteRange struct {
	begin int64
	end   int64
}

func getContents(url string, r byteRange) ([]byte, http.Header, int, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, nil, 0, err
	}
	if r.begin != 0 || r.end != 0 {
		req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", r.begin, r.end))
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, nil, 0, err
	}
	defer resp.Body.Close()

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

	return contents, resp.Header, resp.StatusCode, nil
}

func copyHeaders(dst http.Header, src http.Header) {
	for k := range dst {
		delete(dst, k)
	}
	for k, v := range src {
		dst[k] = v
	}
}

func directProxy(rw http.ResponseWriter, url string) {
	resp, err := http.Get(url)
	if err != nil {
		rw.WriteHeader(http.StatusInternalServerError)
		return
	}
	defer resp.Body.Close()

	copyHeaders(rw.Header(), resp.Header)
	rw.WriteHeader(resp.StatusCode)
	io.Copy(rw, resp.Body)
}

func isDirect(req *http.Request) bool {
	if req.Method != "GET" {
		return true
	}

	headReq := sanitizeRequest(req)
	headReq.Method = "HEAD"
	headReq.Body = nil
	headReq.ContentLength = 0
	headResp, err := http.DefaultClient.Do(headReq)
	if err != nil {
		log.Printf("Head request to %s failed: %v", req.URL, err)
		return true
	}
	headResp.Body.Close()
	if headResp.ContentLength < SizeCutoff || headResp.ContentLength == -1 {
		// If the server can't return a Content-Length, chances are high the page changes often
		// even if the server claims otherwise.
		return true
	}
	// FIXME: Cache-Control
	return false
}

func (p *Proxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if req.Method != "GET" {
		rw.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	url := req.URL.String()
	log.Printf("Serving %s", url)

	if isDirect(req) {
		log.Printf("Serving directly")
		directProxy(rw, url)
		return
	}
	chunked, err := p.Cr.GetChunked(url)
	if err != nil {
		directProxy(rw, url)
		return
	}

	// TODO: deal with status and headers
	// TODO: emit content-length!

	for _, chunk := range chunked.Chunks {
		r, err := p.Cache.Read(string(chunk.Digest))
		if err == nil {
			io.Copy(rw, r)
		} else {
			// TODO: handle cache write failures
			contents, _, _, err := getContents(string(url), byteRange{int64(chunk.Offset), int64(chunk.Offset + chunk.Length)})
			if err != nil {
				log.Printf("Error retrieving missing chunk: %v", err)
				return
			}
			// errors fixme
			go p.Cache.Write(string(chunk.Digest), bytes.NewBuffer(contents))
			rw.Write(contents)
		}
	}
}
<|endoftext|>


<commit_before>package proxy

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/sorah/etcvault/engine"
	"io"
	"io/ioutil"
	"log"
	"net/http"
)

type ClosableBuffer struct {
	*bytes.Buffer
}

func (buf ClosableBuffer) Close() error {
	return nil
}

// Hop-by-hop headers (borrowed from httputil.ReverseProxy)
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var singleHopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te",
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

type Proxy struct {
	Transport    *http.Transport
	Router       *Router
	Engine       engine.Transformable
	AdvertiseUrl string
}

func NewProxy(transport *http.Transport, router *Router, e engine.Transformable, advertiseUrl string) http.Handler {
	return &Proxy{
		Transport:    transport,
		Router:       router,
		Engine:       e,
		AdvertiseUrl: advertiseUrl,
	}
}

func (proxy *Proxy) ServeHTTP(response http.ResponseWriter, request *http.Request) {
	if request.URL.Path == "/v2/members" {
		proxy.serveMembersRequest(response, request)
	} else if request.URL.Path == "/v2/machines" {
		proxy.serveMachinesRequest(response, request)
	} else if request.URL.Path == "/_etcvault/keys" {
		proxy.serveEtcvaultKeysRequest(response, request)
	} else {
		proxy.serveProxyRequest(response, request)
	}
}

func (proxy *Proxy) serveProxyRequest(response http.ResponseWriter, request *http.Request) {
	backendRequest := new(http.Request)
	// copy
	*backendRequest = *request
	backendRequest.Header = make(http.Header)

	backendRequest.Proto = "HTTP/1.1"
	backendRequest.ProtoMajor = 1
	backendRequest.ProtoMinor = 1
	backendRequest.Close = false

	copyHeader(request.Header, backendRequest.Header)
	removeSingleHopHeaders(&backendRequest.Header)

	if (backendRequest.Method == "POST" || backendRequest.Method == "PUT" || backendRequest.Method == "PATCH") && backendRequest.Body != nil {
		origBody := backendRequest.Body
		defer origBody.Close()

		if err := backendRequest.ParseForm(); err != nil {
			log.Printf("couldn't parse form: %s", err.Error())
			http.Error(response, "couldn't parse form", 400)
			return
		}

		if backendRequest.PostForm != nil {
			origValue := backendRequest.PostForm.Get("value")
			value, err := proxy.Engine.Transform(origValue)
			if err == nil {
				backendRequest.PostForm.Set("value", value)
			} else {
				log.Printf("failed to transform value: %s", err.Error())
			}
			newFormString := backendRequest.PostForm.Encode()
			backendRequest.Body = ClosableBuffer{bytes.NewBufferString(newFormString)}
			backendRequest.ContentLength = int64(len(newFormString))
		}
	}

	var backendResponse *http.Response

	backends := proxy.Router.ShuffledAvailableBackends()
	for _, backend := range backends {
		backendRequest.URL.Scheme = backend.Url.Scheme
		backendRequest.URL.Host = backend.Url.Host

		var err error
		backendResponse, err = proxy.Transport.RoundTrip(backendRequest)
		if err != nil {
			log.Printf("backend %s response error: %s", backend.Url.String(), err.Error())
			backend.Fail()
			continue
		}
		backend.Ok()
		break
	}

	if backendResponse == nil {
		log.Printf("all backends not available...")
		http.Error(response, "backends all unavailable", http.StatusBadGateway)
		return
	}

	defer backendResponse.Body.Close()

	removeSingleHopHeaders(&backendResponse.Header)
	copyHeader(backendResponse.Header, response.Header())

	if backendResponse.Header.Get("Content-Type") == "application/json" {
		json, err := ioutil.ReadAll(backendResponse.Body)
		if err != nil {
			panic(err)
		}

		transformedJson, err := proxy.Engine.TransformEtcdJsonResponse(json)
		if err == nil {
			response.Header().Set("Content-Length", fmt.Sprintf("%d", len(transformedJson)+1))
			response.WriteHeader(backendResponse.StatusCode)
			response.Write(transformedJson)
			response.Write([]byte("\n"))
		} else {
			fmt.Printf("transform error %s\n", err.Error())
			response.WriteHeader(backendResponse.StatusCode)
			response.Write(json)
		}
	} else {
		response.WriteHeader(backendResponse.StatusCode)
		io.Copy(response, backendResponse.Body)
	}
}

func (proxy *Proxy) serveMembersRequest(response http.ResponseWriter, request *http.Request) {
	if request.Method != "GET" {
		http.Error(response, "not supported; communicate with etcd directly", http.StatusMethodNotAllowed)
		return
	}

	type memberT struct {
		ClientURLs []string
		PeerURLs   []string
		Name       string
		Id         string
	}

	jsonBytes, err := json.Marshal(struct {
		Members []memberT
	}{
		Members: []memberT{
			{
				ClientURLs: []string{proxy.AdvertiseUrl},
				Name:       "etcvault",
				Id:         "deadbeef",
			},
		},
	})

	if err != nil {
		http.Error(response, "failed to marshal", 500)
		log.Printf("failed to marshal /v2/members: %s", err.Error())
		return
	}

	response.Header().Add("Content-Type", "application/json")
	response.Header().Add("Server", "etcvault")
	response.WriteHeader(200)
	response.Write(jsonBytes)
}

func (proxy *Proxy) serveMachinesRequest(response http.ResponseWriter, request *http.Request) {
	if request.Method != "GET" {
		http.Error(response, "not supported; communicate with etcd directly", http.StatusMethodNotAllowed)
		return
	}

	response.Header().Add("Content-Type", "text/plain")
	response.Header().Add("Server", "etcvault")
	response.WriteHeader(200)
	response.Write([]byte(proxy.AdvertiseUrl))
}

func (proxy *Proxy) serveEtcvaultKeysRequest(response http.ResponseWriter, request *http.Request) {
	if request.Method != "GET" {
		http.Error(response, "not found", http.StatusNotFound)
		return
	}

	request.ParseForm()
	var list []string
	if request.FormValue("encryption") != "" {
		list = proxy.Engine.GetKeychain().ListForEncryption()
	} else {
		list = proxy.Engine.GetKeychain().List()
	}

	response.Header().Add("Content-Type", "text/plain")
	response.Header().Add("Server", "etcvault")
	response.WriteHeader(200)

	for _, name := range list {
		response.Write([]byte(name))
		response.Write([]byte("\n"))
	}
}

func copyHeader(source, destination http.Header) {
	for key, values := range source {
		for _, value := range values {
			destination.Add(key, value)
		}
	}
}

func removeSingleHopHeaders(header *http.Header) {
	for _, name := range singleHopHeaders {
		header.Del(name)
	}
}
<commit_msg>Close backend connection when client disconnects<commit_after>package proxy

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/sorah/etcvault/engine"
	"io"
	"io/ioutil"
	"log"
	"net/http"
)

type ClosableBuffer struct {
	*bytes.Buffer
}

func (buf ClosableBuffer) Close() error {
	return nil
}

// Hop-by-hop headers (borrowed from httputil.ReverseProxy)
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var singleHopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te",
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

type Proxy struct {
	Transport    *http.Transport
	Router       *Router
	Engine       engine.Transformable
	AdvertiseUrl string
}

func NewProxy(transport *http.Transport, router *Router, e engine.Transformable, advertiseUrl string) http.Handler {
	return &Proxy{
		Transport:    transport,
		Router:       router,
		Engine:       e,
		AdvertiseUrl: advertiseUrl,
	}
}

func (proxy *Proxy) ServeHTTP(response http.ResponseWriter, request *http.Request) {
	if request.URL.Path == "/v2/members" {
		proxy.serveMembersRequest(response, request)
	} else if request.URL.Path == "/v2/machines" {
		proxy.serveMachinesRequest(response, request)
	} else if request.URL.Path == "/_etcvault/keys" {
		proxy.serveEtcvaultKeysRequest(response, request)
	} else {
		proxy.serveProxyRequest(response, request)
	}
}

func (proxy *Proxy) serveProxyRequest(response http.ResponseWriter, request *http.Request) {
	backendRequest := new(http.Request)
	// copy
	*backendRequest = *request
	backendRequest.Header = make(http.Header)

	backendRequest.Proto = "HTTP/1.1"
	backendRequest.ProtoMajor = 1
	backendRequest.ProtoMinor = 1
	backendRequest.Close = false

	copyHeader(request.Header, backendRequest.Header)
	removeSingleHopHeaders(&backendRequest.Header)

	if (backendRequest.Method == "POST" || backendRequest.Method == "PUT" || backendRequest.Method == "PATCH") && backendRequest.Body != nil {
		origBody := backendRequest.Body
		defer origBody.Close()

		if err := backendRequest.ParseForm(); err != nil {
			log.Printf("couldn't parse form: %s", err.Error())
			http.Error(response, "couldn't parse form", 400)
			return
		}

		if backendRequest.PostForm != nil {
			origValue := backendRequest.PostForm.Get("value")
			value, err := proxy.Engine.Transform(origValue)
			if err == nil {
				backendRequest.PostForm.Set("value", value)
			} else {
				log.Printf("failed to transform value: %s", err.Error())
			}
			newFormString := backendRequest.PostForm.Encode()
			backendRequest.Body = ClosableBuffer{bytes.NewBufferString(newFormString)}
			backendRequest.ContentLength = int64(len(newFormString))
		}
	}

	var backendResponse *http.Response

	var closeNotifyCh <-chan bool
	closeNotifier, ok := response.(http.CloseNotifier)
	if ok {
		closeNotifyCh = closeNotifier.CloseNotify()
	} else {
		closeNotifyCh = make(<-chan bool)
	}

	completeCh := make(chan bool, 2)
	closed := false
	go func() {
		select {
		case <-closeNotifyCh:
			log.Printf("Request connection closed; cancelling ongoing backend request")
			closed = true
			proxy.Transport.CancelRequest(backendRequest)
		case <-completeCh:
		}
		if backendResponse != nil {
			backendResponse.Body.Close()
		}
	}()
	defer func() {
		completeCh <- true
	}()

	backends := proxy.Router.ShuffledAvailableBackends()
	for _, backend := range backends {
		backendRequest.URL.Scheme = backend.Url.Scheme
		backendRequest.URL.Host = backend.Url.Host

		var err error
		backendResponse, err = proxy.Transport.RoundTrip(backendRequest)
		if err != nil {
			log.Printf("backend %s response error: %s", backend.Url.String(), err.Error())
			backend.Fail()
			continue
		}
		backend.Ok()
		break
	}

	if backendResponse == nil {
		log.Printf("all backends not available...")
		http.Error(response, "backends all unavailable", http.StatusBadGateway)
		return
	}

	defer backendResponse.Body.Close()

	removeSingleHopHeaders(&backendResponse.Header)
	copyHeader(backendResponse.Header, response.Header())

	if backendResponse.Header.Get("Content-Type") == "application/json" {
		json, err := ioutil.ReadAll(backendResponse.Body)
		if closed {
			return
		}
		if err != nil {
			panic(err)
		}

		transformedJson, err := proxy.Engine.TransformEtcdJsonResponse(json)
		if err == nil {
			response.Header().Set("Content-Length", fmt.Sprintf("%d", len(transformedJson)+1))
			response.WriteHeader(backendResponse.StatusCode)
			response.Write(transformedJson)
			response.Write([]byte("\n"))
		} else {
			fmt.Printf("transform error %s\n", err.Error())
			response.WriteHeader(backendResponse.StatusCode)
			response.Write(json)
		}
	} else {
		response.WriteHeader(backendResponse.StatusCode)
		io.Copy(response, backendResponse.Body)
	}
}

func (proxy *Proxy) serveMembersRequest(response http.ResponseWriter, request *http.Request) {
	if request.Method != "GET" {
		http.Error(response, "not supported; communicate with etcd directly", http.StatusMethodNotAllowed)
		return
	}

	type memberT struct {
		ClientURLs []string
		PeerURLs   []string
		Name       string
		Id         string
	}

	jsonBytes, err := json.Marshal(struct {
		Members []memberT
	}{
		Members: []memberT{
			{
				ClientURLs: []string{proxy.AdvertiseUrl},
				Name:       "etcvault",
				Id:         "deadbeef",
			},
		},
	})

	if err != nil {
		http.Error(response, "failed to marshal", 500)
		log.Printf("failed to marshal /v2/members: %s", err.Error())
		return
	}

	response.Header().Add("Content-Type", "application/json")
	response.Header().Add("Server", "etcvault")
	response.WriteHeader(200)
	response.Write(jsonBytes)
}

func (proxy *Proxy) serveMachinesRequest(response http.ResponseWriter, request *http.Request) {
	if request.Method != "GET" {
		http.Error(response, "not supported; communicate with etcd directly", http.StatusMethodNotAllowed)
		return
	}

	response.Header().Add("Content-Type", "text/plain")
	response.Header().Add("Server", "etcvault")
	response.WriteHeader(200)
	response.Write([]byte(proxy.AdvertiseUrl))
}

func (proxy *Proxy) serveEtcvaultKeysRequest(response http.ResponseWriter, request *http.Request) {
	if request.Method != "GET" {
		http.Error(response, "not found", http.StatusNotFound)
		return
	}

	request.ParseForm()
	var list []string
	if request.FormValue("encryption") != "" {
		list = proxy.Engine.GetKeychain().ListForEncryption()
	} else {
		list = proxy.Engine.GetKeychain().List()
	}

	response.Header().Add("Content-Type", "text/plain")
	response.Header().Add("Server", "etcvault")
	response.WriteHeader(200)

	for _, name := range list {
		response.Write([]byte(name))
		response.Write([]byte("\n"))
	}
}

func copyHeader(source, destination http.Header) {
	for key, values := range source {
		for _, value := range values {
			destination.Add(key, value)
		}
	}
}

func removeSingleHopHeaders(header *http.Header) {
	for _, name := range singleHopHeaders {
		header.Del(name)
	}
}
<|endoftext|>


<commit_before>// Package proxy is a cli proxy
package proxy

import (
	"os"
	"strings"
	"time"

	"github.com/micro/cli"
	"github.com/micro/go-micro"
	bmem "github.com/micro/go-micro/broker/memory"
	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/config/options"
	"github.com/micro/go-micro/proxy"
	"github.com/micro/go-micro/proxy/grpc"
	"github.com/micro/go-micro/proxy/http"
	"github.com/micro/go-micro/proxy/mucp"
	"github.com/micro/go-micro/registry"
	rmem "github.com/micro/go-micro/registry/memory"
	"github.com/micro/go-micro/router"
	rs "github.com/micro/go-micro/router/service"
	"github.com/micro/go-micro/server"
	sgrpc "github.com/micro/go-micro/server/grpc"
	smucp "github.com/micro/go-micro/server/mucp"
	"github.com/micro/go-micro/util/log"
	"github.com/micro/go-micro/util/mux"
)

var (
	// Name of the proxy
	Name = "go.micro.proxy"
	// The address of the proxy
	Address = ":8081"
	// the proxy protocol
	Protocol = "mucp"
	// The endpoint host to route to
	Endpoint string
)

func run(ctx *cli.Context, srvOpts ...micro.Option) {
	log.Name("proxy")

	if len(ctx.GlobalString("server_name")) > 0 {
		Name = ctx.GlobalString("server_name")
	}
	if len(ctx.String("address")) > 0 {
		Address = ctx.String("address")
	}
	if len(ctx.String("endpoint")) > 0 {
		Endpoint = ctx.String("endpoint")
	}
	if len(ctx.String("protocol")) > 0 {
		Protocol = ctx.String("protocol")
	}

	// Init plugins
	for _, p := range Plugins() {
		p.Init(ctx)
	}

	// service opts
	srvOpts = append(srvOpts, micro.Name(Name))
	if i := time.Duration(ctx.GlobalInt("register_ttl")); i > 0 {
		srvOpts = append(srvOpts, micro.RegisterTTL(i*time.Second))
	}
	if i := time.Duration(ctx.GlobalInt("register_interval")); i > 0 {
		srvOpts = append(srvOpts, micro.RegisterInterval(i*time.Second))
	}

	// set the context
	var popts []options.Option

	// create new router
	var r router.Router

	routerName := ctx.String("router")
	routerAddr := ctx.String("router_address")

	ropts := []router.Option{
		router.Id(server.DefaultId),
		router.Client(client.DefaultClient),
		router.Address(routerAddr),
		router.Registry(registry.DefaultRegistry),
	}

	// check if we need to use the router service
	switch {
	case routerName == "go.micro.router":
		r = rs.NewRouter(ropts...)
	case len(routerAddr) > 0:
		r = rs.NewRouter(ropts...)
	default:
		r = router.NewRouter(ropts...)
	}

	// start the router
	if err := r.Start(); err != nil {
		log.Logf("Proxy error starting router: %s", err)
		os.Exit(1)
	}

	popts = append(popts, proxy.WithRouter(r))

	// new proxy
	var p proxy.Proxy
	var s server.Server

	// set endpoint
	if len(Endpoint) > 0 {
		switch {
		case strings.HasPrefix(Endpoint, "grpc://"):
			ep := strings.TrimPrefix(Endpoint, "grpc://")
			popts = append(popts, proxy.WithEndpoint(ep))
			p = grpc.NewProxy(popts...)
		case strings.HasPrefix(Endpoint, "http://"):
			// TODO: strip prefix?
			popts = append(popts, proxy.WithEndpoint(Endpoint))
			p = http.NewProxy(popts...)
		default:
			// TODO: strip prefix?
			popts = append(popts, proxy.WithEndpoint(Endpoint))
			p = mucp.NewProxy(popts...)
		}
	}

	// set based on protocol
	if p == nil && len(Protocol) > 0 {
		switch Protocol {
		case "http":
			p = http.NewProxy(popts...)
			// TODO: http server
		case "grpc":
			p = grpc.NewProxy(popts...)
			s = sgrpc.NewServer()
		default:
			p = mucp.NewProxy(popts...)
			s = smucp.NewServer()
		}
	}

	if len(Endpoint) > 0 {
		log.Logf("Proxy [%s] serving endpoint: %s", p.String(), Endpoint)
	} else {
		log.Logf("Proxy [%s] serving protocol: %s", p.String(), Protocol)
	}

	// prepend the server
	if s != nil {
		srvOpts = append([]micro.Option{micro.Server(s)}, srvOpts...)
	}

	// new service
	service := micro.NewService(srvOpts...)

	// create a new proxy muxer which includes the debug handler
	muxer := mux.New(Name, p)

	// set the router
	service.Server().Init(
		server.WithRouter(muxer),
	)

	// gen the straight up http server
	// set address
	srv := server.NewServer(
		server.Address(Address),
		// reset registry to memory
		server.Registry(rmem.NewRegistry()),
		// reset broker to memory
		server.Broker(bmem.NewBroker()),
		// hande it the router
		server.WithRouter(p),
	)

	// Start the server
	if err := srv.Start(); err != nil {
		log.Fatal(err)
	}

	// Run internal service
	if err := service.Run(); err != nil {
		log.Fatal(err)
	}

	// Stop the server
	if err := srv.Stop(); err != nil {
		log.Fatal(err)
	}
}

func Commands(options ...micro.Option) []cli.Command {
	command := cli.Command{
		Name:  "proxy",
		Usage: "Run the service proxy",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:   "router",
				Usage:  "Set the router to use e.g default, go.micro.router",
				EnvVar: "MICRO_ROUTER",
			},
			cli.StringFlag{
				Name:   "router_address",
				Usage:  "Set the router address",
				EnvVar: "MICRO_ROUTER_ADDRESS",
			},
			cli.StringFlag{
				Name:   "address",
				Usage:  "Set the proxy http address e.g 0.0.0.0:8081",
				EnvVar: "MICRO_PROXY_ADDRESS",
			},
			cli.StringFlag{
				Name:   "protocol",
				Usage:  "Set the protocol used for proxying e.g mucp, grpc, http",
				EnvVar: "MICRO_PROXY_PROTOCOL",
			},
			cli.StringFlag{
				Name:   "endpoint",
				Usage:  "Set the endpoint to route to e.g greeter or localhost:9090",
				EnvVar: "MICRO_PROXY_ENDPOINT",
			},
		},
		Action: func(ctx *cli.Context) {
			run(ctx, options...)
		},
	}

	for _, p := range Plugins() {
		if cmds := p.Commands(); len(cmds) > 0 {
			command.Subcommands = append(command.Subcommands, cmds...)
		}

		if flags := p.Flags(); len(flags) > 0 {
			command.Flags = append(command.Flags, flags...)
		}
	}

	return []cli.Command{command}
}
<commit_msg>make proxy respect protocol flag (#444)<commit_after>// Package proxy is a cli proxy
package proxy

import (
	"os"
	"strings"
	"time"

	"github.com/micro/cli"
	"github.com/micro/go-micro"
	bmem "github.com/micro/go-micro/broker/memory"
	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/config/options"
	"github.com/micro/go-micro/proxy"
	"github.com/micro/go-micro/proxy/grpc"
	"github.com/micro/go-micro/proxy/http"
	"github.com/micro/go-micro/proxy/mucp"
	"github.com/micro/go-micro/registry"
	rmem "github.com/micro/go-micro/registry/memory"
	"github.com/micro/go-micro/router"
	rs "github.com/micro/go-micro/router/service"
	"github.com/micro/go-micro/server"
	sgrpc "github.com/micro/go-micro/server/grpc"
	smucp "github.com/micro/go-micro/server/mucp"
	"github.com/micro/go-micro/util/log"
	"github.com/micro/go-micro/util/mux"
)

var (
	// Name of the proxy
	Name = "go.micro.proxy"
	// The address of the proxy
	Address = ":8081"
	// the proxy protocol
	Protocol = "mucp"
	// The endpoint host to route to
	Endpoint string
)

func run(ctx *cli.Context, srvOpts ...micro.Option) {
	log.Name("proxy")

	if len(ctx.GlobalString("server_name")) > 0 {
		Name = ctx.GlobalString("server_name")
	}
	if len(ctx.String("address")) > 0 {
		Address = ctx.String("address")
	}
	if len(ctx.String("endpoint")) > 0 {
		Endpoint = ctx.String("endpoint")
	}
	if len(ctx.String("protocol")) > 0 {
		Protocol = ctx.String("protocol")
	}

	// Init plugins
	for _, p := range Plugins() {
		p.Init(ctx)
	}

	// service opts
	srvOpts = append(srvOpts, micro.Name(Name))
	if i := time.Duration(ctx.GlobalInt("register_ttl")); i > 0 {
		srvOpts = append(srvOpts, micro.RegisterTTL(i*time.Second))
	}
	if i := time.Duration(ctx.GlobalInt("register_interval")); i > 0 {
		srvOpts = append(srvOpts, micro.RegisterInterval(i*time.Second))
	}

	// set the context
	var popts []options.Option

	// create new router
	var r router.Router

	routerName := ctx.String("router")
	routerAddr := ctx.String("router_address")

	ropts := []router.Option{
		router.Id(server.DefaultId),
		router.Client(client.DefaultClient),
		router.Address(routerAddr),
		router.Registry(registry.DefaultRegistry),
	}

	// check if we need to use the router service
	switch {
	case routerName == "go.micro.router":
		r = rs.NewRouter(ropts...)
	case len(routerAddr) > 0:
		r = rs.NewRouter(ropts...)
	default:
		r = router.NewRouter(ropts...)
	}

	// start the router
	if err := r.Start(); err != nil {
		log.Logf("Proxy error starting router: %s", err)
		os.Exit(1)
	}

	popts = append(popts, proxy.WithRouter(r))

	// new proxy
	var p proxy.Proxy
	var s server.Server
	var srv server.Server

	// set endpoint
	if len(Endpoint) > 0 {
		switch {
		case strings.HasPrefix(Endpoint, "grpc://"):
			ep := strings.TrimPrefix(Endpoint, "grpc://")
			popts = append(popts, proxy.WithEndpoint(ep))
			p = grpc.NewProxy(popts...)
		case strings.HasPrefix(Endpoint, "http://"):
			// TODO: strip prefix?
			popts = append(popts, proxy.WithEndpoint(Endpoint))
			p = http.NewProxy(popts...)
		default:
			// TODO: strip prefix?
			popts = append(popts, proxy.WithEndpoint(Endpoint))
			p = mucp.NewProxy(popts...)
		}
	}

	// set based on protocol
	if p == nil && len(Protocol) > 0 {
		switch Protocol {
		case "http":
			p = http.NewProxy(popts...)
			// TODO: http server
		case "grpc":
			p = grpc.NewProxy(popts...)
			s = sgrpc.NewServer()
			srv = sgrpc.NewServer(
				server.Address(Address),
				// reset registry to memory
				server.Registry(rmem.NewRegistry()),
				// reset broker to memory
				server.Broker(bmem.NewBroker()),
				// hande it the router
				server.WithRouter(p),
			)
		default:
			p = mucp.NewProxy(popts...)
			s = smucp.NewServer()
			srv = server.NewServer(
				server.Address(Address),
				// reset registry to memory
				server.Registry(rmem.NewRegistry()),
				// reset broker to memory
				server.Broker(bmem.NewBroker()),
				// hande it the router
				server.WithRouter(p),
			)
		}
	}

	if len(Endpoint) > 0 {
		log.Logf("Proxy [%s] serving endpoint: %s", p.String(), Endpoint)
	} else {
		log.Logf("Proxy [%s] serving protocol: %s", p.String(), Protocol)
	}

	// prepend the server
	if s != nil {
		srvOpts = append([]micro.Option{micro.Server(s)}, srvOpts...)
	}

	// new service
	service := micro.NewService(srvOpts...)

	// create a new proxy muxer which includes the debug handler
	muxer := mux.New(Name, p)

	// set the router
	service.Server().Init(
		server.WithRouter(muxer),
	)

	// Start the server
	if err := srv.Start(); err != nil {
		log.Fatal(err)
	}

	// Run internal service
	if err := service.Run(); err != nil {
		log.Fatal(err)
	}

	// Stop the server
	if err := srv.Stop(); err != nil {
		log.Fatal(err)
	}
}

func Commands(options ...micro.Option) []cli.Command {
	command := cli.Command{
		Name:  "proxy",
		Usage: "Run the service proxy",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:   "router",
				Usage:  "Set the router to use e.g default, go.micro.router",
				EnvVar: "MICRO_ROUTER",
			},
			cli.StringFlag{
				Name:   "router_address",
				Usage:  "Set the router address",
				EnvVar: "MICRO_ROUTER_ADDRESS",
			},
			cli.StringFlag{
				Name:   "address",
				Usage:  "Set the proxy http address e.g 0.0.0.0:8081",
				EnvVar: "MICRO_PROXY_ADDRESS",
			},
			cli.StringFlag{
				Name:   "protocol",
				Usage:  "Set the protocol used for proxying e.g mucp, grpc, http",
				EnvVar: "MICRO_PROXY_PROTOCOL",
			},
			cli.StringFlag{
				Name:   "endpoint",
				Usage:  "Set the endpoint to route to e.g greeter or localhost:9090",
				EnvVar: "MICRO_PROXY_ENDPOINT",
			},
		},
		Action: func(ctx *cli.Context) {
			run(ctx, options...)
		},
	}

	for _, p := range Plugins() {
		if cmds := p.Commands(); len(cmds) > 0 {
			command.Subcommands = append(command.Subcommands, cmds...)
		}

		if flags := p.Flags(); len(flags) > 0 {
			command.Flags = append(command.Flags, flags...)
		}
	}

	return []cli.Command{command}
}
<|endoftext|>