<commit_before>package dockeripv6nat

import (
	"errors"
	"net"
	"strconv"

	"github.com/coreos/go-iptables/iptables"
)

type managedNetwork struct {
	id         string
	bridge     string
	subnet     net.IPNet
	icc        bool
	masquerade bool
	internal   bool
	binding    net.IP
}

type managedContainer struct {
	id      string
	bridge  string
	address net.IP
	ports   []managedPort
}

type managedPort struct {
	port        uint16
	proto       string
	hostAddress net.IP
	hostPort    uint16
}

type manager struct {
	fw          *firewall
	hairpinMode bool
}

func NewManager() (*manager, error) {
	hairpinMode, err := detectHairpinMode()
	if err != nil {
		return nil, err
	}

	fw, err := NewFirewall()
	if err != nil {
		return nil, err
	}

	if err := fw.EnsureTableChains(getCustomTableChains()); err != nil {
		return nil, err
	}

	if err := fw.EnsureRules(getBaseRules(hairpinMode)); err != nil {
		return nil, err
	}

	return &manager{
		fw:          fw,
		hairpinMode: hairpinMode,
	}, nil
}

func detectHairpinMode() (bool, error) {
	// Use the IPv4 firewall to detect if the docker daemon is started with --userland-proxy=false.

	ipt, err := iptables.NewWithProtocol(iptables.ProtocolIPv4)
	if err != nil {
		return false, err
	}

	hairpinModeOffRulespec := []string{
		"!", "-d", "127.0.0.0/8",
		"-m", "addrtype",
		"--dst-type", "LOCAL",
		"-j", "DOCKER",
	}

	hairpinModeOnRulespec := hairpinModeOffRulespec[3:]

	hairpinModeOn, err := ipt.Exists(TableNat, ChainOutput, hairpinModeOnRulespec...)
	if err != nil {
		return false, err
	} else if hairpinModeOn {
		return true, nil
	}

	hairpinModeOff, err := ipt.Exists(TableNat, ChainOutput, hairpinModeOffRulespec...)
	if err != nil {
		return false, err
	} else if hairpinModeOff {
		return false, nil
	}

	return false, errors.New("unable to detect hairpin mode (is the docker daemon running?)")
}

func (m *manager) Cleanup() error {
	if err := m.fw.RemoveRules(getBaseRules(m.hairpinMode)); err != nil {
		return err
	}

	if err := m.fw.RemoveTableChains(getCustomTableChains()); err != nil {
		return err
	}

	return nil
}

func (m *manager) ReplaceNetwork(oldNetwork, newNetwork *managedNetwork) error {
	return m.applyRules(getRulesForNetwork(oldNetwork, m.hairpinMode), getRulesForNetwork(newNetwork, m.hairpinMode))
}

func (m *manager) ReplaceContainer(oldContainer, newContainer *managedContainer) error {
	return m.applyRules(getRulesForContainer(oldContainer, m.hairpinMode), getRulesForContainer(newContainer, m.hairpinMode))
}

func (m *manager) EnsureInterconnectionRules(network *managedNetwork, otherNetworks []*managedNetwork) error {
	return m.fw.EnsureRules(getInterconnectionRules(network, otherNetworks))
}

func (m *manager) RemoveInterconnectionRules(network *managedNetwork, otherNetworks []*managedNetwork) error {
	return m.fw.RemoveRules(getInterconnectionRules(network, otherNetworks))
}

func (m *manager) applyRules(oldRules, newRules *Ruleset) error {
	oldRules = oldRules.Diff(newRules)

	if err := m.fw.EnsureRules(newRules); err != nil {
		return err
	}

	if err := m.fw.RemoveRules(oldRules); err != nil {
		return err
	}

	return nil
}

func getCustomTableChains() []TableChain {
	return []TableChain{
		TableChain{TableFilter, ChainDocker},
		TableChain{TableFilter, ChainDockerIsolation},
		TableChain{TableNat, ChainDocker},
	}
}

func getBaseRules(hairpinMode bool) *Ruleset {
	outputRule := NewRule(TableNat, ChainOutput,
		"-m", "addrtype",
		"--dst-type", "LOCAL",
		"-j", ChainDocker)

	if !hairpinMode {
		outputRule.spec = append(outputRule.spec, "!", "-d", "::1")
	}

	return &Ruleset{
		NewPrependRule(TableFilter, ChainForward,
			"-j", ChainDockerIsolation),
		NewRule(TableFilter, ChainDockerIsolation,
			"-j", "RETURN"),
		NewRule(TableNat, ChainPrerouting,
			"-m", "addrtype",
			"--dst-type", "LOCAL",
			"-j", ChainDocker),
		outputRule,
	}
}

func getRulesForNetwork(network *managedNetwork, hairpinMode bool) *Ruleset {
	if network == nil {
		return &Ruleset{}
	}

	if network.internal {
		return &Ruleset{
			NewPrependRule(TableFilter, ChainDockerIsolation,
				"!", "-s", network.subnet.String(),
				"-o", network.bridge,
				"-j", "DROP"),
			NewPrependRule(TableFilter, ChainDockerIsolation,
				"!", "-d", network.subnet.String(),
				"-i", network.bridge,
				"-j", "DROP"),
		}
	}

	iccAction := "ACCEPT"
	if !network.icc {
		iccAction = "DROP"
	}

	rs := Ruleset{
		NewRule(TableFilter, ChainForward,
			"-o", network.bridge,
			"-j", ChainDocker),
		NewRule(TableFilter, ChainForward,
			"-o", network.bridge,
			"-m", "conntrack",
			"--ctstate", "RELATED,ESTABLISHED",
			"-j", "ACCEPT"),
		NewRule(TableFilter, ChainForward,
			"-i", network.bridge,
			"!", "-o", network.bridge,
			"-j", "ACCEPT"),
		NewRule(TableFilter, ChainForward,
			"-i", network.bridge,
			"-o", network.bridge,
			"-j", iccAction),
	}

	if network.masquerade {
		rs = append(rs, NewPrependRule(TableNat, ChainPostrouting,
			"-s", network.subnet.String(),
			"!", "-o", network.bridge,
			"-j", "MASQUERADE"))
	}

	if !hairpinMode {
		rs = append(rs, NewPrependRule(TableNat, ChainDocker,
			"-i", network.bridge,
			"-j", "RETURN"))
	}

	return &rs
}

func getRulesForContainer(container *managedContainer, hairpinMode bool) *Ruleset {
	if container == nil {
		return &Ruleset{}
	}

	rs := make(Ruleset, 0, len(container.ports)*3)
	for _, port := range container.ports {
		rs = append(rs, *getRulesForPort(&port, container, hairpinMode)...)
	}

	return &rs
}

func getRulesForPort(port *managedPort, container *managedContainer, hairpinMode bool) *Ruleset {
	containerPortString := strconv.Itoa(int(port.port))
	hostPortString := strconv.Itoa(int(port.hostPort))
	hostAddressString := "0/0"
	if !port.hostAddress.IsUnspecified() {
		hostAddressString = port.hostAddress.String()
	}

	dnatRule := NewRule(TableNat, ChainDocker,
		"-d", hostAddressString,
		"-p", port.proto,
		"-m", port.proto,
		"--dport", hostPortString,
		"-j", "DNAT",
		"--to-destination", net.JoinHostPort(container.address.String(), containerPortString))

	if !hairpinMode {
		dnatRule.spec = append(dnatRule.spec, "!", "-i", container.bridge)
	}

	return &Ruleset{
		NewRule(TableFilter, ChainDocker,
			"-d", container.address.String(),
			"!", "-i", container.bridge,
			"-o", container.bridge,
			"-p", port.proto,
			"-m", port.proto,
			"--dport", containerPortString,
			"-j", "ACCEPT"),
		NewRule(TableNat, ChainPostrouting,
			"-s", container.address.String(),
			"-d", container.address.String(),
			"-p", port.proto,
			"-m", port.proto,
			"--dport", containerPortString,
			"-j", "MASQUERADE"),
		dnatRule,
	}
}

func getInterconnectionRules(network *managedNetwork, otherNetworks []*managedNetwork) *Ruleset {
	if network.internal {
		return &Ruleset{}
	}

	rs := make(Ruleset, 0, len(otherNetworks)*2)
	for _, otherNetwork := range otherNetworks {
		if otherNetwork.id == network.id {
			continue
		}

		if otherNetwork.internal {
			continue
		}

		rs = append(rs, NewPrependRule(TableFilter, ChainDockerIsolation,
			"-i", network.bridge,
			"-o", otherNetwork.bridge,
			"-j", "DROP"))
		rs = append(rs, NewPrependRule(TableFilter, ChainDockerIsolation,
			"-i", otherNetwork.bridge,
			"-o", network.bridge,
			"-j", "DROP"))
	}

	return &rs
}
<commit_msg>Check ICC flag before creating Ruleset for internal networks<commit_after>package dockeripv6nat

import (
	"errors"
	"net"
	"strconv"

	"github.com/coreos/go-iptables/iptables"
)

type managedNetwork struct {
	id         string
	bridge     string
	subnet     net.IPNet
	icc        bool
	masquerade bool
	internal   bool
	binding    net.IP
}

type managedContainer struct {
	id      string
	bridge  string
	address net.IP
	ports   []managedPort
}

type managedPort struct {
	port        uint16
	proto       string
	hostAddress net.IP
	hostPort    uint16
}

type manager struct {
	fw          *firewall
	hairpinMode bool
}

func NewManager() (*manager, error) {
	hairpinMode, err := detectHairpinMode()
	if err != nil {
		return nil, err
	}

	fw, err := NewFirewall()
	if err != nil {
		return nil, err
	}

	if err := fw.EnsureTableChains(getCustomTableChains()); err != nil {
		return nil, err
	}

	if err := fw.EnsureRules(getBaseRules(hairpinMode)); err != nil {
		return nil, err
	}

	return &manager{
		fw:          fw,
		hairpinMode: hairpinMode,
	}, nil
}

func detectHairpinMode() (bool, error) {
	// Use the IPv4 firewall to detect if the docker daemon is started with --userland-proxy=false.

	ipt, err := iptables.NewWithProtocol(iptables.ProtocolIPv4)
	if err != nil {
		return false, err
	}

	hairpinModeOffRulespec := []string{
		"!", "-d", "127.0.0.0/8",
		"-m", "addrtype",
		"--dst-type", "LOCAL",
		"-j", "DOCKER",
	}

	hairpinModeOnRulespec := hairpinModeOffRulespec[3:]

	hairpinModeOn, err := ipt.Exists(TableNat, ChainOutput, hairpinModeOnRulespec...)
	if err != nil {
		return false, err
	} else if hairpinModeOn {
		return true, nil
	}

	hairpinModeOff, err := ipt.Exists(TableNat, ChainOutput, hairpinModeOffRulespec...)
	if err != nil {
		return false, err
	} else if hairpinModeOff {
		return false, nil
	}

	return false, errors.New("unable to detect hairpin mode (is the docker daemon running?)")
}

func (m *manager) Cleanup() error {
	if err := m.fw.RemoveRules(getBaseRules(m.hairpinMode)); err != nil {
		return err
	}

	if err := m.fw.RemoveTableChains(getCustomTableChains()); err != nil {
		return err
	}

	return nil
}

func (m *manager) ReplaceNetwork(oldNetwork, newNetwork *managedNetwork) error {
	return m.applyRules(getRulesForNetwork(oldNetwork, m.hairpinMode), getRulesForNetwork(newNetwork, m.hairpinMode))
}

func (m *manager) ReplaceContainer(oldContainer, newContainer *managedContainer) error {
	return m.applyRules(getRulesForContainer(oldContainer, m.hairpinMode), getRulesForContainer(newContainer, m.hairpinMode))
}

func (m *manager) EnsureInterconnectionRules(network *managedNetwork, otherNetworks []*managedNetwork) error {
	return m.fw.EnsureRules(getInterconnectionRules(network, otherNetworks))
}

func (m *manager) RemoveInterconnectionRules(network *managedNetwork, otherNetworks []*managedNetwork) error {
	return m.fw.RemoveRules(getInterconnectionRules(network, otherNetworks))
}

func (m *manager) applyRules(oldRules, newRules *Ruleset) error {
	oldRules = oldRules.Diff(newRules)

	if err := m.fw.EnsureRules(newRules); err != nil {
		return err
	}

	if err := m.fw.RemoveRules(oldRules); err != nil {
		return err
	}

	return nil
}

func getCustomTableChains() []TableChain {
	return []TableChain{
		TableChain{TableFilter, ChainDocker},
		TableChain{TableFilter, ChainDockerIsolation},
		TableChain{TableNat, ChainDocker},
	}
}

func getBaseRules(hairpinMode bool) *Ruleset {
	outputRule := NewRule(TableNat, ChainOutput,
		"-m", "addrtype",
		"--dst-type", "LOCAL",
		"-j", ChainDocker)

	if !hairpinMode {
		outputRule.spec = append(outputRule.spec, "!", "-d", "::1")
	}

	return &Ruleset{
		NewPrependRule(TableFilter, ChainForward,
			"-j", ChainDockerIsolation),
		NewRule(TableFilter, ChainDockerIsolation,
			"-j", "RETURN"),
		NewRule(TableNat, ChainPrerouting,
			"-m", "addrtype",
			"--dst-type", "LOCAL",
			"-j", ChainDocker),
		outputRule,
	}
}

func getRulesForNetwork(network *managedNetwork, hairpinMode bool) *Ruleset {
	if network == nil {
		return &Ruleset{}
	}

	iccAction := "ACCEPT"
	if !network.icc {
		iccAction = "DROP"
	}

	if network.internal {
		return &Ruleset{
			NewPrependRule(TableFilter, ChainDockerIsolation,
				"!", "-s", network.subnet.String(),
				"-o", network.bridge,
				"-j", "DROP"),
			NewPrependRule(TableFilter, ChainDockerIsolation,
				"!", "-d", network.subnet.String(),
				"-i", network.bridge,
				"-j", "DROP"),
		}
	}

	rs := Ruleset{
		NewRule(TableFilter, ChainForward,
			"-o", network.bridge,
			"-j", ChainDocker),
		NewRule(TableFilter, ChainForward,
			"-o", network.bridge,
			"-m", "conntrack",
			"--ctstate", "RELATED,ESTABLISHED",
			"-j", "ACCEPT"),
		NewRule(TableFilter, ChainForward,
			"-i", network.bridge,
			"!", "-o", network.bridge,
			"-j", "ACCEPT"),
		NewRule(TableFilter, ChainForward,
			"-i", network.bridge,
			"-o", network.bridge,
			"-j", iccAction),
	}

	if network.masquerade {
		rs = append(rs, NewPrependRule(TableNat, ChainPostrouting,
			"-s", network.subnet.String(),
			"!", "-o", network.bridge,
			"-j", "MASQUERADE"))
	}

	if !hairpinMode {
		rs = append(rs, NewPrependRule(TableNat, ChainDocker,
			"-i", network.bridge,
			"-j", "RETURN"))
	}

	return &rs
}

func getRulesForContainer(container *managedContainer, hairpinMode bool) *Ruleset {
	if container == nil {
		return &Ruleset{}
	}

	rs := make(Ruleset, 0, len(container.ports)*3)
	for _, port := range container.ports {
		rs = append(rs, *getRulesForPort(&port, container, hairpinMode)...)
	}

	return &rs
}

func getRulesForPort(port *managedPort, container *managedContainer, hairpinMode bool) *Ruleset {
	containerPortString := strconv.Itoa(int(port.port))
	hostPortString := strconv.Itoa(int(port.hostPort))
	hostAddressString := "0/0"
	if !port.hostAddress.IsUnspecified() {
		hostAddressString = port.hostAddress.String()
	}

	dnatRule := NewRule(TableNat, ChainDocker,
		"-d", hostAddressString,
		"-p", port.proto,
		"-m", port.proto,
		"--dport", hostPortString,
		"-j", "DNAT",
		"--to-destination", net.JoinHostPort(container.address.String(), containerPortString))

	if !hairpinMode {
		dnatRule.spec = append(dnatRule.spec, "!", "-i", container.bridge)
	}

	return &Ruleset{
		NewRule(TableFilter, ChainDocker,
			"-d", container.address.String(),
			"!", "-i", container.bridge,
			"-o", container.bridge,
			"-p", port.proto,
			"-m", port.proto,
			"--dport", containerPortString,
			"-j", "ACCEPT"),
		NewRule(TableNat, ChainPostrouting,
			"-s", container.address.String(),
			"-d", container.address.String(),
			"-p", port.proto,
			"-m", port.proto,
			"--dport", containerPortString,
			"-j", "MASQUERADE"),
		dnatRule,
	}
}

func getInterconnectionRules(network *managedNetwork, otherNetworks []*managedNetwork) *Ruleset {
	if network.internal {
		return &Ruleset{}
	}

	rs := make(Ruleset, 0, len(otherNetworks)*2)
	for _, otherNetwork := range otherNetworks {
		if otherNetwork.id == network.id {
			continue
		}

		if otherNetwork.internal {
			continue
		}

		rs = append(rs, NewPrependRule(TableFilter, ChainDockerIsolation,
			"-i", network.bridge,
			"-o", otherNetwork.bridge,
			"-j", "DROP"))
		rs = append(rs, NewPrependRule(TableFilter, ChainDockerIsolation,
			"-i", otherNetwork.bridge,
			"-o", network.bridge,
			"-j", "DROP"))
	}

	return &rs
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2017 Intel Corporation
//
// 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 virtcontainers

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
)

const testPodID = "7f49d00d-1995-4156-8c79-5f5ab24ce138"
const testKernel = "kernel"
const testImage = "image"
const testHypervisor = "hypervisor"
const testBundle = "bundle"

const testDisabledAsNonRoot = "Test disabled as requires root privileges"

// package variables set in TestMain
var testDir = ""
var podDirConfig = ""
var podFileConfig = ""
var podDirState = ""
var podDirLock = ""
var podFileState = ""
var podFileLock = ""
var testQemuKernelPath = ""
var testQemuImagePath = ""
var testQemuPath = ""
var testHyperstartCtlSocket = ""
var testHyperstartTtySocket = ""

// cleanUp Removes any stale pod/container state that can affect
// the next test to run.
func cleanUp() {
	for _, dir := range []string{testDir, defaultSharedDir} {
		os.RemoveAll(dir)
		os.MkdirAll(dir, dirMode)
	}

	os.Mkdir(filepath.Join(testDir, testBundle), dirMode)

	_, err := os.Create(filepath.Join(testDir, testImage))
	if err != nil {
		fmt.Println("Could not recreate test image:", err)
		os.Exit(1)
	}
}

// TestMain is the common main function used by ALL the test functions
// for this package.
func TestMain(m *testing.M) {
	var err error

	flag.Parse()

	testDir, err = ioutil.TempDir("", "virtcontainers-tmp-")
	if err != nil {
		panic(err)
	}

	err = os.MkdirAll(testDir, dirMode)
	if err != nil {
		fmt.Println("Could not create test directories:", err)
		os.Exit(1)
	}

	_, err = os.Create(filepath.Join(testDir, testKernel))
	if err != nil {
		fmt.Println("Could not create test kernel:", err)
		os.RemoveAll(testDir)
		os.Exit(1)
	}

	_, err = os.Create(filepath.Join(testDir, testImage))
	if err != nil {
		fmt.Println("Could not create test image:", err)
		os.RemoveAll(testDir)
		os.Exit(1)
	}

	_, err = os.Create(filepath.Join(testDir, testHypervisor))
	if err != nil {
		fmt.Println("Could not create test hypervisor:", err)
		os.RemoveAll(testDir)
		os.Exit(1)
	}

	err = os.Mkdir(filepath.Join(testDir, testBundle), dirMode)
	if err != nil {
		fmt.Println("Could not create test bundle directory:", err)
		os.RemoveAll(testDir)
		os.Exit(1)
	}

	// allow the tests to run without affecting the host system.
	configStoragePath = filepath.Join(testDir, storagePathSuffix, "config")
	runStoragePath = filepath.Join(testDir, storagePathSuffix, "run")

	// set now that configStoragePath has been overridden.
	podDirConfig = filepath.Join(configStoragePath, testPodID)
	podFileConfig = filepath.Join(configStoragePath, testPodID, configFile)
	podDirState = filepath.Join(runStoragePath, testPodID)
	podDirLock = filepath.Join(runStoragePath, testPodID)
	podFileState = filepath.Join(runStoragePath, testPodID, stateFile)
	podFileLock = filepath.Join(runStoragePath, testPodID, lockFileName)

	testQemuKernelPath = filepath.Join(testDir, testKernel)
	testQemuImagePath = filepath.Join(testDir, testImage)
	testQemuPath = filepath.Join(testDir, testHypervisor)

	testHyperstartCtlSocket = filepath.Join(testDir, "test_hyper.sock")
	testHyperstartTtySocket = filepath.Join(testDir, "test_tty.sock")

	ret := m.Run()

	os.RemoveAll(testDir)

	os.Exit(ret)
}
<commit_msg>test: Add custom flag to enable debug logs for benchmarks<commit_after>//
// Copyright (c) 2017 Intel Corporation
//
// 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 virtcontainers

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/Sirupsen/logrus"
)

const testPodID = "7f49d00d-1995-4156-8c79-5f5ab24ce138"
const testKernel = "kernel"
const testImage = "image"
const testHypervisor = "hypervisor"
const testBundle = "bundle"

const testDisabledAsNonRoot = "Test disabled as requires root privileges"

// package variables set in TestMain
var testDir = ""
var podDirConfig = ""
var podFileConfig = ""
var podDirState = ""
var podDirLock = ""
var podFileState = ""
var podFileLock = ""
var testQemuKernelPath = ""
var testQemuImagePath = ""
var testQemuPath = ""
var testHyperstartCtlSocket = ""
var testHyperstartTtySocket = ""

// cleanUp Removes any stale pod/container state that can affect
// the next test to run.
func cleanUp() {
	for _, dir := range []string{testDir, defaultSharedDir} {
		os.RemoveAll(dir)
		os.MkdirAll(dir, dirMode)
	}

	os.Mkdir(filepath.Join(testDir, testBundle), dirMode)

	_, err := os.Create(filepath.Join(testDir, testImage))
	if err != nil {
		fmt.Println("Could not recreate test image:", err)
		os.Exit(1)
	}
}

// TestMain is the common main function used by ALL the test functions
// for this package.
func TestMain(m *testing.M) {
	var err error

	flag.Parse()

	virtLog.Level = logrus.ErrorLevel
	for _, arg := range flag.Args() {
		if arg == "debug-logs" {
			virtLog.Level = logrus.DebugLevel
		}
	}

	testDir, err = ioutil.TempDir("", "virtcontainers-tmp-")
	if err != nil {
		panic(err)
	}

	err = os.MkdirAll(testDir, dirMode)
	if err != nil {
		fmt.Println("Could not create test directories:", err)
		os.Exit(1)
	}

	_, err = os.Create(filepath.Join(testDir, testKernel))
	if err != nil {
		fmt.Println("Could not create test kernel:", err)
		os.RemoveAll(testDir)
		os.Exit(1)
	}

	_, err = os.Create(filepath.Join(testDir, testImage))
	if err != nil {
		fmt.Println("Could not create test image:", err)
		os.RemoveAll(testDir)
		os.Exit(1)
	}

	_, err = os.Create(filepath.Join(testDir, testHypervisor))
	if err != nil {
		fmt.Println("Could not create test hypervisor:", err)
		os.RemoveAll(testDir)
		os.Exit(1)
	}

	err = os.Mkdir(filepath.Join(testDir, testBundle), dirMode)
	if err != nil {
		fmt.Println("Could not create test bundle directory:", err)
		os.RemoveAll(testDir)
		os.Exit(1)
	}

	// allow the tests to run without affecting the host system.
	configStoragePath = filepath.Join(testDir, storagePathSuffix, "config")
	runStoragePath = filepath.Join(testDir, storagePathSuffix, "run")

	// set now that configStoragePath has been overridden.
	podDirConfig = filepath.Join(configStoragePath, testPodID)
	podFileConfig = filepath.Join(configStoragePath, testPodID, configFile)
	podDirState = filepath.Join(runStoragePath, testPodID)
	podDirLock = filepath.Join(runStoragePath, testPodID)
	podFileState = filepath.Join(runStoragePath, testPodID, stateFile)
	podFileLock = filepath.Join(runStoragePath, testPodID, lockFileName)

	testQemuKernelPath = filepath.Join(testDir, testKernel)
	testQemuImagePath = filepath.Join(testDir, testImage)
	testQemuPath = filepath.Join(testDir, testHypervisor)

	testHyperstartCtlSocket = filepath.Join(testDir, "test_hyper.sock")
	testHyperstartTtySocket = filepath.Join(testDir, "test_tty.sock")

	ret := m.Run()

	os.RemoveAll(testDir)

	os.Exit(ret)
}
<|endoftext|>


<commit_before>package stellarbase

import (
	"errors"
	"fmt"
	"math/big"
	"strconv"
	"strings"
)

const alphabet = "gsphnaf39wBUDNEGHJKLM4PQRST7VWXYZ2bcdeCr65jkm8oFqi1tuvAxyz"

var decodeMap [256]byte

func init() {
	for i := 0; i < len(decodeMap); i++ {
		decodeMap[i] = 0xFF
	}
	for i := 0; i < len(alphabet); i++ {
		decodeMap[alphabet[i]] = byte(i)
	}
}

var (
	NotCheckEncodedError = errors.New("base58: input is not check encoded")
)

type CorruptInputError int64

func (e CorruptInputError) Error() string {
	return "illegal base58 data at input byte " + strconv.FormatInt(int64(e), 10)
}

type InvalidVersionByteError struct {
	Expected VersionByte
	Actual   VersionByte
}

func (e InvalidVersionByteError) Error() string {
	return fmt.Sprintf("illegal base58 version byte expected:%d actual:%d", e.Expected, e.Actual)
}

type VersionByte byte

const (
	VersionByteAccountID VersionByte = 0
	VersionByteNone                  = 1
	VersionByteSeed                  = 33
)

func EncodeBase58(src []byte) string {
	bigInt := new(big.Int)
	bigInt.SetBytes(src)
	leadingZeroes := strings.Repeat("g", leadingZeroCount(src))

	var resultSlice []byte = make([]byte, 0, 256)
	var results = []string{
		leadingZeroes,
		string(EncodeBigToBase58(resultSlice, bigInt)),
	}

	return strings.Join(results, "")
}

func EncodeBase58Check(version VersionByte, src []byte) string {
	start := []byte{}
	withVersion := append(start, byte(version))
	withPayload := append(withVersion, src...)
	withChecksum := append(withPayload, base58CheckSum(withPayload)...)
	// _ = withChecksum
	// return fmt.Sprintf("b: %s\n", withChecksum)
	return string(EncodeBase58(withChecksum))
}

func DecodeBase58(src string) ([]byte, error) {
	leadingGs := make([]byte, leadingGCount(src))

	bigInt, err := DecodeBase58ToBig([]byte(src))

	if err != nil {
		return []byte{}, err
	}

	return append(leadingGs, bigInt.Bytes()...), nil
}

func DecodeBase58Check(version VersionByte, src string) ([]byte, error) {

	decoded, err := DecodeBase58(src)

	if err != nil {
		return []byte{}, err
	}

	if len(decoded) < 5 {
		return []byte{}, NotCheckEncodedError
	}

	decodedVersion := VersionByte(decoded[0])
	payload := decoded[1 : len(decoded)-4]
	checksum := decoded[len(decoded)-4:]

	if decodedVersion != version {
		return []byte{}, InvalidVersionByteError{version, decodedVersion}
	}

	_ = checksum

	return payload, nil
}

// Decode a big integer from the bytes. Returns an error on corrupt
// input.
func DecodeBase58ToBig(src []byte) (*big.Int, error) {
	n := new(big.Int)
	radix := big.NewInt(58)
	for i := 0; i < len(src); i++ {
		b := decodeMap[src[i]]
		if b == 0xFF {
			return nil, CorruptInputError(i)
		}
		n.Mul(n, radix)
		n.Add(n, big.NewInt(int64(b)))
	}
	return n, nil
}

// Encode encodes src, appending to dst. Be sure to use the returned
// new value of dst.
func EncodeBigToBase58(dst []byte, src *big.Int) []byte {
	start := len(dst)
	n := new(big.Int)
	n.Set(src)
	radix := big.NewInt(58)
	zero := big.NewInt(0)

	for n.Cmp(zero) > 0 {
		mod := new(big.Int)
		n.DivMod(n, radix, mod)
		dst = append(dst, alphabet[mod.Int64()])
	}

	// reverse string
	for i, j := start, len(dst)-1; i < j; i, j = i+1, j-1 {
		dst[i], dst[j] = dst[j], dst[i]
	}
	return dst
}

func base58CheckSum(message []byte) []byte {
	inner := Hash(message)
	outer := Hash(inner[:])
	return outer[0:4]
}

func leadingZeroCount(src []byte) (result int) {
	for _, val := range src {
		if val != 0x00 {
			return
		}
		result++
	}
	return
}

func leadingGCount(src string) (result int) {
	for _, val := range src {
		if val != 'g' {
			return
		}
		result++
	}
	return
}
<commit_msg>Add todo note in base58<commit_after>package stellarbase

import (
	"errors"
	"fmt"
	"math/big"
	"strconv"
	"strings"
)

const alphabet = "gsphnaf39wBUDNEGHJKLM4PQRST7VWXYZ2bcdeCr65jkm8oFqi1tuvAxyz"

var decodeMap [256]byte

func init() {
	for i := 0; i < len(decodeMap); i++ {
		decodeMap[i] = 0xFF
	}
	for i := 0; i < len(alphabet); i++ {
		decodeMap[alphabet[i]] = byte(i)
	}
}

var (
	NotCheckEncodedError = errors.New("base58: input is not check encoded")
)

type CorruptInputError int64

func (e CorruptInputError) Error() string {
	return "illegal base58 data at input byte " + strconv.FormatInt(int64(e), 10)
}

type InvalidVersionByteError struct {
	Expected VersionByte
	Actual   VersionByte
}

func (e InvalidVersionByteError) Error() string {
	return fmt.Sprintf("illegal base58 version byte expected:%d actual:%d", e.Expected, e.Actual)
}

type VersionByte byte

const (
	VersionByteAccountID VersionByte = 0
	VersionByteNone                  = 1
	VersionByteSeed                  = 33
)

func EncodeBase58(src []byte) string {
	bigInt := new(big.Int)
	bigInt.SetBytes(src)
	leadingZeroes := strings.Repeat("g", leadingZeroCount(src))

	var resultSlice []byte = make([]byte, 0, 256)
	var results = []string{
		leadingZeroes,
		string(EncodeBigToBase58(resultSlice, bigInt)),
	}

	return strings.Join(results, "")
}

func EncodeBase58Check(version VersionByte, src []byte) string {
	start := []byte{}
	withVersion := append(start, byte(version))
	withPayload := append(withVersion, src...)
	withChecksum := append(withPayload, base58CheckSum(withPayload)...)
	// _ = withChecksum
	// return fmt.Sprintf("b: %s\n", withChecksum)
	return string(EncodeBase58(withChecksum))
}

func DecodeBase58(src string) ([]byte, error) {
	leadingGs := make([]byte, leadingGCount(src))

	bigInt, err := DecodeBase58ToBig([]byte(src))

	if err != nil {
		return []byte{}, err
	}

	return append(leadingGs, bigInt.Bytes()...), nil
}

func DecodeBase58Check(version VersionByte, src string) ([]byte, error) {

	decoded, err := DecodeBase58(src)

	if err != nil {
		return []byte{}, err
	}

	if len(decoded) < 5 {
		return []byte{}, NotCheckEncodedError
	}

	decodedVersion := VersionByte(decoded[0])
	payload := decoded[1 : len(decoded)-4]
	checksum := decoded[len(decoded)-4:]

	if decodedVersion != version {
		return []byte{}, InvalidVersionByteError{version, decodedVersion}
	}

	_ = checksum //TODO

	return payload, nil
}

// Decode a big integer from the bytes. Returns an error on corrupt
// input.
func DecodeBase58ToBig(src []byte) (*big.Int, error) {
	n := new(big.Int)
	radix := big.NewInt(58)
	for i := 0; i < len(src); i++ {
		b := decodeMap[src[i]]
		if b == 0xFF {
			return nil, CorruptInputError(i)
		}
		n.Mul(n, radix)
		n.Add(n, big.NewInt(int64(b)))
	}
	return n, nil
}

// Encode encodes src, appending to dst. Be sure to use the returned
// new value of dst.
func EncodeBigToBase58(dst []byte, src *big.Int) []byte {
	start := len(dst)
	n := new(big.Int)
	n.Set(src)
	radix := big.NewInt(58)
	zero := big.NewInt(0)

	for n.Cmp(zero) > 0 {
		mod := new(big.Int)
		n.DivMod(n, radix, mod)
		dst = append(dst, alphabet[mod.Int64()])
	}

	// reverse string
	for i, j := start, len(dst)-1; i < j; i, j = i+1, j-1 {
		dst[i], dst[j] = dst[j], dst[i]
	}
	return dst
}

func base58CheckSum(message []byte) []byte {
	inner := Hash(message)
	outer := Hash(inner[:])
	return outer[0:4]
}

func leadingZeroCount(src []byte) (result int) {
	for _, val := range src {
		if val != 0x00 {
			return
		}
		result++
	}
	return
}

func leadingGCount(src string) (result int) {
	for _, val := range src {
		if val != 'g' {
			return
		}
		result++
	}
	return
}
<|endoftext|>


<commit_before>package docker

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path"
	"strings"
	"time"
)

type Builder struct {
	runtime      *Runtime
	repositories *TagStore
	graph        *Graph
}

func NewBuilder(runtime *Runtime) *Builder {
	return &Builder{
		runtime:      runtime,
		graph:        runtime.graph,
		repositories: runtime.repositories,
	}
}

func (builder *Builder) mergeConfig(userConf, imageConf *Config) {
	if userConf.Hostname != "" {
		userConf.Hostname = imageConf.Hostname
	}
	if userConf.User != "" {
		userConf.User = imageConf.User
	}
	if userConf.Memory == 0 {
		userConf.Memory = imageConf.Memory
	}
	if userConf.MemorySwap == 0 {
		userConf.MemorySwap = imageConf.MemorySwap
	}
	if userConf.PortSpecs == nil || len(userConf.PortSpecs) == 0 {
		userConf.PortSpecs = imageConf.PortSpecs
	}
	if !userConf.Tty {
		userConf.Tty = userConf.Tty
	}
	if !userConf.OpenStdin {
		userConf.OpenStdin = imageConf.OpenStdin
	}
	if !userConf.StdinOnce {
		userConf.StdinOnce = imageConf.StdinOnce
	}
	if userConf.Env == nil || len(userConf.Env) == 0 {
		userConf.Env = imageConf.Env
	}
	if userConf.Cmd == nil || len(userConf.Cmd) == 0 {
		userConf.Cmd = imageConf.Cmd
	}
	if userConf.Dns == nil || len(userConf.Dns) == 0 {
		userConf.Dns = imageConf.Dns
	}
}

func (builder *Builder) Create(config *Config) (*Container, error) {
	// Lookup image
	img, err := builder.repositories.LookupImage(config.Image)
	if err != nil {
		return nil, err
	}

	if img.Config != nil {
		builder.mergeConfig(config, img.Config)
	}

	if config.Cmd == nil {
		return nil, fmt.Errorf("No command specified")
	}

	// Generate id
	id := GenerateId()
	// Generate default hostname
	// FIXME: the lxc template no longer needs to set a default hostname
	if config.Hostname == "" {
		config.Hostname = id[:12]
	}

	container := &Container{
		// FIXME: we should generate the ID here instead of receiving it as an argument
		Id:              id,
		Created:         time.Now(),
		Path:            config.Cmd[0],
		Args:            config.Cmd[1:], //FIXME: de-duplicate from config
		Config:          config,
		Image:           img.Id, // Always use the resolved image id
		NetworkSettings: &NetworkSettings{},
		// FIXME: do we need to store this in the container?
		SysInitPath: sysInitPath,
	}
	container.root = builder.runtime.containerRoot(container.Id)
	// Step 1: create the container directory.
	// This doubles as a barrier to avoid race conditions.
	if err := os.Mkdir(container.root, 0700); err != nil {
		return nil, err
	}

	// If custom dns exists, then create a resolv.conf for the container
	if len(config.Dns) > 0 {
		container.ResolvConfPath = path.Join(container.root, "resolv.conf")
		f, err := os.Create(container.ResolvConfPath)
		if err != nil {
			return nil, err
		}
		defer f.Close()
		for _, dns := range config.Dns {
			if _, err := f.Write([]byte("nameserver " + dns + "\n")); err != nil {
				return nil, err
			}
		}
	} else {
		container.ResolvConfPath = "/etc/resolv.conf"
	}

	// Step 2: save the container json
	if err := container.ToDisk(); err != nil {
		return nil, err
	}
	// Step 3: register the container
	if err := builder.runtime.Register(container); err != nil {
		return nil, err
	}
	return container, nil
}

// Commit creates a new filesystem image from the current state of a container.
// The image can optionally be tagged into a repository
func (builder *Builder) Commit(container *Container, repository, tag, comment, author string, config *Config) (*Image, error) {
	// FIXME: freeze the container before copying it to avoid data corruption?
	// FIXME: this shouldn't be in commands.
	rwTar, err := container.ExportRw()
	if err != nil {
		return nil, err
	}
	// Create a new image from the container's base layers + a new layer from container changes
	img, err := builder.graph.Create(rwTar, container, comment, author, config)
	if err != nil {
		return nil, err
	}
	// Register the image if needed
	if repository != "" {
		if err := builder.repositories.Set(repository, tag, img.Id, true); err != nil {
			return img, err
		}
	}
	return img, nil
}

func (builder *Builder) clearTmp(containers, images map[string]struct{}) {
	for c := range containers {
		tmp := builder.runtime.Get(c)
		builder.runtime.Destroy(tmp)
		Debugf("Removing container %s", c)
	}
	for i := range images {
		builder.runtime.graph.Delete(i)
		Debugf("Removing image %s", i)
	}
}

func (builder *Builder) getCachedImage(image *Image, config *Config) (*Image, error) {
	// Retrieve all images
	images, err := builder.graph.All()
	if err != nil {
		return nil, err
	}

	// Store the tree in a map of map (map[parentId][childId])
	imageMap := make(map[string]map[string]struct{})
	for _, img := range images {
		if _, exists := imageMap[img.Parent]; !exists {
			imageMap[img.Parent] = make(map[string]struct{})
		}
		imageMap[img.Parent][img.Id] = struct{}{}
	}

	// Loop on the children of the given image and check the config
	for elem := range imageMap[image.Id] {
		img, err := builder.graph.Get(elem)
		if err != nil {
			return nil, err
		}
		if CompareConfig(&img.ContainerConfig, config) {
			return img, nil
		}
	}
	return nil, nil
}

func (builder *Builder) Build(dockerfile io.Reader, stdout io.Writer) (*Image, error) {
	var (
		image, base   *Image
		config        *Config
		maintainer    string
		tmpContainers map[string]struct{} = make(map[string]struct{})
		tmpImages     map[string]struct{} = make(map[string]struct{})
	)
	defer builder.clearTmp(tmpContainers, tmpImages)

	file := bufio.NewReader(dockerfile)
	for {
		line, err := file.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, err
		}
		line = strings.Replace(strings.TrimSpace(line), "	", " ", 1)
		// Skip comments and empty line
		if len(line) == 0 || line[0] == '#' {
			continue
		}
		tmp := strings.SplitN(line, " ", 2)
		if len(tmp) != 2 {
			return nil, fmt.Errorf("Invalid Dockerfile format")
		}
		instruction := strings.Trim(tmp[0], " ")
		arguments := strings.Trim(tmp[1], " ")
		switch strings.ToLower(instruction) {
		case "from":
			fmt.Fprintf(stdout, "FROM %s\n", arguments)
			image, err = builder.runtime.repositories.LookupImage(arguments)
			if err != nil {
				if builder.runtime.graph.IsNotExist(err) {

					var tag, remote string
					if strings.Contains(arguments, ":") {
						remoteParts := strings.Split(arguments, ":")
						tag = remoteParts[1]
						remote = remoteParts[0]
					} else {
						remote = arguments
					}

					if err := builder.runtime.graph.PullRepository(stdout, remote, tag, builder.runtime.repositories, builder.runtime.authConfig); err != nil {
						return nil, err
					}

					image, err = builder.runtime.repositories.LookupImage(arguments)
					if err != nil {
						return nil, err
					}
				} else {
					return nil, err
				}
			}
			config = &Config{}

			break
		case "maintainer":
			fmt.Fprintf(stdout, "MAINTAINER %s\n", arguments)
			maintainer = arguments
			break
		case "run":
			fmt.Fprintf(stdout, "RUN %s\n", arguments)
			if image == nil {
				return nil, fmt.Errorf("Please provide a source image with `from` prior to run")
			}
			config, err := ParseRun([]string{image.Id, "/bin/sh", "-c", arguments}, nil, builder.runtime.capabilities)
			if err != nil {
				return nil, err
			}

			if cache, err := builder.getCachedImage(image, config); err != nil {
				return nil, err
			} else if cache != nil {
				image = cache
				fmt.Fprintf(stdout, "===> %s\n", image.ShortId())
				break
			}

			// Create the container and start it
			c, err := builder.Create(config)
			if err != nil {
				return nil, err
			}
			if err := c.Start(); err != nil {
				return nil, err
			}
			tmpContainers[c.Id] = struct{}{}

			// Wait for it to finish
			if result := c.Wait(); result != 0 {
				return nil, fmt.Errorf("!!! '%s' return non-zero exit code '%d'. Aborting.", arguments, result)
			}

			// Commit the container
			base, err = builder.Commit(c, "", "", "", maintainer, nil)
			if err != nil {
				return nil, err
			}
			tmpImages[base.Id] = struct{}{}

			fmt.Fprintf(stdout, "===> %s\n", base.ShortId())

			// use the base as the new image
			image = base

			break
		case "cmd":
			fmt.Fprintf(stdout, "CMD %s\n", arguments)

			// Create the container and start it
			c, err := builder.Create(&Config{Image: image.Id, Cmd: []string{"", ""}})
			if err != nil {
				return nil, err
			}
			if err := c.Start(); err != nil {
				return nil, err
			}
			tmpContainers[c.Id] = struct{}{}

			cmd := []string{}
			if err := json.Unmarshal([]byte(arguments), &cmd); err != nil {
				return nil, err
			}
			config.Cmd = cmd

			// Commit the container
			base, err = builder.Commit(c, "", "", "", maintainer, config)
			if err != nil {
				return nil, err
			}
			tmpImages[base.Id] = struct{}{}

			fmt.Fprintf(stdout, "===> %s\n", base.ShortId())
			image = base
			break
		case "expose":
			ports := strings.Split(arguments, " ")

			fmt.Fprintf(stdout, "EXPOSE %v\n", ports)
			if image == nil {
				return nil, fmt.Errorf("Please provide a source image with `from` prior to copy")
			}

			// Create the container and start it
			c, err := builder.Create(&Config{Image: image.Id, Cmd: []string{"", ""}})
			if err != nil {
				return nil, err
			}
			if err := c.Start(); err != nil {
				return nil, err
			}
			tmpContainers[c.Id] = struct{}{}

			config.PortSpecs = append(ports, config.PortSpecs...)

			// Commit the container
			base, err = builder.Commit(c, "", "", "", maintainer, config)
			if err != nil {
				return nil, err
			}
			tmpImages[base.Id] = struct{}{}

			fmt.Fprintf(stdout, "===> %s\n", base.ShortId())
			image = base
			break
		case "insert":
			if image == nil {
				return nil, fmt.Errorf("Please provide a source image with `from` prior to copy")
			}
			tmp = strings.SplitN(arguments, " ", 2)
			if len(tmp) != 2 {
				return nil, fmt.Errorf("Invalid INSERT format")
			}
			sourceUrl := strings.Trim(tmp[0], " ")
			destPath := strings.Trim(tmp[1], " ")
			fmt.Fprintf(stdout, "COPY %s to %s in %s\n", sourceUrl, destPath, base.ShortId())

			file, err := Download(sourceUrl, stdout)
			if err != nil {
				return nil, err
			}
			defer file.Body.Close()

			config, err := ParseRun([]string{base.Id, "echo", "insert", sourceUrl, destPath}, nil, builder.runtime.capabilities)
			if err != nil {
				return nil, err
			}
			c, err := builder.Create(config)
			if err != nil {
				return nil, err
			}

			if err := c.Start(); err != nil {
				return nil, err
			}

			// Wait for echo to finish
			if result := c.Wait(); result != 0 {
				return nil, fmt.Errorf("!!! '%s' return non-zero exit code '%d'. Aborting.", arguments, result)
			}

			if err := c.Inject(file.Body, destPath); err != nil {
				return nil, err
			}

			base, err = builder.Commit(c, "", "", "", maintainer, nil)
			if err != nil {
				return nil, err
			}
			fmt.Fprintf(stdout, "===> %s\n", base.ShortId())

			image = base

			break
		default:
			fmt.Fprintf(stdout, "Skipping unknown instruction %s\n", strings.ToUpper(instruction))
		}
	}
	if image != nil {
		// The build is successful, keep the temporary containers and images
		for i := range tmpImages {
			delete(tmpImages, i)
		}
		for i := range tmpContainers {
			delete(tmpContainers, i)
		}
		fmt.Fprintf(stdout, "Build finished. image id: %s\n", image.ShortId())
		return image, nil
	}
	return nil, fmt.Errorf("An error occured during the build\n")
}
<commit_msg>Implement ENV within docker builder<commit_after>package docker

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path"
	"strings"
	"time"
)

type Builder struct {
	runtime      *Runtime
	repositories *TagStore
	graph        *Graph
}

func NewBuilder(runtime *Runtime) *Builder {
	return &Builder{
		runtime:      runtime,
		graph:        runtime.graph,
		repositories: runtime.repositories,
	}
}

func (builder *Builder) mergeConfig(userConf, imageConf *Config) {
	if userConf.Hostname != "" {
		userConf.Hostname = imageConf.Hostname
	}
	if userConf.User != "" {
		userConf.User = imageConf.User
	}
	if userConf.Memory == 0 {
		userConf.Memory = imageConf.Memory
	}
	if userConf.MemorySwap == 0 {
		userConf.MemorySwap = imageConf.MemorySwap
	}
	if userConf.PortSpecs == nil || len(userConf.PortSpecs) == 0 {
		userConf.PortSpecs = imageConf.PortSpecs
	}
	if !userConf.Tty {
		userConf.Tty = userConf.Tty
	}
	if !userConf.OpenStdin {
		userConf.OpenStdin = imageConf.OpenStdin
	}
	if !userConf.StdinOnce {
		userConf.StdinOnce = imageConf.StdinOnce
	}
	if userConf.Env == nil || len(userConf.Env) == 0 {
		userConf.Env = imageConf.Env
	}
	if userConf.Cmd == nil || len(userConf.Cmd) == 0 {
		userConf.Cmd = imageConf.Cmd
	}
	if userConf.Dns == nil || len(userConf.Dns) == 0 {
		userConf.Dns = imageConf.Dns
	}
}

func (builder *Builder) Create(config *Config) (*Container, error) {
	// Lookup image
	img, err := builder.repositories.LookupImage(config.Image)
	if err != nil {
		return nil, err
	}

	if img.Config != nil {
		builder.mergeConfig(config, img.Config)
	}

	if config.Cmd == nil {
		return nil, fmt.Errorf("No command specified")
	}

	// Generate id
	id := GenerateId()
	// Generate default hostname
	// FIXME: the lxc template no longer needs to set a default hostname
	if config.Hostname == "" {
		config.Hostname = id[:12]
	}

	container := &Container{
		// FIXME: we should generate the ID here instead of receiving it as an argument
		Id:              id,
		Created:         time.Now(),
		Path:            config.Cmd[0],
		Args:            config.Cmd[1:], //FIXME: de-duplicate from config
		Config:          config,
		Image:           img.Id, // Always use the resolved image id
		NetworkSettings: &NetworkSettings{},
		// FIXME: do we need to store this in the container?
		SysInitPath: sysInitPath,
	}
	container.root = builder.runtime.containerRoot(container.Id)
	// Step 1: create the container directory.
	// This doubles as a barrier to avoid race conditions.
	if err := os.Mkdir(container.root, 0700); err != nil {
		return nil, err
	}

	// If custom dns exists, then create a resolv.conf for the container
	if len(config.Dns) > 0 {
		container.ResolvConfPath = path.Join(container.root, "resolv.conf")
		f, err := os.Create(container.ResolvConfPath)
		if err != nil {
			return nil, err
		}
		defer f.Close()
		for _, dns := range config.Dns {
			if _, err := f.Write([]byte("nameserver " + dns + "\n")); err != nil {
				return nil, err
			}
		}
	} else {
		container.ResolvConfPath = "/etc/resolv.conf"
	}

	// Step 2: save the container json
	if err := container.ToDisk(); err != nil {
		return nil, err
	}
	// Step 3: register the container
	if err := builder.runtime.Register(container); err != nil {
		return nil, err
	}
	return container, nil
}

// Commit creates a new filesystem image from the current state of a container.
// The image can optionally be tagged into a repository
func (builder *Builder) Commit(container *Container, repository, tag, comment, author string, config *Config) (*Image, error) {
	// FIXME: freeze the container before copying it to avoid data corruption?
	// FIXME: this shouldn't be in commands.
	rwTar, err := container.ExportRw()
	if err != nil {
		return nil, err
	}
	// Create a new image from the container's base layers + a new layer from container changes
	img, err := builder.graph.Create(rwTar, container, comment, author, config)
	if err != nil {
		return nil, err
	}
	// Register the image if needed
	if repository != "" {
		if err := builder.repositories.Set(repository, tag, img.Id, true); err != nil {
			return img, err
		}
	}
	return img, nil
}

func (builder *Builder) clearTmp(containers, images map[string]struct{}) {
	for c := range containers {
		tmp := builder.runtime.Get(c)
		builder.runtime.Destroy(tmp)
		Debugf("Removing container %s", c)
	}
	for i := range images {
		builder.runtime.graph.Delete(i)
		Debugf("Removing image %s", i)
	}
}

func (builder *Builder) getCachedImage(image *Image, config *Config) (*Image, error) {
	// Retrieve all images
	images, err := builder.graph.All()
	if err != nil {
		return nil, err
	}

	// Store the tree in a map of map (map[parentId][childId])
	imageMap := make(map[string]map[string]struct{})
	for _, img := range images {
		if _, exists := imageMap[img.Parent]; !exists {
			imageMap[img.Parent] = make(map[string]struct{})
		}
		imageMap[img.Parent][img.Id] = struct{}{}
	}

	// Loop on the children of the given image and check the config
	for elem := range imageMap[image.Id] {
		img, err := builder.graph.Get(elem)
		if err != nil {
			return nil, err
		}
		if CompareConfig(&img.ContainerConfig, config) {
			return img, nil
		}
	}
	return nil, nil
}

func (builder *Builder) Build(dockerfile io.Reader, stdout io.Writer) (*Image, error) {
	var (
		image, base   *Image
		config        *Config
		maintainer    string
		env           map[string]string   = make(map[string]string)
		tmpContainers map[string]struct{} = make(map[string]struct{})
		tmpImages     map[string]struct{} = make(map[string]struct{})
	)
	defer builder.clearTmp(tmpContainers, tmpImages)

	file := bufio.NewReader(dockerfile)
	for {
		line, err := file.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, err
		}
		line = strings.Replace(strings.TrimSpace(line), "	", " ", 1)
		// Skip comments and empty line
		if len(line) == 0 || line[0] == '#' {
			continue
		}
		tmp := strings.SplitN(line, " ", 2)
		if len(tmp) != 2 {
			return nil, fmt.Errorf("Invalid Dockerfile format")
		}
		instruction := strings.Trim(tmp[0], " ")
		arguments := strings.Trim(tmp[1], " ")
		switch strings.ToLower(instruction) {
		case "from":
			fmt.Fprintf(stdout, "FROM %s\n", arguments)
			image, err = builder.runtime.repositories.LookupImage(arguments)
			if err != nil {
				if builder.runtime.graph.IsNotExist(err) {

					var tag, remote string
					if strings.Contains(arguments, ":") {
						remoteParts := strings.Split(arguments, ":")
						tag = remoteParts[1]
						remote = remoteParts[0]
					} else {
						remote = arguments
					}

					if err := builder.runtime.graph.PullRepository(stdout, remote, tag, builder.runtime.repositories, builder.runtime.authConfig); err != nil {
						return nil, err
					}

					image, err = builder.runtime.repositories.LookupImage(arguments)
					if err != nil {
						return nil, err
					}
				} else {
					return nil, err
				}
			}
			config = &Config{}

			break
		case "maintainer":
			fmt.Fprintf(stdout, "MAINTAINER %s\n", arguments)
			maintainer = arguments
			break
		case "run":
			fmt.Fprintf(stdout, "RUN %s\n", arguments)
			if image == nil {
				return nil, fmt.Errorf("Please provide a source image with `from` prior to run")
			}
			config, err := ParseRun([]string{image.Id, "/bin/sh", "-c", arguments}, nil, builder.runtime.capabilities)
			if err != nil {
				return nil, err
			}

			for key, value := range env {
				config.Env = append(config.Env, fmt.Sprintf("%s=%s", key, value))
			}

			if cache, err := builder.getCachedImage(image, config); err != nil {
				return nil, err
			} else if cache != nil {
				image = cache
				fmt.Fprintf(stdout, "===> %s\n", image.ShortId())
				break
			}

			Debugf("Env -----> %v ------ %v\n", config.Env, env)

			// Create the container and start it
			c, err := builder.Create(config)
			if err != nil {
				return nil, err
			}

			if os.Getenv("DEBUG") != "" {
				out, _ := c.StdoutPipe()
				err2, _ := c.StderrPipe()
				go io.Copy(os.Stdout, out)
				go io.Copy(os.Stdout, err2)
			}

			if err := c.Start(); err != nil {
				return nil, err
			}
			tmpContainers[c.Id] = struct{}{}

			// Wait for it to finish
			if result := c.Wait(); result != 0 {
				return nil, fmt.Errorf("!!! '%s' return non-zero exit code '%d'. Aborting.", arguments, result)
			}

			// Commit the container
			base, err = builder.Commit(c, "", "", "", maintainer, nil)
			if err != nil {
				return nil, err
			}
			tmpImages[base.Id] = struct{}{}

			fmt.Fprintf(stdout, "===> %s\n", base.ShortId())

			// use the base as the new image
			image = base

			break
		case "env":
			tmp := strings.SplitN(arguments, " ", 2)
			if len(tmp) != 2 {
				return nil, fmt.Errorf("Invalid ENV format")
			}
			key := strings.Trim(tmp[0], " ")
			value := strings.Trim(tmp[1], " ")
			fmt.Fprintf(stdout, "ENV %s %s\n", key, value)
			env[key] = value
			if image != nil {
				fmt.Fprintf(stdout, "===> %s\n", image.ShortId())
			} else {
				fmt.Fprintf(stdout, "===> <nil>\n")
			}
			break
		case "cmd":
			fmt.Fprintf(stdout, "CMD %s\n", arguments)

			// Create the container and start it
			c, err := builder.Create(&Config{Image: image.Id, Cmd: []string{"", ""}})
			if err != nil {
				return nil, err
			}
			if err := c.Start(); err != nil {
				return nil, err
			}
			tmpContainers[c.Id] = struct{}{}

			cmd := []string{}
			if err := json.Unmarshal([]byte(arguments), &cmd); err != nil {
				return nil, err
			}
			config.Cmd = cmd

			// Commit the container
			base, err = builder.Commit(c, "", "", "", maintainer, config)
			if err != nil {
				return nil, err
			}
			tmpImages[base.Id] = struct{}{}

			fmt.Fprintf(stdout, "===> %s\n", base.ShortId())
			image = base
			break
		case "expose":
			ports := strings.Split(arguments, " ")

			fmt.Fprintf(stdout, "EXPOSE %v\n", ports)
			if image == nil {
				return nil, fmt.Errorf("Please provide a source image with `from` prior to copy")
			}

			// Create the container and start it
			c, err := builder.Create(&Config{Image: image.Id, Cmd: []string{"", ""}})
			if err != nil {
				return nil, err
			}
			if err := c.Start(); err != nil {
				return nil, err
			}
			tmpContainers[c.Id] = struct{}{}

			config.PortSpecs = append(ports, config.PortSpecs...)

			// Commit the container
			base, err = builder.Commit(c, "", "", "", maintainer, config)
			if err != nil {
				return nil, err
			}
			tmpImages[base.Id] = struct{}{}

			fmt.Fprintf(stdout, "===> %s\n", base.ShortId())
			image = base
			break
		case "insert":
			if image == nil {
				return nil, fmt.Errorf("Please provide a source image with `from` prior to copy")
			}
			tmp = strings.SplitN(arguments, " ", 2)
			if len(tmp) != 2 {
				return nil, fmt.Errorf("Invalid INSERT format")
			}
			sourceUrl := strings.Trim(tmp[0], " ")
			destPath := strings.Trim(tmp[1], " ")
			fmt.Fprintf(stdout, "COPY %s to %s in %s\n", sourceUrl, destPath, base.ShortId())

			file, err := Download(sourceUrl, stdout)
			if err != nil {
				return nil, err
			}
			defer file.Body.Close()

			config, err := ParseRun([]string{base.Id, "echo", "insert", sourceUrl, destPath}, nil, builder.runtime.capabilities)
			if err != nil {
				return nil, err
			}
			c, err := builder.Create(config)
			if err != nil {
				return nil, err
			}

			if err := c.Start(); err != nil {
				return nil, err
			}

			// Wait for echo to finish
			if result := c.Wait(); result != 0 {
				return nil, fmt.Errorf("!!! '%s' return non-zero exit code '%d'. Aborting.", arguments, result)
			}

			if err := c.Inject(file.Body, destPath); err != nil {
				return nil, err
			}

			base, err = builder.Commit(c, "", "", "", maintainer, nil)
			if err != nil {
				return nil, err
			}
			fmt.Fprintf(stdout, "===> %s\n", base.ShortId())

			image = base

			break
		default:
			fmt.Fprintf(stdout, "Skipping unknown instruction %s\n", strings.ToUpper(instruction))
		}
	}
	if image != nil {
		// The build is successful, keep the temporary containers and images
		for i := range tmpImages {
			delete(tmpImages, i)
		}
		for i := range tmpContainers {
			delete(tmpContainers, i)
		}
		fmt.Fprintf(stdout, "Build finished. image id: %s\n", image.ShortId())
		return image, nil
	}
	return nil, fmt.Errorf("An error occured during the build\n")
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package tests_test

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"net"
	"strconv"
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	ini "gopkg.in/ini.v1"

	"kubevirt.io/kubevirt/pkg/api/v1"
	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/rest"
	"kubevirt.io/kubevirt/tests"
)

var _ = Describe("Vmlifecycle", func() {

	flag.Parse()

	virtClient, err := kubecli.GetKubevirtClient()
	tests.PanicOnError(err)
	var vm *v1.VirtualMachine

	getVmNode := func() string {
		obj, err := virtClient.RestClient().Get().Resource("virtualmachines").Namespace(tests.NamespaceTestDefault).Name(vm.GetObjectMeta().GetName()).Do().Get()
		Expect(err).ToNot(HaveOccurred())
		return obj.(*v1.VirtualMachine).Status.NodeName
	}

	checkSpiceConnection := func() {
		raw, err := virtClient.RestClient().Get().Resource("virtualmachines").SetHeader("Accept", rest.MIME_INI).SubResource("spice").Namespace(tests.NamespaceTestDefault).Name(vm.GetObjectMeta().GetName()).Do().Raw()
		Expect(err).To(BeNil())
		spiceINI, err := ini.Load(raw)
		Expect(err).NotTo(HaveOccurred())
		spice := v1.SpiceInfo{}
		Expect(spiceINI.Section("virt-viewer").MapTo(&spice)).To(Succeed())

		proxy := strings.TrimPrefix(spice.Proxy, "http://")
		host := fmt.Sprintf("%s:%d", spice.Host, spice.Port)

		// Let's see if we can connect to the spice port through the proxy
		conn, err := net.Dial("tcp", proxy)
		Expect(err).To(BeNil())
		conn.Write([]byte("CONNECT " + host + " HTTP/1.1\r\n"))
		conn.Write([]byte("Host: " + host + "\r\n"))
		conn.Write([]byte("\r\n"))
		line, err := bufio.NewReader(conn).ReadString('\n')
		Expect(err).To(BeNil())
		Expect(strings.TrimSpace(line)).To(Equal("HTTP/1.1 200 Connection established"))

		// Let's send a spice handshake
		conn.Write(newSpiceHandshake())

		// Let's parse the response
		var i int32
		x := make([]byte, 4, 4)
		io.ReadFull(conn, x)
		Expect(string(x)).To(Equal("REDQ")) // spice magic
		binary.Read(conn, binary.LittleEndian, &i)
		Expect(i).To(Equal(int32(2)), "Major version does not match.")
		binary.Read(conn, binary.LittleEndian, &i)
		Expect(i).To(Equal(int32(2)), "Minor version does not match.")
		binary.Read(conn, binary.LittleEndian, &i)
		Expect(i).To(BeNumerically(">", 4), "Message not long enough.")
		binary.Read(conn, binary.LittleEndian, &i)
		Expect(i).To(Equal(int32(0)), "Message status is not OK.") // 0 is equal to OK
	}

	BeforeEach(func() {
		vm = tests.NewRandomVM()
		tests.BeforeTestCleanup()
	})

	Context("New VM with a spice connection given", func() {

		It("should return no connection details if VM does not exist", func(done Done) {
			result := virtClient.RestClient().Get().Resource("virtualmachines").SubResource("spice").Namespace(tests.NamespaceTestDefault).Name("something-random").Do()
			Expect(result.Error()).NotTo(BeNil())
			close(done)
		}, 3)

		It("should return connection details for running VMs in ini format", func(done Done) {
			// Create the VM
			result := virtClient.RestClient().Post().Resource("virtualmachines").Namespace(tests.NamespaceTestDefault).Body(vm).Do()
			obj, err := result.Get()
			Expect(err).To(BeNil())

			// Block until the VM is running
			tests.WaitForSuccessfulVMStart(obj)

			raw, err := virtClient.RestClient().Get().Resource("virtualmachines").SetHeader("Accept", rest.MIME_INI).SubResource("spice").Namespace(tests.NamespaceTestDefault).Name(vm.GetObjectMeta().GetName()).Do().Raw()
			spice, err := ini.Load(raw)
			Expect(err).To(Not(HaveOccurred()))

			Expect(spice.Section("virt-viewer")).NotTo(BeNil())
			section := spice.Section("virt-viewer")
			Expect(section.HasKey("type")).To(BeTrue())
			Expect(section.HasKey("host")).To(BeTrue())
			Expect(section.HasKey("port")).To(BeTrue())
			Expect(strconv.Atoi(section.Key("port").Value())).To(BeNumerically(">=", int32(5900)))
			Expect(section.HasKey("proxy")).To(BeTrue())
			close(done)
		}, 30)

		It("should return connection details for running VMs in json format", func(done Done) {
			// Create the VM
			result := virtClient.RestClient().Post().Resource("virtualmachines").Namespace(tests.NamespaceTestDefault).Body(vm).Do()
			obj, err := result.Get()
			Expect(err).To(BeNil())

			// Block until the VM is running
			tests.WaitForSuccessfulVMStart(obj)

			obj, err = virtClient.RestClient().Get().Resource("virtualmachines").SetHeader("Accept", rest.MIME_JSON).SubResource("spice").Namespace(tests.NamespaceTestDefault).Name(vm.GetObjectMeta().GetName()).Do().Get()
			Expect(err).To(BeNil())
			spice := obj.(*v1.Spice).Info
			Expect(spice.Type).To(Equal("spice"))
			Expect(spice.Port).To(BeNumerically(">=", int32(5900)))
			close(done)
		}, 30)

		It("should allow accessing the spice device on the VM", func(done Done) {
			// Create the VM
			result := virtClient.RestClient().Post().Resource("virtualmachines").Namespace(tests.NamespaceTestDefault).Body(vm).Do()
			obj, err := result.Get()
			Expect(err).To(BeNil())

			// Block until the VM is running
			tests.WaitForSuccessfulVMStart(obj)

			checkSpiceConnection()

			close(done)
		}, 30)
	})

	Context("Two new VMs scheduled on the same node with a spice graphics device", func() {

		It("should start without port clashes", func(done Done) {
			// Create the VM
			result := virtClient.RestClient().Post().Resource("virtualmachines").Namespace(tests.NamespaceTestDefault).Body(vm).Do()
			obj, err := result.Get()
			Expect(err).To(BeNil())

			// Block until the VM is running
			nodeName := tests.WaitForSuccessfulVMStart(obj)

			vm1 := tests.NewRandomVM()
			vm1.Spec.NodeSelector = map[string]string{"kubernetes.io/hostname": nodeName}
			result = virtClient.RestClient().Post().Resource("virtualmachines").Namespace(tests.NamespaceTestDefault).Body(vm1).Do()
			obj1, err := result.Get()
			Expect(err).To(BeNil())

			// Block until the VM is running
			tests.WaitForSuccessfulVMStart(obj1)

			obj, err = virtClient.RestClient().Get().Resource("virtualmachines").SetHeader("Accept", rest.MIME_JSON).SubResource("spice").Namespace(tests.NamespaceTestDefault).Name(vm.GetObjectMeta().GetName()).Do().Get()
			Expect(err).To(BeNil())
			obj1, err = virtClient.RestClient().Get().Resource("virtualmachines").SetHeader("Accept", rest.MIME_JSON).SubResource("spice").Namespace(tests.NamespaceTestDefault).Name(vm1.GetObjectMeta().GetName()).Do().Get()
			Expect(err).To(BeNil())
			Expect(obj.(*v1.Spice).Info.Port).ToNot(BeNumerically("==", obj1.(*v1.Spice).Info.Port))
			close(done)
		}, 30)
	})

	Context("Migrate VM with spice connection", func() {

		BeforeEach(func() {
			if len(tests.GetReadyNodes()) < 2 {
				Skip("To test migrations, at least two nodes need to be active")
			}
			// Create the VM
			result := virtClient.RestClient().Post().Resource("virtualmachines").Namespace(tests.NamespaceTestDefault).Body(vm).Do()
			obj, err := result.Get()
			Expect(err).To(BeNil())

			// Block until the VM is running
			tests.WaitForSuccessfulVMStart(obj)
		})

		It("should allow accessing the spice device on the VM", func() {
			sourceNode := getVmNode()

			migration := tests.NewRandomMigrationForVm(vm)
			err = virtClient.RestClient().Post().Resource("migrations").Namespace(tests.NamespaceTestDefault).Body(migration).Do().Error()
			Expect(err).ToNot(HaveOccurred())

			Eventually(getVmNode, 30*time.Second, time.Second).ShouldNot(Equal(sourceNode))

			Eventually(func() v1.VMPhase {
				obj, err := virtClient.RestClient().Get().Resource("virtualmachines").Namespace(tests.NamespaceTestDefault).Name(vm.GetObjectMeta().GetName()).Do().Get()
				Expect(err).ToNot(HaveOccurred())
				fetchedVM := obj.(*v1.VirtualMachine)
				return fetchedVM.Status.Phase
			}, 10*time.Second, time.Second).Should(Equal(v1.Running))

			checkSpiceConnection()
		})
	})
})

func newSpiceHandshake() []byte {
	var b []byte
	bb := bytes.NewBuffer(b)
	bb.Write([]byte("REDQ"))                                  // spice magic
	binary.Write(bb, binary.LittleEndian, uint32(2))          // protocol major version
	binary.Write(bb, binary.LittleEndian, uint32(2))          // protocol minor version
	binary.Write(bb, binary.LittleEndian, uint32(22))         // message size
	binary.Write(bb, binary.LittleEndian, uint32(rand.Int())) // session id
	binary.Write(bb, binary.LittleEndian, uint8(3))           // channel type
	binary.Write(bb, binary.LittleEndian, uint8(0))           // channel id
	binary.Write(bb, binary.LittleEndian, uint32(1))          // number of common capabilities
	binary.Write(bb, binary.LittleEndian, uint32(0))          // number of channel capabilities
	binary.Write(bb, binary.LittleEndian, uint32(18))         // capabilities offset
	binary.Write(bb, binary.LittleEndian, uint32(13))         // client common capabilities
	return bb.Bytes()
}
<commit_msg>remove spice tests in favor of vnc<commit_after><|endoftext|>


<commit_before>// Package prerender provides a Prerender.io handler implementation and a
// Negroni middleware.
package prerender

import (
	"compress/gzip"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"

	e "github.com/jqatampa/gadget-arm/errors"
	"google.golang.org/appengine"
	"google.golang.org/appengine/urlfetch"
)

// Options provides you with the ability to specify a custom Prerender.io URL
// as well as a Prerender.io Token to include as an X-Prerender-Token header
// to the upstream server.
type Options struct {
	PrerenderURL   *url.URL
	Token          string
	BlackList      []regexp.Regexp
	WhiteList      []regexp.Regexp
	UsingAppEngine bool
}

// NewOptions generates a default Options struct pointing to the Prerender.io
// service, obtaining a Token from the environment variable PRERENDER_TOKEN.
// No blacklist/whitelist is created.
func NewOptions() *Options {
	url, _ := url.Parse("https://service.prerender.io/")
	return &Options{
		PrerenderURL: url,
		Token:        os.Getenv("PRERENDER_TOKEN"),
		BlackList:    nil,
		WhiteList:    nil,
		UsingAppEngine: false,
	}
}

// Prerender exposes methods to validate and serve content from a Prerender.io
// upstream server.
type Prerender struct {
	Options *Options
}

// NewPrerender generates a new Prerender instance.
func (o *Options) NewPrerender() *Prerender {
	return &Prerender{Options: o}
}

// ServeHTTP allows Prerender to act as a Negroni middleware.
func (p *Prerender) ServeHTTP(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	fmt.Println("Prerender")
	if p.ShouldPrerender(r) {
		p.PreRenderHandler(rw, r)
	} else if next != nil {
		next(rw, r)
	}
}

// ShouldPrerender analyzes the request to determine whether it should be routed
// to a Prerender.io upstream server.
func (p *Prerender) ShouldPrerender(or *http.Request) bool {
	fmt.Println(or)
	userAgent := strings.ToLower(or.Header.Get("User-Agent"))
	bufferAgent := or.Header.Get("X-Bufferbot")
	isRequestingPrerenderedPage := false
	reqURL := strings.ToLower(or.URL.String())

	// No user agent, don't prerender
	if userAgent == "" {
		return false
	}

	// Not a GET or HEAD request, don't prerender
	if or.Method != "GET" && or.Method != "HEAD" {
		return false
	}

	// Static resource, don't prerender
	for _, extension := range skippedTypes {
		if strings.HasSuffix(reqURL, strings.ToLower(extension)) {
			return false
		}
	}

	// Buffer Agent or requesting an excaped fragment, request prerender
	if bufferAgent != "" || or.URL.Query().Get("_escaped_fragment_") != "" {
		isRequestingPrerenderedPage = true
	}

	// Cralwer, request prerender
	for _, crawlerAgent := range crawlerUserAgents {
		if strings.Contains(crawlerAgent, strings.ToLower(userAgent)) {
			isRequestingPrerenderedPage = true
			break
		}
	}

	// If it's a bot/crawler/escaped fragment request apply Blacklist/Whitelist logic
	if isRequestingPrerenderedPage {
		if p.Options.WhiteList != nil {
			matchFound := false
			for _, val := range p.Options.WhiteList {
				if val.MatchString(reqURL) {
					matchFound = true
					break
				}
			}
			if !matchFound {
				return false
			}
		}

		if p.Options.BlackList != nil {
			matchFound := false
			for _, val := range p.Options.BlackList {
				if val.MatchString(reqURL) {
					matchFound = true
					break
				}
			}
			if matchFound {
				return false
			}
		}
	}

	return isRequestingPrerenderedPage
}

func (p *Prerender) buildURL(or *http.Request) string {
	url := p.Options.PrerenderURL

	if !strings.HasSuffix(url.String(), "/") {
		url.Path = url.Path + "/"
	}

	var protocol = or.URL.Scheme

	if cf := or.Header.Get("CF-Visitor"); cf != "" {
		match := cfSchemeRegex.FindStringSubmatch(cf)
		if len(match) > 1 {
			protocol = match[1]
		}
	}

	if len(protocol) == 0 {
		protocol = "http"
	}

	if fp := or.Header.Get("X-Forwarded-Proto"); fp != "" {
		protocol = strings.Split(fp, ",")[0]
	}

	apiURL := url.String() + protocol + "://" + or.Host + or.URL.Path + "?" +
		or.URL.RawQuery
	return apiURL
}

// PreRenderHandler is a net/http compatible handler that proxies a request to
// the configured Prerender.io URL.  All upstream requests are made with an
// Accept-Encoding=gzip header.  Responses are provided either uncompressed or
// gzip compressed based on the downstream requests Accept-Encoding header
func (p *Prerender) PreRenderHandler(rw http.ResponseWriter, or *http.Request) {
	client := &http.Client{}

	req, err := http.NewRequest("GET", p.buildURL(or), nil)
	e.Check(err)

	if p.Options.Token != "" {
		req.Header.Set("X-Prerender-Token", p.Options.Token)
	}
	req.Header.Set("User-Agent", or.Header.Get("User-Agent"))
	req.Header.Set("Content-Type", or.Header.Get("Content-Type"))
	req.Header.Set("Accept-Encoding", "gzip")

	if p.Options.UsingAppEngine {
		ctx := appengine.NewContext(or)
		client = urlfetch.Client(ctx)
	}

	res, err := client.Do(req)

	fmt.Println(res)
	e.Check(err)

	rw.Header().Set("Content-Type", res.Header.Get("Content-Type"))

	defer res.Body.Close()

	//Figure out whether the client accepts gzip responses
	doGzip := strings.Contains(or.Header.Get("Accept-Encoding"), "gzip")
	isGzip := strings.Contains(res.Header.Get("Content-Encoding"), "gzip")

	if doGzip && !isGzip {
		// gzip raw response
		rw.Header().Set("Content-Encoding", "gzip")
		gz := gzip.NewWriter(rw)
		defer gz.Close()
		io.Copy(gz, res.Body)
		gz.Flush()

	} else if !doGzip && isGzip {
		// gunzip response
		gz, err := gzip.NewReader(res.Body)
		e.Check(err)
		defer gz.Close()
		io.Copy(rw, gz)
	} else {
		// Pass through, gzip/gzip or raw/raw
		rw.Header().Set("Content-Encoding", res.Header.Get("Content-Encoding"))
		io.Copy(rw, res.Body)

	}
}
<commit_msg>Support escaped_fragment with empty string<commit_after>// Package prerender provides a Prerender.io handler implementation and a
// Negroni middleware.
package prerender

import (
	"compress/gzip"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"

	e "github.com/jqatampa/gadget-arm/errors"
	"google.golang.org/appengine"
	"google.golang.org/appengine/urlfetch"
)

// Options provides you with the ability to specify a custom Prerender.io URL
// as well as a Prerender.io Token to include as an X-Prerender-Token header
// to the upstream server.
type Options struct {
	PrerenderURL   *url.URL
	Token          string
	BlackList      []regexp.Regexp
	WhiteList      []regexp.Regexp
	UsingAppEngine bool
}

// NewOptions generates a default Options struct pointing to the Prerender.io
// service, obtaining a Token from the environment variable PRERENDER_TOKEN.
// No blacklist/whitelist is created.
func NewOptions() *Options {
	url, _ := url.Parse("https://service.prerender.io/")
	return &Options{
		PrerenderURL: url,
		Token:        os.Getenv("PRERENDER_TOKEN"),
		BlackList:    nil,
		WhiteList:    nil,
		UsingAppEngine: false,
	}
}

// Prerender exposes methods to validate and serve content from a Prerender.io
// upstream server.
type Prerender struct {
	Options *Options
}

// NewPrerender generates a new Prerender instance.
func (o *Options) NewPrerender() *Prerender {
	return &Prerender{Options: o}
}

// ServeHTTP allows Prerender to act as a Negroni middleware.
func (p *Prerender) ServeHTTP(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	fmt.Println("Prerender")
	if p.ShouldPrerender(r) {
		p.PreRenderHandler(rw, r)
	} else if next != nil {
		next(rw, r)
	}
}

// ShouldPrerender analyzes the request to determine whether it should be routed
// to a Prerender.io upstream server.
func (p *Prerender) ShouldPrerender(or *http.Request) bool {
	fmt.Println(or)
	userAgent := strings.ToLower(or.Header.Get("User-Agent"))
	bufferAgent := or.Header.Get("X-Bufferbot")
	isRequestingPrerenderedPage := false
	reqURL := strings.ToLower(or.URL.String())

	// No user agent, don't prerender
	if userAgent == "" {
		return false
	}

	// Not a GET or HEAD request, don't prerender
	if or.Method != "GET" && or.Method != "HEAD" {
		return false
	}

	// Static resource, don't prerender
	for _, extension := range skippedTypes {
		if strings.HasSuffix(reqURL, strings.ToLower(extension)) {
			return false
		}
	}

	// Buffer Agent or requesting an excaped fragment, request prerender
	if _, ok := or.URL.Query()["_escaped_fragment_"]; bufferAgent != "" || ok {
		isRequestingPrerenderedPage = true
	}

	// Cralwer, request prerender
	for _, crawlerAgent := range crawlerUserAgents {
		if strings.Contains(crawlerAgent, strings.ToLower(userAgent)) {
			isRequestingPrerenderedPage = true
			break
		}
	}

	// If it's a bot/crawler/escaped fragment request apply Blacklist/Whitelist logic
	if isRequestingPrerenderedPage {
		if p.Options.WhiteList != nil {
			matchFound := false
			for _, val := range p.Options.WhiteList {
				if val.MatchString(reqURL) {
					matchFound = true
					break
				}
			}
			if !matchFound {
				return false
			}
		}

		if p.Options.BlackList != nil {
			matchFound := false
			for _, val := range p.Options.BlackList {
				if val.MatchString(reqURL) {
					matchFound = true
					break
				}
			}
			if matchFound {
				return false
			}
		}
	}

	return isRequestingPrerenderedPage
}

func (p *Prerender) buildURL(or *http.Request) string {
	url := p.Options.PrerenderURL

	if !strings.HasSuffix(url.String(), "/") {
		url.Path = url.Path + "/"
	}

	var protocol = or.URL.Scheme

	if cf := or.Header.Get("CF-Visitor"); cf != "" {
		match := cfSchemeRegex.FindStringSubmatch(cf)
		if len(match) > 1 {
			protocol = match[1]
		}
	}

	if len(protocol) == 0 {
		protocol = "http"
	}

	if fp := or.Header.Get("X-Forwarded-Proto"); fp != "" {
		protocol = strings.Split(fp, ",")[0]
	}

	apiURL := url.String() + protocol + "://" + or.Host + or.URL.Path + "?" +
		or.URL.RawQuery
	return apiURL
}

// PreRenderHandler is a net/http compatible handler that proxies a request to
// the configured Prerender.io URL.  All upstream requests are made with an
// Accept-Encoding=gzip header.  Responses are provided either uncompressed or
// gzip compressed based on the downstream requests Accept-Encoding header
func (p *Prerender) PreRenderHandler(rw http.ResponseWriter, or *http.Request) {
	client := &http.Client{}

	req, err := http.NewRequest("GET", p.buildURL(or), nil)
	e.Check(err)

	if p.Options.Token != "" {
		req.Header.Set("X-Prerender-Token", p.Options.Token)
	}
	req.Header.Set("User-Agent", or.Header.Get("User-Agent"))
	req.Header.Set("Content-Type", or.Header.Get("Content-Type"))
	req.Header.Set("Accept-Encoding", "gzip")

	if p.Options.UsingAppEngine {
		ctx := appengine.NewContext(or)
		client = urlfetch.Client(ctx)
	}

	res, err := client.Do(req)

	fmt.Println(res)
	e.Check(err)

	rw.Header().Set("Content-Type", res.Header.Get("Content-Type"))

	defer res.Body.Close()

	//Figure out whether the client accepts gzip responses
	doGzip := strings.Contains(or.Header.Get("Accept-Encoding"), "gzip")
	isGzip := strings.Contains(res.Header.Get("Content-Encoding"), "gzip")

	if doGzip && !isGzip {
		// gzip raw response
		rw.Header().Set("Content-Encoding", "gzip")
		gz := gzip.NewWriter(rw)
		defer gz.Close()
		io.Copy(gz, res.Body)
		gz.Flush()

	} else if !doGzip && isGzip {
		// gunzip response
		gz, err := gzip.NewReader(res.Body)
		e.Check(err)
		defer gz.Close()
		io.Copy(rw, gz)
	} else {
		// Pass through, gzip/gzip or raw/raw
		rw.Header().Set("Content-Encoding", res.Header.Get("Content-Encoding"))
		io.Copy(rw, res.Body)

	}
}
<|endoftext|>


<commit_before>package cachetype

import (
	"context"
	"fmt"
	"reflect"
	"sort"
	"strings"
	"time"

	"github.com/hashicorp/go-bexpr"
	"github.com/hashicorp/go-hclog"

	"github.com/hashicorp/consul/agent/cache"
	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/agent/submatview"
	"github.com/hashicorp/consul/lib/retry"
	"github.com/hashicorp/consul/proto/pbservice"
	"github.com/hashicorp/consul/proto/pbsubscribe"
)

const (
	// Recommended name for registration.
	StreamingHealthServicesName = "streaming-health-services"
)

// StreamingHealthServices supports fetching discovering service instances via the
// catalog using the streaming gRPC endpoint.
type StreamingHealthServices struct {
	RegisterOptionsBlockingRefresh
	deps MaterializerDeps
}

// RegisterOptions returns options with a much shorter LastGetTTL than the default.
// Unlike other cache-types, StreamingHealthServices runs a materialized view in
// the background which will receive streamed events from a server. If the cache
// is not being used, that stream uses memory on the server and network transfer
// between the client and the server.
// The materialize view and the stream are stopped when the cache entry expires,
// so using a shorter TTL ensures the cache entry expires sooner.
func (c *StreamingHealthServices) RegisterOptions() cache.RegisterOptions {
	opts := c.RegisterOptionsBlockingRefresh.RegisterOptions()
	opts.LastGetTTL = 20 * time.Minute
	return opts
}

// NewStreamingHealthServices creates a cache-type for watching for service
// health results via streaming updates.
func NewStreamingHealthServices(deps MaterializerDeps) *StreamingHealthServices {
	return &StreamingHealthServices{deps: deps}
}

type MaterializerDeps struct {
	Client submatview.StreamClient
	Logger hclog.Logger
}

// Fetch service health from the materialized view. If no materialized view
// exists, create one and start it running in a goroutine. The goroutine will
// exit when the cache entry storing the result is expired, the cache will call
// Close on the result.State.
//
// Fetch implements part of the cache.Type interface, and assumes that the
// caller ensures that only a single call to Fetch is running at any time.
func (c *StreamingHealthServices) Fetch(opts cache.FetchOptions, req cache.Request) (cache.FetchResult, error) {
	if opts.LastResult != nil && opts.LastResult.State != nil {
		return opts.LastResult.State.(*streamingHealthState).Fetch(opts)
	}

	srvReq := req.(*structs.ServiceSpecificRequest)
	newReqFn := func(index uint64) pbsubscribe.SubscribeRequest {
		req := pbsubscribe.SubscribeRequest{
			Topic:      pbsubscribe.Topic_ServiceHealth,
			Key:        srvReq.ServiceName,
			Token:      srvReq.Token,
			Datacenter: srvReq.Datacenter,
			Index:      index,
			Namespace:  srvReq.EnterpriseMeta.NamespaceOrEmpty(),
		}
		if srvReq.Connect {
			req.Topic = pbsubscribe.Topic_ServiceHealthConnect
		}
		return req
	}

	materializer, err := newMaterializer(c.deps, newReqFn, srvReq)
	if err != nil {
		return cache.FetchResult{}, err
	}
	ctx, cancel := context.WithCancel(context.TODO())
	go materializer.Run(ctx)

	state := &streamingHealthState{
		materializer: materializer,
		done:         ctx.Done(),
		cancel:       cancel,
	}
	return state.Fetch(opts)
}

func newMaterializer(
	deps MaterializerDeps,
	newRequestFn func(uint64) pbsubscribe.SubscribeRequest,
	req *structs.ServiceSpecificRequest,
) (*submatview.Materializer, error) {
	view, err := newHealthView(req)
	if err != nil {
		return nil, err
	}
	return submatview.NewMaterializer(submatview.Deps{
		View:   view,
		Client: deps.Client,
		Logger: deps.Logger,
		Waiter: &retry.Waiter{
			MinFailures: 1,
			MinWait:     0,
			MaxWait:     60 * time.Second,
			Jitter:      retry.NewJitter(100),
		},
		Request: newRequestFn,
	}), nil
}

// streamingHealthState wraps a Materializer to manage its lifecycle, and to
// add itself to the FetchResult.State.
type streamingHealthState struct {
	materializer *submatview.Materializer
	done         <-chan struct{}
	cancel       func()
}

func (s *streamingHealthState) Close() error {
	s.cancel()
	return nil
}

func (s *streamingHealthState) Fetch(opts cache.FetchOptions) (cache.FetchResult, error) {
	result, err := s.materializer.Fetch(s.done, opts)
	result.State = s
	return result, err
}

func newHealthView(req *structs.ServiceSpecificRequest) (*healthView, error) {
	fe, err := newFilterEvaluator(req)
	if err != nil {
		return nil, err
	}
	return &healthView{
		state:  make(map[string]structs.CheckServiceNode),
		filter: fe,
	}, nil
}

// healthView implements submatview.View for storing the view state
// of a service health result. We store it as a map to make updates and
// deletions a little easier but we could just store a result type
// (IndexedCheckServiceNodes) and update it in place for each event - that
// involves re-sorting each time etc. though.
type healthView struct {
	state  map[string]structs.CheckServiceNode
	filter filterEvaluator
}

// Update implements View
func (s *healthView) Update(events []*pbsubscribe.Event) error {
	for _, event := range events {
		serviceHealth := event.GetServiceHealth()
		if serviceHealth == nil {
			return fmt.Errorf("unexpected event type for service health view: %T",
				event.GetPayload())
		}

		id := serviceHealth.CheckServiceNode.UniqueID()
		switch serviceHealth.Op {
		case pbsubscribe.CatalogOp_Register:
			csn := *pbservice.CheckServiceNodeToStructs(serviceHealth.CheckServiceNode)
			passed, err := s.filter.Evaluate(csn)
			switch {
			case err != nil:
				return err
			case passed:
				s.state[id] = csn
			}

		case pbsubscribe.CatalogOp_Deregister:
			delete(s.state, id)
		}
	}
	return nil
}

type filterEvaluator interface {
	Evaluate(datum interface{}) (bool, error)
}

func newFilterEvaluator(req *structs.ServiceSpecificRequest) (filterEvaluator, error) {
	var evaluators []filterEvaluator

	typ := reflect.TypeOf(structs.CheckServiceNode{})
	if req.Filter != "" {
		e, err := bexpr.CreateEvaluatorForType(req.Filter, nil, typ)
		if err != nil {
			return nil, err
		}
		evaluators = append(evaluators, e)
	}

	if req.ServiceTag != "" {
		// Handle backwards compat with old field
		req.ServiceTags = []string{req.ServiceTag}
	}

	if req.TagFilter && len(req.ServiceTags) > 0 {
		evaluators = append(evaluators, serviceTagEvaluator{tags: req.ServiceTags})
	}

	for key, value := range req.NodeMetaFilters {
		expr := fmt.Sprintf(`"%s" in Node.Meta.%s`, value, key)
		e, err := bexpr.CreateEvaluatorForType(expr, nil, typ)
		if err != nil {
			return nil, err
		}
		evaluators = append(evaluators, e)
	}

	switch len(evaluators) {
	case 0:
		return noopFilterEvaluator{}, nil
	case 1:
		return evaluators[0], nil
	default:
		return &multiFilterEvaluator{evaluators: evaluators}, nil
	}
}

// noopFilterEvaluator may be used in place of a bexpr.Evaluator. The Evaluate
// method always return true, so no items will be filtered out.
type noopFilterEvaluator struct{}

func (noopFilterEvaluator) Evaluate(_ interface{}) (bool, error) {
	return true, nil
}

type multiFilterEvaluator struct {
	evaluators []filterEvaluator
}

func (m multiFilterEvaluator) Evaluate(data interface{}) (bool, error) {
	for _, e := range m.evaluators {
		match, err := e.Evaluate(data)
		if !match || err != nil {
			return match, err
		}
	}
	return true, nil
}

// sortCheckServiceNodes sorts the results to match memdb semantics
// Sort results by Node.Node, if 2 instances match, order by Service.ID
// Will allow result to be stable sorted and match queries without cache
func sortCheckServiceNodes(serviceNodes *structs.IndexedCheckServiceNodes) {
	sort.SliceStable(serviceNodes.Nodes, func(i, j int) bool {
		left := serviceNodes.Nodes[i]
		right := serviceNodes.Nodes[j]
		if left.Node.Node == right.Node.Node {
			return left.Service.ID < right.Service.ID
		}
		return left.Node.Node < right.Node.Node
	})
}

// Result returns the structs.IndexedCheckServiceNodes stored by this view.
func (s *healthView) Result(index uint64) (interface{}, error) {
	result := structs.IndexedCheckServiceNodes{
		Nodes: make(structs.CheckServiceNodes, 0, len(s.state)),
		QueryMeta: structs.QueryMeta{
			Index: index,
		},
	}
	for _, node := range s.state {
		result.Nodes = append(result.Nodes, node)
	}
	sortCheckServiceNodes(&result)

	return &result, nil
}

func (s *healthView) Reset() {
	s.state = make(map[string]structs.CheckServiceNode)
}

// serviceTagEvaluator implements the filterEvaluator to perform filtering
// by service tags. bexpr can not be used at this time, because the filtering
// must be case insensitive for backwards compatibility. In the future this
// may be replaced with bexpr once case insensitive support is added.
type serviceTagEvaluator struct {
	tags []string
}

func (m serviceTagEvaluator) Evaluate(data interface{}) (bool, error) {
	csn, ok := data.(structs.CheckServiceNode)
	if !ok {
		return false, fmt.Errorf("unexpected type %T for structs.CheckServiceNode filter", data)
	}
	for _, tag := range m.tags {
		if !serviceHasTag(csn.Service, tag) {
			// If any one of the expected tags was not found, filter the service
			return false, nil
		}
	}
	return true, nil
}

func serviceHasTag(sn *structs.NodeService, tag string) bool {
	for _, t := range sn.Tags {
		if strings.EqualFold(t, tag) {
			return true
		}
	}
	return false
}
<commit_msg>Tune streaming backoff on errors to retry a bit faster when TCP connections drop<commit_after>package cachetype

import (
	"context"
	"fmt"
	"reflect"
	"sort"
	"strings"
	"time"

	"github.com/hashicorp/go-bexpr"
	"github.com/hashicorp/go-hclog"

	"github.com/hashicorp/consul/agent/cache"
	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/agent/submatview"
	"github.com/hashicorp/consul/lib/retry"
	"github.com/hashicorp/consul/proto/pbservice"
	"github.com/hashicorp/consul/proto/pbsubscribe"
)

const (
	// Recommended name for registration.
	StreamingHealthServicesName = "streaming-health-services"
)

// StreamingHealthServices supports fetching discovering service instances via the
// catalog using the streaming gRPC endpoint.
type StreamingHealthServices struct {
	RegisterOptionsBlockingRefresh
	deps MaterializerDeps
}

// RegisterOptions returns options with a much shorter LastGetTTL than the default.
// Unlike other cache-types, StreamingHealthServices runs a materialized view in
// the background which will receive streamed events from a server. If the cache
// is not being used, that stream uses memory on the server and network transfer
// between the client and the server.
// The materialize view and the stream are stopped when the cache entry expires,
// so using a shorter TTL ensures the cache entry expires sooner.
func (c *StreamingHealthServices) RegisterOptions() cache.RegisterOptions {
	opts := c.RegisterOptionsBlockingRefresh.RegisterOptions()
	opts.LastGetTTL = 20 * time.Minute
	return opts
}

// NewStreamingHealthServices creates a cache-type for watching for service
// health results via streaming updates.
func NewStreamingHealthServices(deps MaterializerDeps) *StreamingHealthServices {
	return &StreamingHealthServices{deps: deps}
}

type MaterializerDeps struct {
	Client submatview.StreamClient
	Logger hclog.Logger
}

// Fetch service health from the materialized view. If no materialized view
// exists, create one and start it running in a goroutine. The goroutine will
// exit when the cache entry storing the result is expired, the cache will call
// Close on the result.State.
//
// Fetch implements part of the cache.Type interface, and assumes that the
// caller ensures that only a single call to Fetch is running at any time.
func (c *StreamingHealthServices) Fetch(opts cache.FetchOptions, req cache.Request) (cache.FetchResult, error) {
	if opts.LastResult != nil && opts.LastResult.State != nil {
		return opts.LastResult.State.(*streamingHealthState).Fetch(opts)
	}

	srvReq := req.(*structs.ServiceSpecificRequest)
	newReqFn := func(index uint64) pbsubscribe.SubscribeRequest {
		req := pbsubscribe.SubscribeRequest{
			Topic:      pbsubscribe.Topic_ServiceHealth,
			Key:        srvReq.ServiceName,
			Token:      srvReq.Token,
			Datacenter: srvReq.Datacenter,
			Index:      index,
			Namespace:  srvReq.EnterpriseMeta.NamespaceOrEmpty(),
		}
		if srvReq.Connect {
			req.Topic = pbsubscribe.Topic_ServiceHealthConnect
		}
		return req
	}

	materializer, err := newMaterializer(c.deps, newReqFn, srvReq)
	if err != nil {
		return cache.FetchResult{}, err
	}
	ctx, cancel := context.WithCancel(context.TODO())
	go materializer.Run(ctx)

	state := &streamingHealthState{
		materializer: materializer,
		done:         ctx.Done(),
		cancel:       cancel,
	}
	return state.Fetch(opts)
}

func newMaterializer(
	deps MaterializerDeps,
	newRequestFn func(uint64) pbsubscribe.SubscribeRequest,
	req *structs.ServiceSpecificRequest,
) (*submatview.Materializer, error) {
	view, err := newHealthView(req)
	if err != nil {
		return nil, err
	}
	return submatview.NewMaterializer(submatview.Deps{
		View:   view,
		Client: deps.Client,
		Logger: deps.Logger,
		Waiter: &retry.Waiter{
			MinFailures: 1,
			// Start backing off with small increments (200-400ms) which will double
			// each attempt. (200-400, 400-800, 800-1600, 1600-3200, 3200-6000, 6000
			// after that). (retry.Wait applies Max limit after jitter right now).
			Factor:  200 * time.Millisecond,
			MinWait: 0,
			MaxWait: 60 * time.Second,
			Jitter:  retry.NewJitter(100),
		},
		Request: newRequestFn,
	}), nil
}

// streamingHealthState wraps a Materializer to manage its lifecycle, and to
// add itself to the FetchResult.State.
type streamingHealthState struct {
	materializer *submatview.Materializer
	done         <-chan struct{}
	cancel       func()
}

func (s *streamingHealthState) Close() error {
	s.cancel()
	return nil
}

func (s *streamingHealthState) Fetch(opts cache.FetchOptions) (cache.FetchResult, error) {
	result, err := s.materializer.Fetch(s.done, opts)
	result.State = s
	return result, err
}

func newHealthView(req *structs.ServiceSpecificRequest) (*healthView, error) {
	fe, err := newFilterEvaluator(req)
	if err != nil {
		return nil, err
	}
	return &healthView{
		state:  make(map[string]structs.CheckServiceNode),
		filter: fe,
	}, nil
}

// healthView implements submatview.View for storing the view state
// of a service health result. We store it as a map to make updates and
// deletions a little easier but we could just store a result type
// (IndexedCheckServiceNodes) and update it in place for each event - that
// involves re-sorting each time etc. though.
type healthView struct {
	state  map[string]structs.CheckServiceNode
	filter filterEvaluator
}

// Update implements View
func (s *healthView) Update(events []*pbsubscribe.Event) error {
	for _, event := range events {
		serviceHealth := event.GetServiceHealth()
		if serviceHealth == nil {
			return fmt.Errorf("unexpected event type for service health view: %T",
				event.GetPayload())
		}

		id := serviceHealth.CheckServiceNode.UniqueID()
		switch serviceHealth.Op {
		case pbsubscribe.CatalogOp_Register:
			csn := *pbservice.CheckServiceNodeToStructs(serviceHealth.CheckServiceNode)
			passed, err := s.filter.Evaluate(csn)
			switch {
			case err != nil:
				return err
			case passed:
				s.state[id] = csn
			}

		case pbsubscribe.CatalogOp_Deregister:
			delete(s.state, id)
		}
	}
	return nil
}

type filterEvaluator interface {
	Evaluate(datum interface{}) (bool, error)
}

func newFilterEvaluator(req *structs.ServiceSpecificRequest) (filterEvaluator, error) {
	var evaluators []filterEvaluator

	typ := reflect.TypeOf(structs.CheckServiceNode{})
	if req.Filter != "" {
		e, err := bexpr.CreateEvaluatorForType(req.Filter, nil, typ)
		if err != nil {
			return nil, err
		}
		evaluators = append(evaluators, e)
	}

	if req.ServiceTag != "" {
		// Handle backwards compat with old field
		req.ServiceTags = []string{req.ServiceTag}
	}

	if req.TagFilter && len(req.ServiceTags) > 0 {
		evaluators = append(evaluators, serviceTagEvaluator{tags: req.ServiceTags})
	}

	for key, value := range req.NodeMetaFilters {
		expr := fmt.Sprintf(`"%s" in Node.Meta.%s`, value, key)
		e, err := bexpr.CreateEvaluatorForType(expr, nil, typ)
		if err != nil {
			return nil, err
		}
		evaluators = append(evaluators, e)
	}

	switch len(evaluators) {
	case 0:
		return noopFilterEvaluator{}, nil
	case 1:
		return evaluators[0], nil
	default:
		return &multiFilterEvaluator{evaluators: evaluators}, nil
	}
}

// noopFilterEvaluator may be used in place of a bexpr.Evaluator. The Evaluate
// method always return true, so no items will be filtered out.
type noopFilterEvaluator struct{}

func (noopFilterEvaluator) Evaluate(_ interface{}) (bool, error) {
	return true, nil
}

type multiFilterEvaluator struct {
	evaluators []filterEvaluator
}

func (m multiFilterEvaluator) Evaluate(data interface{}) (bool, error) {
	for _, e := range m.evaluators {
		match, err := e.Evaluate(data)
		if !match || err != nil {
			return match, err
		}
	}
	return true, nil
}

// sortCheckServiceNodes sorts the results to match memdb semantics
// Sort results by Node.Node, if 2 instances match, order by Service.ID
// Will allow result to be stable sorted and match queries without cache
func sortCheckServiceNodes(serviceNodes *structs.IndexedCheckServiceNodes) {
	sort.SliceStable(serviceNodes.Nodes, func(i, j int) bool {
		left := serviceNodes.Nodes[i]
		right := serviceNodes.Nodes[j]
		if left.Node.Node == right.Node.Node {
			return left.Service.ID < right.Service.ID
		}
		return left.Node.Node < right.Node.Node
	})
}

// Result returns the structs.IndexedCheckServiceNodes stored by this view.
func (s *healthView) Result(index uint64) (interface{}, error) {
	result := structs.IndexedCheckServiceNodes{
		Nodes: make(structs.CheckServiceNodes, 0, len(s.state)),
		QueryMeta: structs.QueryMeta{
			Index: index,
		},
	}
	for _, node := range s.state {
		result.Nodes = append(result.Nodes, node)
	}
	sortCheckServiceNodes(&result)

	return &result, nil
}

func (s *healthView) Reset() {
	s.state = make(map[string]structs.CheckServiceNode)
}

// serviceTagEvaluator implements the filterEvaluator to perform filtering
// by service tags. bexpr can not be used at this time, because the filtering
// must be case insensitive for backwards compatibility. In the future this
// may be replaced with bexpr once case insensitive support is added.
type serviceTagEvaluator struct {
	tags []string
}

func (m serviceTagEvaluator) Evaluate(data interface{}) (bool, error) {
	csn, ok := data.(structs.CheckServiceNode)
	if !ok {
		return false, fmt.Errorf("unexpected type %T for structs.CheckServiceNode filter", data)
	}
	for _, tag := range m.tags {
		if !serviceHasTag(csn.Service, tag) {
			// If any one of the expected tags was not found, filter the service
			return false, nil
		}
	}
	return true, nil
}

func serviceHasTag(sn *structs.NodeService, tag string) bool {
	for _, t := range sn.Tags {
		if strings.EqualFold(t, tag) {
			return true
		}
	}
	return false
}
<|endoftext|>