<commit_before>package softlayer

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"log"
	"strings"

	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/common/uuid"
	"github.com/mitchellh/packer/packer"
	"golang.org/x/crypto/ssh"
)

type stepCreateSshKey struct {
	keyId          int64
	PrivateKeyFile string
}

func (self *stepCreateSshKey) Run(state multistep.StateBag) multistep.StepAction {
	ui := state.Get("ui").(packer.Ui)
	if self.PrivateKeyFile != "" {
		ui.Say(fmt.Sprintf("Reading private key file (%s)...", self.PrivateKeyFile))

		privateKeyBytes, err := ioutil.ReadFile(self.PrivateKeyFile)
		if err != nil {
			state.Put("error", fmt.Errorf("Error loading configured private key file: %s", err))
			return multistep.ActionHalt
		}

		state.Put("ssh_key_id", self.keyId)
		state.Put("ssh_private_key", string(privateKeyBytes))

		return multistep.ActionContinue
	}

	client := state.Get("client").(*SoftlayerClient)
	ui.Say("Creating temporary ssh key for the instance...")

	rsaKey, err := rsa.GenerateKey(rand.Reader, 2014)
	if err != nil {
		ui.Error(err.Error())
		state.Put("error", err)
		return multistep.ActionHalt
	}

	// ASN.1 DER encoded form
	privDer := x509.MarshalPKCS1PrivateKey(rsaKey)
	privBlk := pem.Block{
		Type:    "RSA PRIVATE KEY",
		Headers: nil,
		Bytes:   privDer,
	}

	// Set the private key in the statebag for later
	state.Put("ssh_private_key", string(pem.EncodeToMemory(&privBlk)))

	pub, err := ssh.NewPublicKey(&rsaKey.PublicKey)
	if err != nil {
		ui.Error(err.Error())
		state.Put("error", err)
		return multistep.ActionHalt
	}

	publicKey := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(pub)))

	// The name of the public key
	label := fmt.Sprintf("packer-%s", uuid.TimeOrderedUUID())
	keyId, err := client.UploadSshKey(label, publicKey)
	if err != nil {
		ui.Error(err.Error())
		state.Put("error", err)
		return multistep.ActionHalt
	}

	self.keyId = keyId
	state.Put("ssh_key_id", keyId)

	ui.Say(fmt.Sprintf("Created SSH key with id '%d'", keyId))

	return multistep.ActionContinue
}

func (self *stepCreateSshKey) Cleanup(state multistep.StateBag) {
	// If no key name is set, then we never created it, so just return
	if self.keyId == 0 {
		return
	}

	client := state.Get("client").(*SoftlayerClient)
	ui := state.Get("ui").(packer.Ui)

	ui.Say("Deleting temporary ssh key...")
	err := client.DestroySshKey(self.keyId)

	if err != nil {
		log.Printf("Error cleaning up ssh key: %v", err.Error())
		ui.Error(fmt.Sprintf("Error cleaning up ssh key. Please delete the key (%d) manually", self.keyId))
	}
}
<commit_msg>builder: fix use of existing private key<commit_after>package softlayer

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/common/uuid"
	"github.com/mitchellh/packer/packer"
	"golang.org/x/crypto/ssh"
)

type stepCreateSshKey struct {
	keyId          int64
	temporary      bool
	PrivateKeyFile string
}

func (self *stepCreateSshKey) Run(state multistep.StateBag) multistep.StepAction {
	ui := state.Get("ui").(packer.Ui)
	client := state.Get("client").(*SoftlayerClient)
	if self.PrivateKeyFile != "" {
		ui.Say(fmt.Sprintf("Reading private key file (%s)...", self.PrivateKeyFile))

		privateKeyBytes, err := ioutil.ReadFile(self.PrivateKeyFile)
		if err != nil {
			state.Put("error", fmt.Errorf("Error loading configured private key file: %s", err))
			return multistep.ActionHalt
		}

		if self.keyId == 0 {
			key, err := ssh.ParseRawPrivateKey(privateKeyBytes)
			if err != nil {
				return self.error(state, ui, err)
			}

			rsaKey, ok := key.(*rsa.PrivateKey)
			if !ok {
				return self.error(state, ui, errors.New("the private key is not RSA one"))
			}

			keyId, err := self.uploadSshKey(client, rsaKey)
			if err != nil {
				return self.error(state, ui, err)
			}

			self.keyId = keyId
		} else {
			ui.Say(fmt.Sprintf("Attaching existing sshkey ID to the instance (%d)...", self.keyId))
		}

		state.Put("ssh_key_id", self.keyId)
		state.Put("ssh_private_key", string(privateKeyBytes))

		return multistep.ActionContinue
	}

	ui.Say("Creating temporary ssh key for the instance...")

	rsaKey, err := rsa.GenerateKey(rand.Reader, 2014)
	if err != nil {
		return self.error(state, ui, err)
	}

	// ASN.1 DER encoded form
	privDer := x509.MarshalPKCS1PrivateKey(rsaKey)
	privBlk := pem.Block{
		Type:    "RSA PRIVATE KEY",
		Headers: nil,
		Bytes:   privDer,
	}

	keyId, err := self.uploadSshKey(client, rsaKey)
	if err != nil {
		return self.error(state, ui, err)
	}

	self.temporary = true
	self.keyId = keyId

	// Set the private key in the statebag for later
	state.Put("ssh_private_key", string(pem.EncodeToMemory(&privBlk)))
	state.Put("ssh_key_id", keyId)

	ui.Say(fmt.Sprintf("Created SSH key with id '%d'", keyId))

	return multistep.ActionContinue
}

func (self *stepCreateSshKey) Cleanup(state multistep.StateBag) {
	if !self.temporary {
		return
	}

	client := state.Get("client").(*SoftlayerClient)
	ui := state.Get("ui").(packer.Ui)

	ui.Say("Deleting temporary ssh key...")
	err := client.DestroySshKey(self.keyId)

	if err != nil {
		self.error(nil, ui, fmt.Errorf("Error cleaning up ssh key. Please delete the key (%d) manually", self.keyId))
	}
}

func (self *stepCreateSshKey) uploadSshKey(client *SoftlayerClient, rsaKey *rsa.PrivateKey) (keyId int64, err error) {
	pub, err := ssh.NewPublicKey(&rsaKey.PublicKey)
	if err != nil {
		return 0, err
	}

	publicKey := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(pub)))

	// The name of the public key
	label := fmt.Sprintf("packer-%s", uuid.TimeOrderedUUID())
	keyId, err = client.UploadSshKey(label, publicKey)
	if err != nil {
		return 0, err
	}

	return keyId, nil
}

func (self *stepCreateSshKey) error(state multistep.StateBag, ui packer.Ui, err error) multistep.StepAction {
	if ui != nil {
		ui.Error(err.Error())
	}
	if state != nil {
		state.Put("error", err)
	}
	return multistep.ActionHalt
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"

	"github.com/smartystreets/go-disruptor"
)

const Mod = 1000000 * 10 // 1 million * N

func consume(reader *disruptor.Reader) {
	// runtime.LockOSThread()
	started := time.Now()

	for {
		sequence, remaining := reader.Receive()
		if remaining >= 0 {
			for remaining >= 0 {

				if sequence%Mod == 0 {
					finished := time.Now()
					fmt.Println(sequence, finished.Sub(started))
					started = time.Now()
				}

				// if sequence != ringBuffer[sequence&RingMask] {
				// 	message := ringBuffer[sequence&RingMask]
				// 	panic(fmt.Sprintf("Sequence: %d, Message %d\n", sequence, message))
				// }

				remaining--
				sequence++
			}

			reader.Commit(sequence)

		} else {

		}
	}
}
<commit_msg>Producing a single item has huge effect on latency.<commit_after>package main

import (
	"fmt"
	"time"

	"github.com/smartystreets/go-disruptor"
)

const Mod = 1000000 * 10 // 1 million * N

func consume(reader *disruptor.Reader) {
	// runtime.LockOSThread()
	started := time.Now()

	for {
		sequence, remaining := reader.Receive()
		if remaining >= 0 {
			for remaining >= 0 {

				if sequence%Mod == 0 {
					finished := time.Now()
					fmt.Println(sequence, finished.Sub(started))
					started = time.Now()
				}

				if sequence != ringBuffer[sequence&RingMask] {
					message := ringBuffer[sequence&RingMask]
					panic(fmt.Sprintf("Sequence: %d, Message %d\n", sequence, message))
				}

				remaining--
				sequence++
			}

			reader.Commit(sequence)

		} else {

		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/jameycribbs/hare"
)

func main() {
	// Open the database and return a handle to it.
	db, err := hare.OpenDB("./data")
	if err != nil {
		panic(err)
	}
	defer db.Close()

	// Here is how to check if a table exists in your database.
	if !db.TableExists("contacts") {
		fmt.Println("Table 'contacts' does not exist!")
	}

	// Here is how to create a new table in the database and get back a
	// handle to it.
	tbl, err := db.CreateTable("contacts")
	if err != nil {
		panic(err)
	}

	// If the table already exists, you can get a handle to it by
	// calling GetTable.
	tbl, err = db.GetTable("contacts")
	if err != nil {
		panic(err)
	}

	fmt.Println("Table handle:", tbl)

	// Here is how to drop a table.
	err = db.DropTable("contacts")
	if err != nil {
		panic(err)
	}
}
<commit_msg>fixed errors in dbadmin.go example<commit_after>package main

import (
	"fmt"

	"github.com/jameycribbs/hare"
	"github.com/jameycribbs/hare/datastores/disk"
)

func main() {
	ds, err := disk.New("./data", ".json")
	if err != nil {
		panic(err)
	}

	db, err := hare.New(ds)
	if err != nil {
		panic(err)
	}
	defer db.Close()

	// Here is how to create a new table in the database.
	err = db.CreateTable("contacts")
	if err != nil {
		panic(err)
	}

	// Here is how to check if a table exists in your database.
	if !db.TableExists("contacts") {
		fmt.Println("Table 'contacts' does not exist!")
	}

	// Here is how to drop a table.
	err = db.DropTable("contacts")
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Intel Corporation.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package nat

import (
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/intel-go/yanff/common"
	"github.com/intel-go/yanff/flow"
	"github.com/intel-go/yanff/packet"
)

// Tuple is a pair of address and port.
type Tuple struct {
	addr uint32
	port uint16
}

func (t *Tuple) String() string {
	return fmt.Sprintf("addr = %d.%d.%d.%d:%d",
		(t.addr>>24)&0xff,
		(t.addr>>16)&0xff,
		(t.addr>>8)&0xff,
		t.addr&0xff,
		t.port)
}

var (
	// Main lookup table which contains entries
	pri2pubTable []sync.Map
	pub2priTable []sync.Map
	mutex        sync.Mutex

	emptyEntry = Tuple{addr: 0, port: 0}

	// Debug variables
	debugDump = false
	fdump     []*os.File
)

func init() {
	pri2pubTable = make([]sync.Map, common.UDPNumber+1)
	pub2priTable = make([]sync.Map, common.UDPNumber+1)
}

func allocateNewEgressConnection(protocol uint8, privEntry Tuple, publicAddr uint32) (Tuple, error) {
	mutex.Lock()

	port, err := allocNewPort(protocol)
	if err != nil {
		mutex.Unlock()
		return Tuple{}, err
	}

	pubEntry := Tuple{
		addr: publicAddr,
		port: uint16(port),
	}

	portmap[protocol][port] = portMapEntry{
		lastused:             time.Now(),
		addr:                 publicAddr,
		finCount:             0,
		terminationDirection: 0,
	}

	// Add lookup entries for packet translation
	pri2pubTable[protocol].Store(privEntry, pubEntry)
	pub2priTable[protocol].Store(pubEntry, privEntry)

	mutex.Unlock()
	return pubEntry, nil
}

func dumpInput(pkt *packet.Packet, index int) {
	if debugDump {
		// Dump input packet
		if fdump[index] == nil {
			fdump[index], _ = os.Create(fmt.Sprintf("%ddump.pcap", index))
			packet.WritePcapGlobalHdr(fdump[index])
			pkt.WritePcapOnePacket(fdump[index])
		}

		pkt.WritePcapOnePacket(fdump[index])
	}
}

func dumpOutput(pkt *packet.Packet, index int) {
	if debugDump {
		pkt.WritePcapOnePacket(fdump[index])
	}
}

// PublicToPrivateTranslation does ingress translation.
func PublicToPrivateTranslation(pkt *packet.Packet, ctx flow.UserContext) uint {
	pi := ctx.(pairIndex)

	dumpInput(pkt, pi.index)

	// Parse packet type and address
	pkt.ParseL3()
	pktIPv4 := pkt.GetIPv4()
	if pktIPv4 == nil {
		arp := pkt.GetARP()
		if arp != nil {
			port := handleARP(pkt, &Natconfig.PortPairs[pi.index].PublicPort)
			if port != flowDrop {
				dumpOutput(pkt, pi.index)
			}
			return port
		}
		// We don't currently support anything except for IPv4 and ARP
		return flowDrop
	}

	pktTCP, pktUDP, pktICMP := pkt.ParseAllKnownL4ForIPv4()
	// Create a lookup key
	protocol := pktIPv4.NextProtoID
	pub2priKey := Tuple{
		addr: packet.SwapBytesUint32(pktIPv4.DstAddr),
	}
	// Parse packet destination port
	if pktTCP != nil {
		pub2priKey.port = packet.SwapBytesUint16(pktTCP.DstPort)
	} else if pktUDP != nil {
		pub2priKey.port = packet.SwapBytesUint16(pktUDP.DstPort)
	} else if pktICMP != nil {
		pub2priKey.port = pktICMP.Identifier
	} else {
		return flowDrop
	}

	// Do lookup
	v, found := pub2priTable[protocol].Load(pub2priKey)
	// For ingress connections packets are allowed only if a
	// connection has been previosly established with a egress
	// (private to public) packet. So if lookup fails, this incoming
	// packet is ignored.
	if !found {
		return flowDrop
	}
	value := v.(Tuple)

	// Check whether connection is too old
	if portmap[protocol][pub2priKey.port].lastused.Add(connectionTimeout).After(time.Now()) {
		portmap[protocol][pub2priKey.port].lastused = time.Now()
	} else {
		// There was no transfer on this port for too long
		// time. We don't allow it any more
		mutex.Lock()
		deleteOldConnection(protocol, int(pub2priKey.port))
		mutex.Unlock()
		return flowDrop
	}

	// Check whether TCP connection could be reused
	if protocol == common.TCPNumber {
		checkTCPTermination(pktTCP, int(pub2priKey.port), pub2pri)
	}

	// Do packet translation
	pkt.Ether.DAddr = getMACForIP(&Natconfig.PortPairs[pi.index].PrivatePort, value.addr)
	pkt.Ether.SAddr = Natconfig.PortPairs[pi.index].PrivatePort.SrcMACAddress
	pktIPv4.DstAddr = packet.SwapBytesUint32(value.addr)

	if pktTCP != nil {
		pktTCP.DstPort = packet.SwapBytesUint16(value.port)
		setIPv4TCPChecksum(pktIPv4, pktTCP, CalculateChecksum, HWTXChecksum)
	} else if pktUDP != nil {
		pktUDP.DstPort = packet.SwapBytesUint16(value.port)
		setIPv4UDPChecksum(pktIPv4, pktUDP, CalculateChecksum, HWTXChecksum)
	} else {
		setIPv4ICMPChecksum(pktIPv4, pktICMP, CalculateChecksum, HWTXChecksum)
	}

	dumpOutput(pkt, pi.index)
	return flowOut
}

// PrivateToPublicTranslation does egress translation.
func PrivateToPublicTranslation(pkt *packet.Packet, ctx flow.UserContext) uint {
	pi := ctx.(pairIndex)

	dumpInput(pkt, pi.index)

	// Parse packet type and address
	pkt.ParseL3()
	pktIPv4 := pkt.GetIPv4()
	if pktIPv4 == nil {
		arp := pkt.GetARP()
		if arp != nil {
			port := handleARP(pkt, &Natconfig.PortPairs[pi.index].PrivatePort)
			if port != flowDrop {
				dumpOutput(pkt, pi.index)
			}
			return port
		}
		// We don't currently support anything except for IPv4 and ARP
		return flowDrop
	}

	pktTCP, pktUDP, pktICMP := pkt.ParseAllKnownL4ForIPv4()

	// Create a lookup key
	protocol := pktIPv4.NextProtoID
	pri2pubKey := Tuple{
		addr: packet.SwapBytesUint32(pktIPv4.SrcAddr),
	}

	// Parse packet source port
	if pktTCP != nil {
		pri2pubKey.port = packet.SwapBytesUint16(pktTCP.SrcPort)
	} else if pktUDP != nil {
		pri2pubKey.port = packet.SwapBytesUint16(pktUDP.SrcPort)
	} else if pktICMP != nil {
		pri2pubKey.port = pktICMP.Identifier
	} else {
		return flowDrop
	}

	// Do lookup
	var value Tuple
	v, found := pri2pubTable[protocol].Load(pri2pubKey)
	if !found {
		var err error
		// Store new local network entry in ARP cache
		Natconfig.PortPairs[pi.index].PrivatePort.ArpTable.Store(pri2pubKey.addr, pkt.Ether.SAddr)
		// Allocate new connection from private to public network
		value, err = allocateNewEgressConnection(protocol, pri2pubKey,
			Natconfig.PortPairs[pi.index].PublicPort.Subnet.Addr)

		if err != nil {
			println("Warning! Failed to allocate new connection", err)
			return flowDrop
		}
	} else {
		value = v.(Tuple)
		portmap[protocol][value.port].lastused = time.Now()
	}

	// Check whether TCP connection could be reused
	if pktTCP != nil {
		checkTCPTermination(pktTCP, int(value.port), pri2pub)
	}

	// Do packet translation
	pkt.Ether.DAddr = Natconfig.PortPairs[pi.index].PublicPort.DstMACAddress
	pkt.Ether.SAddr = Natconfig.PortPairs[pi.index].PublicPort.SrcMACAddress
	pktIPv4.SrcAddr = packet.SwapBytesUint32(value.addr)

	if pktTCP != nil {
		pktTCP.SrcPort = packet.SwapBytesUint16(value.port)
		setIPv4TCPChecksum(pktIPv4, pktTCP, CalculateChecksum, HWTXChecksum)
	} else if pktUDP != nil {
		pktUDP.SrcPort = packet.SwapBytesUint16(value.port)
		setIPv4UDPChecksum(pktIPv4, pktUDP, CalculateChecksum, HWTXChecksum)
	} else {
		setIPv4ICMPChecksum(pktIPv4, pktICMP, CalculateChecksum, HWTXChecksum)
	}

	dumpOutput(pkt, pi.index)
	return flowOut
}

// Simple check for FIN or RST in TCP
func checkTCPTermination(hdr *packet.TCPHdr, port int, dir terminationDirection) {
	if hdr.TCPFlags&common.TCPFlagFin != 0 {
		// First check for FIN
		mutex.Lock()

		pme := &portmap[common.TCPNumber][port]
		if pme.finCount == 0 {
			pme.finCount = 1
			pme.terminationDirection = dir
		} else if pme.finCount == 1 && pme.terminationDirection == ^dir {
			pme.finCount = 2
		}

		mutex.Unlock()
	} else if hdr.TCPFlags&common.TCPFlagRst != 0 {
		// RST means that connection is terminated immediately
		mutex.Lock()
		deleteOldConnection(common.TCPNumber, port)
		mutex.Unlock()
	} else if hdr.TCPFlags&common.TCPFlagAck != 0 {
		// Check for ACK last so that if there is also FIN,
		// termination doesn't happen. Last ACK should come without
		// FIN
		mutex.Lock()

		pme := &portmap[common.TCPNumber][port]
		if pme.finCount == 2 {
			deleteOldConnection(common.TCPNumber, port)
		}

		mutex.Unlock()
	}
}

func handleARP(pkt *packet.Packet, port *ipv4Port) uint {
	arp := pkt.GetARP()

	if packet.SwapBytesUint16(arp.Operation) != packet.ARPRequest {
		// We don't care about replies so far
		return flowDrop
	}

	// Check that someone is asking about MAC of my IP address and HW
	// address is blank in request
	if packet.BytesToIPv4(arp.TPA[0], arp.TPA[1], arp.TPA[2], arp.TPA[3]) != packet.SwapBytesUint32(port.Subnet.Addr) ||
		arp.THA != [common.EtherAddrLen]byte{} {
		return flowDrop
	}

	// Prepare an answer to this request
	pkt.Ether.DAddr = pkt.Ether.SAddr
	pkt.Ether.SAddr = port.SrcMACAddress
	arp.Operation = packet.SwapBytesUint16(packet.ARPReply)
	myIP := arp.TPA
	arp.TPA = arp.SPA
	arp.THA = arp.SHA
	arp.SPA = myIP
	arp.SHA = port.SrcMACAddress

	return flowArp
}

func getMACForIP(port *ipv4Port, ip uint32) macAddress {
	v, found := port.ArpTable.Load(ip)
	if found {
		return macAddress(v.([common.EtherAddrLen]byte))
	}
	println("Warning! IP address",
		byte(ip), ".", byte(ip>>8), ".", byte(ip>>16), ".", byte(ip>>24),
		"not found in ARP cache on port", port.Index)
	return macAddress{0xff, 0xff, 0xff, 0xff, 0xff, 0xff}
}
<commit_msg>Fixed NAT-ing of ICMP packets. Pings now pass through one subnetwork to another<commit_after>// Copyright 2017 Intel Corporation.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package nat

import (
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/intel-go/yanff/common"
	"github.com/intel-go/yanff/flow"
	"github.com/intel-go/yanff/packet"
)

// Tuple is a pair of address and port.
type Tuple struct {
	addr uint32
	port uint16
}

func (t *Tuple) String() string {
	return fmt.Sprintf("addr = %d.%d.%d.%d:%d",
		(t.addr>>24)&0xff,
		(t.addr>>16)&0xff,
		(t.addr>>8)&0xff,
		t.addr&0xff,
		t.port)
}

var (
	// Main lookup table which contains entries
	pri2pubTable []sync.Map
	pub2priTable []sync.Map
	mutex        sync.Mutex

	emptyEntry = Tuple{addr: 0, port: 0}

	// Debug variables
	debugDump = false
	fdump     []*os.File
)

func init() {
	pri2pubTable = make([]sync.Map, common.UDPNumber+1)
	pub2priTable = make([]sync.Map, common.UDPNumber+1)
}

func allocateNewEgressConnection(protocol uint8, privEntry Tuple, publicAddr uint32) (Tuple, error) {
	mutex.Lock()

	port, err := allocNewPort(protocol)
	if err != nil {
		mutex.Unlock()
		return Tuple{}, err
	}

	pubEntry := Tuple{
		addr: publicAddr,
		port: uint16(port),
	}

	portmap[protocol][port] = portMapEntry{
		lastused:             time.Now(),
		addr:                 publicAddr,
		finCount:             0,
		terminationDirection: 0,
	}

	// Add lookup entries for packet translation
	pri2pubTable[protocol].Store(privEntry, pubEntry)
	pub2priTable[protocol].Store(pubEntry, privEntry)

	mutex.Unlock()
	return pubEntry, nil
}

func dumpInput(pkt *packet.Packet, index int) {
	if debugDump {
		// Dump input packet
		if fdump[index] == nil {
			fdump[index], _ = os.Create(fmt.Sprintf("%ddump.pcap", index))
			packet.WritePcapGlobalHdr(fdump[index])
			pkt.WritePcapOnePacket(fdump[index])
		}

		pkt.WritePcapOnePacket(fdump[index])
	}
}

func dumpOutput(pkt *packet.Packet, index int) {
	if debugDump {
		pkt.WritePcapOnePacket(fdump[index])
	}
}

// PublicToPrivateTranslation does ingress translation.
func PublicToPrivateTranslation(pkt *packet.Packet, ctx flow.UserContext) uint {
	pi := ctx.(pairIndex)

	dumpInput(pkt, pi.index)

	// Parse packet type and address
	pkt.ParseL3()
	pktIPv4 := pkt.GetIPv4()
	if pktIPv4 == nil {
		arp := pkt.GetARP()
		if arp != nil {
			port := handleARP(pkt, &Natconfig.PortPairs[pi.index].PublicPort)
			if port != flowDrop {
				dumpOutput(pkt, pi.index)
			}
			return port
		}
		// We don't currently support anything except for IPv4 and ARP
		return flowDrop
	}

	pktTCP, pktUDP, pktICMP := pkt.ParseAllKnownL4ForIPv4()
	// Create a lookup key
	protocol := pktIPv4.NextProtoID
	pub2priKey := Tuple{
		addr: packet.SwapBytesUint32(pktIPv4.DstAddr),
	}
	// Parse packet destination port
	if pktTCP != nil {
		pub2priKey.port = packet.SwapBytesUint16(pktTCP.DstPort)
	} else if pktUDP != nil {
		pub2priKey.port = packet.SwapBytesUint16(pktUDP.DstPort)
	} else if pktICMP != nil {
		pub2priKey.port = packet.SwapBytesUint16(pktICMP.Identifier)
	} else {
		return flowDrop
	}

	// Do lookup
	v, found := pub2priTable[protocol].Load(pub2priKey)
	// For ingress connections packets are allowed only if a
	// connection has been previosly established with a egress
	// (private to public) packet. So if lookup fails, this incoming
	// packet is ignored.
	if !found {
		return flowDrop
	}
	value := v.(Tuple)

	// Check whether connection is too old
	if portmap[protocol][pub2priKey.port].lastused.Add(connectionTimeout).After(time.Now()) {
		portmap[protocol][pub2priKey.port].lastused = time.Now()
	} else {
		// There was no transfer on this port for too long
		// time. We don't allow it any more
		mutex.Lock()
		deleteOldConnection(protocol, int(pub2priKey.port))
		mutex.Unlock()
		return flowDrop
	}

	// Check whether TCP connection could be reused
	if protocol == common.TCPNumber {
		checkTCPTermination(pktTCP, int(pub2priKey.port), pub2pri)
	}

	// Do packet translation
	pkt.Ether.DAddr = getMACForIP(&Natconfig.PortPairs[pi.index].PrivatePort, value.addr)
	pkt.Ether.SAddr = Natconfig.PortPairs[pi.index].PrivatePort.SrcMACAddress
	pktIPv4.DstAddr = packet.SwapBytesUint32(value.addr)

	if pktTCP != nil {
		pktTCP.DstPort = packet.SwapBytesUint16(value.port)
		setIPv4TCPChecksum(pktIPv4, pktTCP, CalculateChecksum, HWTXChecksum)
	} else if pktUDP != nil {
		pktUDP.DstPort = packet.SwapBytesUint16(value.port)
		setIPv4UDPChecksum(pktIPv4, pktUDP, CalculateChecksum, HWTXChecksum)
	} else {
		pktICMP.Identifier = packet.SwapBytesUint16(value.port)
		setIPv4ICMPChecksum(pktIPv4, pktICMP, CalculateChecksum, HWTXChecksum)
	}

	dumpOutput(pkt, pi.index)
	return flowOut
}

// PrivateToPublicTranslation does egress translation.
func PrivateToPublicTranslation(pkt *packet.Packet, ctx flow.UserContext) uint {
	pi := ctx.(pairIndex)

	dumpInput(pkt, pi.index)

	// Parse packet type and address
	pkt.ParseL3()
	pktIPv4 := pkt.GetIPv4()
	if pktIPv4 == nil {
		arp := pkt.GetARP()
		if arp != nil {
			port := handleARP(pkt, &Natconfig.PortPairs[pi.index].PrivatePort)
			if port != flowDrop {
				dumpOutput(pkt, pi.index)
			}
			return port
		}
		// We don't currently support anything except for IPv4 and ARP
		return flowDrop
	}

	pktTCP, pktUDP, pktICMP := pkt.ParseAllKnownL4ForIPv4()

	// Create a lookup key
	protocol := pktIPv4.NextProtoID
	pri2pubKey := Tuple{
		addr: packet.SwapBytesUint32(pktIPv4.SrcAddr),
	}

	// Parse packet source port
	if pktTCP != nil {
		pri2pubKey.port = packet.SwapBytesUint16(pktTCP.SrcPort)
	} else if pktUDP != nil {
		pri2pubKey.port = packet.SwapBytesUint16(pktUDP.SrcPort)
	} else if pktICMP != nil {
		pri2pubKey.port = packet.SwapBytesUint16(pktICMP.Identifier)
	} else {
		return flowDrop
	}

	// Do lookup
	var value Tuple
	v, found := pri2pubTable[protocol].Load(pri2pubKey)
	if !found {
		var err error
		// Store new local network entry in ARP cache
		Natconfig.PortPairs[pi.index].PrivatePort.ArpTable.Store(pri2pubKey.addr, pkt.Ether.SAddr)
		// Allocate new connection from private to public network
		value, err = allocateNewEgressConnection(protocol, pri2pubKey,
			Natconfig.PortPairs[pi.index].PublicPort.Subnet.Addr)

		if err != nil {
			println("Warning! Failed to allocate new connection", err)
			return flowDrop
		}
	} else {
		value = v.(Tuple)
		portmap[protocol][value.port].lastused = time.Now()
	}

	// Check whether TCP connection could be reused
	if pktTCP != nil {
		checkTCPTermination(pktTCP, int(value.port), pri2pub)
	}

	// Do packet translation
	pkt.Ether.DAddr = Natconfig.PortPairs[pi.index].PublicPort.DstMACAddress
	pkt.Ether.SAddr = Natconfig.PortPairs[pi.index].PublicPort.SrcMACAddress
	pktIPv4.SrcAddr = packet.SwapBytesUint32(value.addr)

	if pktTCP != nil {
		pktTCP.SrcPort = packet.SwapBytesUint16(value.port)
		setIPv4TCPChecksum(pktIPv4, pktTCP, CalculateChecksum, HWTXChecksum)
	} else if pktUDP != nil {
		pktUDP.SrcPort = packet.SwapBytesUint16(value.port)
		setIPv4UDPChecksum(pktIPv4, pktUDP, CalculateChecksum, HWTXChecksum)
	} else {
		pktICMP.Identifier = packet.SwapBytesUint16(value.port)
		setIPv4ICMPChecksum(pktIPv4, pktICMP, CalculateChecksum, HWTXChecksum)
	}

	dumpOutput(pkt, pi.index)
	return flowOut
}

// Simple check for FIN or RST in TCP
func checkTCPTermination(hdr *packet.TCPHdr, port int, dir terminationDirection) {
	if hdr.TCPFlags&common.TCPFlagFin != 0 {
		// First check for FIN
		mutex.Lock()

		pme := &portmap[common.TCPNumber][port]
		if pme.finCount == 0 {
			pme.finCount = 1
			pme.terminationDirection = dir
		} else if pme.finCount == 1 && pme.terminationDirection == ^dir {
			pme.finCount = 2
		}

		mutex.Unlock()
	} else if hdr.TCPFlags&common.TCPFlagRst != 0 {
		// RST means that connection is terminated immediately
		mutex.Lock()
		deleteOldConnection(common.TCPNumber, port)
		mutex.Unlock()
	} else if hdr.TCPFlags&common.TCPFlagAck != 0 {
		// Check for ACK last so that if there is also FIN,
		// termination doesn't happen. Last ACK should come without
		// FIN
		mutex.Lock()

		pme := &portmap[common.TCPNumber][port]
		if pme.finCount == 2 {
			deleteOldConnection(common.TCPNumber, port)
		}

		mutex.Unlock()
	}
}

func handleARP(pkt *packet.Packet, port *ipv4Port) uint {
	arp := pkt.GetARP()

	if packet.SwapBytesUint16(arp.Operation) != packet.ARPRequest {
		// We don't care about replies so far
		return flowDrop
	}

	// Check that someone is asking about MAC of my IP address and HW
	// address is blank in request
	if packet.BytesToIPv4(arp.TPA[0], arp.TPA[1], arp.TPA[2], arp.TPA[3]) != packet.SwapBytesUint32(port.Subnet.Addr) ||
		arp.THA != [common.EtherAddrLen]byte{} {
		return flowDrop
	}

	// Prepare an answer to this request
	pkt.Ether.DAddr = pkt.Ether.SAddr
	pkt.Ether.SAddr = port.SrcMACAddress
	arp.Operation = packet.SwapBytesUint16(packet.ARPReply)
	myIP := arp.TPA
	arp.TPA = arp.SPA
	arp.THA = arp.SHA
	arp.SPA = myIP
	arp.SHA = port.SrcMACAddress

	return flowArp
}

func getMACForIP(port *ipv4Port, ip uint32) macAddress {
	v, found := port.ArpTable.Load(ip)
	if found {
		return macAddress(v.([common.EtherAddrLen]byte))
	}
	println("Warning! IP address",
		byte(ip), ".", byte(ip>>8), ".", byte(ip>>16), ".", byte(ip>>24),
		"not found in ARP cache on port", port.Index)
	return macAddress{0xff, 0xff, 0xff, 0xff, 0xff, 0xff}
}
<|endoftext|>


<commit_before>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// This LevelDB Go implementation is based on LevelDB C++ implementation.
// Which contains the following header:
//   Copyright (c) 2011 The LevelDB Authors. All rights reserved.
//   Use of this source code is governed by a BSD-style license that can be
//   found in the LEVELDBCPP_LICENSE file. See the LEVELDBCPP_AUTHORS file
//   for names of contributors.

// Package memdb provide in-memory key/value database implementation.
package memdb

import (
	"leveldb/comparer"
	"leveldb/errors"
	"math/rand"
	"sync/atomic"
	"unsafe"
)

const tMaxHeight = 12

var (
	mPtrSize  int
	mNodeSize int
)

func init() {
	node := new(mNode)
	mPtrSize = int(unsafe.Sizeof(node))
	mNodeSize = int(unsafe.Sizeof(*node))
}

type mNode struct {
	key   []byte
	value []byte
	next  []unsafe.Pointer
}

func (p *mNode) getNext(n int) *mNode {
	return (*mNode)(atomic.LoadPointer(&p.next[n]))
}

func (p *mNode) setNext(n int, x *mNode) {
	atomic.StorePointer(&p.next[n], unsafe.Pointer(x))
}

func (p *mNode) getNext_NB(n int) *mNode {
	return (*mNode)(p.next[n])
}

func (p *mNode) setNext_NB(n int, x *mNode) {
	p.next[n] = unsafe.Pointer(x)
}

// DB represent an in-memory key/value database.
type DB struct {
	cmp       comparer.BasicComparer
	rnd       *rand.Rand
	head      *mNode
	maxHeight int
	memSize   int64
	n         int32

	prev [tMaxHeight]*mNode
}

// New create new initalized in-memory key/value database.
func New(cmp comparer.BasicComparer) *DB {
	p := &DB{
		cmp:       cmp,
		rnd:       rand.New(rand.NewSource(0xdeadbeef)),
		maxHeight: 1,
	}
	p.head = p.newNode(nil, nil, tMaxHeight)
	return p
}

// Put insert given key and value to the database. Need external synchronization.
func (p *DB) Put(key []byte, value []byte) {
	p.findGE_NB(key, true)

	h := p.randHeight()
	if h > p.maxHeight {
		for i := p.maxHeight; i < h; i++ {
			p.prev[i] = p.head
		}
		p.maxHeight = h
	}

	x := p.newNode(key, value, h)
	for i, n := range p.prev[:h] {
		x.setNext_NB(i, n.getNext_NB(i))
		n.setNext(i, x)
	}

	atomic.AddInt64(&p.memSize, int64(mNodeSize+(mPtrSize*h)+len(key)+len(value)))
	atomic.AddInt32(&p.n, 1)
}

// Remove remove given key from the database. Need external synchronization.
func (p *DB) Remove(key []byte) {
	x := p.findGE_NB(key, true)
	if x == nil || x == p.head || p.cmp.Compare(x.key, key) != 0 {
		return
	}

	h := len(x.next)
	for i, n := range p.prev[:h] {
		n.setNext(i, n.getNext_NB(i).getNext_NB(i))
	}

	atomic.AddInt64(&p.memSize, -int64(mNodeSize+(mPtrSize*h)+len(x.key)+len(x.value)))
	atomic.AddInt32(&p.n, -1)
}

// Contains return true if given key are in database.
func (p *DB) Contains(key []byte) bool {
	x := p.findGE(key, false)
	if x != nil && x != p.head && p.cmp.Compare(x.key, key) == 0 {
		return true
	}
	return false
}

// Get return key/value equal or greater than given key.
func (p *DB) Get(key []byte) (rkey, value []byte, err error) {
	node := p.findGE(key, false)
	if node == nil || node == p.head {
		err = errors.ErrNotFound
		return
	}
	return node.key, node.value, nil
}

// NewIterator create a new iterator over the database content.
func (p *DB) NewIterator() *Iterator {
	return &Iterator{p: p}
}

// Size return approximate size of memory used by the database.
func (p *DB) Size() int {
	return int(atomic.LoadInt64(&p.memSize))
}

// Len return the number of entries in the database.
func (p *DB) Len() int {
	return int(atomic.LoadInt32(&p.n))
}

func (p *DB) newNode(key, value []byte, height int) *mNode {
	return &mNode{key, value, make([]unsafe.Pointer, height)}
}

func (p *DB) findGE(key []byte, prev bool) *mNode {
	x := p.head
	n := p.maxHeight - 1
	for {
		next := x.getNext(n)
		if next != nil && p.cmp.Compare(next.key, key) < 0 {
			// Keep searching in this list
			x = next
		} else {
			if prev {
				p.prev[n] = x
			}
			if n == 0 {
				return next
			}
			n--
		}
	}
	return nil
}

func (p *DB) findGE_NB(key []byte, prev bool) *mNode {
	x := p.head
	n := p.maxHeight - 1
	for {
		next := x.getNext_NB(n)
		if next != nil && p.cmp.Compare(next.key, key) < 0 {
			// Keep searching in this list
			x = next
		} else {
			if prev {
				p.prev[n] = x
			}
			if n == 0 {
				return next
			}
			n--
		}
	}
	return nil
}

func (p *DB) findLT(key []byte) *mNode {
	x := p.head
	n := p.maxHeight - 1
	for {
		next := x.getNext(n)
		if next == nil || p.cmp.Compare(next.key, key) >= 0 {
			if n == 0 {
				return x
			}
			n--
		} else {
			x = next
		}
	}
	return nil
}

func (p *DB) findLast() *mNode {
	x := p.head
	n := p.maxHeight - 1
	for {
		next := x.getNext(n)
		if next == nil {
			if n == 0 {
				return x
			}
			n--
		} else {
			x = next
		}
	}
	return nil
}

func (p *DB) randHeight() int {
	const branching = 4
	n := 1
	for n < tMaxHeight && p.rnd.Int()%branching == 0 {
		n++
	}
	return n
}

type Iterator struct {
	p      *DB
	node   *mNode
	onLast bool
}

func (i *Iterator) Valid() bool {
	return i.node != nil && i.node != i.p.head
}

func (i *Iterator) First() bool {
	i.node = i.p.head.getNext(0)
	return i.Valid()
}

func (i *Iterator) Last() bool {
	i.node = i.p.findLast()
	return i.Valid()
}

func (i *Iterator) Seek(key []byte) (r bool) {
	i.node = i.p.findGE(key, false)
	return i.Valid()
}

func (i *Iterator) Next() bool {
	if i.node == nil {
		return i.First()
	}
	i.node = i.node.getNext(0)
	res := i.Valid()
	if !res {
		i.onLast = true
	}
	return res
}

func (i *Iterator) Prev() bool {
	if i.node == nil {
		if i.onLast {
			return i.Last()
		}
		return false
	}
	i.node = i.p.findLT(i.node.key)
	if i.node == i.p.head {
		i.node = nil
	}
	return i.Valid()
}

func (i *Iterator) Key() []byte {
	if !i.Valid() {
		return nil
	}
	return i.node.key
}

func (i *Iterator) Value() []byte {
	if !i.Valid() {
		return nil
	}
	return i.node.value
}

func (i *Iterator) Error() error { return nil }
<commit_msg>memdb: fix data race<commit_after>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// This LevelDB Go implementation is based on LevelDB C++ implementation.
// Which contains the following header:
//   Copyright (c) 2011 The LevelDB Authors. All rights reserved.
//   Use of this source code is governed by a BSD-style license that can be
//   found in the LEVELDBCPP_LICENSE file. See the LEVELDBCPP_AUTHORS file
//   for names of contributors.

// Package memdb provide in-memory key/value database implementation.
package memdb

import (
	"leveldb/comparer"
	"leveldb/errors"
	"math/rand"
	"sync/atomic"
	"unsafe"
)

const tMaxHeight = 12

var (
	mPtrSize  int
	mNodeSize int
)

func init() {
	node := new(mNode)
	mPtrSize = int(unsafe.Sizeof(node))
	mNodeSize = int(unsafe.Sizeof(*node))
}

type mNode struct {
	key   []byte
	value []byte
	next  []unsafe.Pointer
}

func (p *mNode) getNext(n int) *mNode {
	return (*mNode)(atomic.LoadPointer(&p.next[n]))
}

func (p *mNode) setNext(n int, x *mNode) {
	atomic.StorePointer(&p.next[n], unsafe.Pointer(x))
}

func (p *mNode) getNext_NB(n int) *mNode {
	return (*mNode)(p.next[n])
}

func (p *mNode) setNext_NB(n int, x *mNode) {
	p.next[n] = unsafe.Pointer(x)
}

// DB represent an in-memory key/value database.
type DB struct {
	cmp       comparer.BasicComparer
	rnd       *rand.Rand
	head      *mNode
	maxHeight int32
	memSize   int64
	n         int32

	prev [tMaxHeight]*mNode
}

// New create new initalized in-memory key/value database.
func New(cmp comparer.BasicComparer) *DB {
	p := &DB{
		cmp:       cmp,
		rnd:       rand.New(rand.NewSource(0xdeadbeef)),
		maxHeight: 1,
	}
	p.head = p.newNode(nil, nil, tMaxHeight)
	return p
}

// Put insert given key and value to the database. Need external synchronization.
func (p *DB) Put(key []byte, value []byte) {
	p.findGE_NB(key, true)

	h := p.randHeight()
	if h > p.maxHeight {
		for i := p.maxHeight; i < h; i++ {
			p.prev[i] = p.head
		}
		atomic.StoreInt32(&p.maxHeight, h)
	}

	x := p.newNode(key, value, h)
	for i, n := range p.prev[:h] {
		x.setNext_NB(i, n.getNext_NB(i))
		n.setNext(i, x)
	}

	atomic.AddInt64(&p.memSize, int64(mNodeSize+(mPtrSize*int(h))+len(key)+len(value)))
	atomic.AddInt32(&p.n, 1)
}

// Remove remove given key from the database. Need external synchronization.
func (p *DB) Remove(key []byte) {
	x := p.findGE_NB(key, true)
	if x == nil || x == p.head || p.cmp.Compare(x.key, key) != 0 {
		return
	}

	h := len(x.next)
	for i, n := range p.prev[:h] {
		n.setNext(i, n.getNext_NB(i).getNext_NB(i))
	}

	atomic.AddInt64(&p.memSize, -int64(mNodeSize+(mPtrSize*h)+len(x.key)+len(x.value)))
	atomic.AddInt32(&p.n, -1)
}

// Contains return true if given key are in database.
func (p *DB) Contains(key []byte) bool {
	x := p.findGE(key, false)
	if x != nil && x != p.head && p.cmp.Compare(x.key, key) == 0 {
		return true
	}
	return false
}

// Get return key/value equal or greater than given key.
func (p *DB) Get(key []byte) (rkey, value []byte, err error) {
	node := p.findGE(key, false)
	if node == nil || node == p.head {
		err = errors.ErrNotFound
		return
	}
	return node.key, node.value, nil
}

// NewIterator create a new iterator over the database content.
func (p *DB) NewIterator() *Iterator {
	return &Iterator{p: p}
}

// Size return approximate size of memory used by the database.
func (p *DB) Size() int {
	return int(atomic.LoadInt64(&p.memSize))
}

// Len return the number of entries in the database.
func (p *DB) Len() int {
	return int(atomic.LoadInt32(&p.n))
}

func (p *DB) newNode(key, value []byte, height int32) *mNode {
	return &mNode{key, value, make([]unsafe.Pointer, height)}
}

func (p *DB) findGE(key []byte, prev bool) *mNode {
	x := p.head
	n := int(atomic.LoadInt32(&p.maxHeight)) - 1
	for {
		next := x.getNext(n)
		if next != nil && p.cmp.Compare(next.key, key) < 0 {
			// Keep searching in this list
			x = next
		} else {
			if prev {
				p.prev[n] = x
			}
			if n == 0 {
				return next
			}
			n--
		}
	}
	return nil
}

func (p *DB) findGE_NB(key []byte, prev bool) *mNode {
	x := p.head
	n := int(p.maxHeight) - 1
	for {
		next := x.getNext_NB(n)
		if next != nil && p.cmp.Compare(next.key, key) < 0 {
			// Keep searching in this list
			x = next
		} else {
			if prev {
				p.prev[n] = x
			}
			if n == 0 {
				return next
			}
			n--
		}
	}
	return nil
}

func (p *DB) findLT(key []byte) *mNode {
	x := p.head
	n := int(atomic.LoadInt32(&p.maxHeight)) - 1
	for {
		next := x.getNext(n)
		if next == nil || p.cmp.Compare(next.key, key) >= 0 {
			if n == 0 {
				return x
			}
			n--
		} else {
			x = next
		}
	}
	return nil
}

func (p *DB) findLast() *mNode {
	x := p.head
	n := int(atomic.LoadInt32(&p.maxHeight)) - 1
	for {
		next := x.getNext(n)
		if next == nil {
			if n == 0 {
				return x
			}
			n--
		} else {
			x = next
		}
	}
	return nil
}

func (p *DB) randHeight() (h int32) {
	const branching = 4
	h = 1
	for h < tMaxHeight && p.rnd.Int()%branching == 0 {
		h++
	}
	return
}

type Iterator struct {
	p      *DB
	node   *mNode
	onLast bool
}

func (i *Iterator) Valid() bool {
	return i.node != nil && i.node != i.p.head
}

func (i *Iterator) First() bool {
	i.node = i.p.head.getNext(0)
	return i.Valid()
}

func (i *Iterator) Last() bool {
	i.node = i.p.findLast()
	return i.Valid()
}

func (i *Iterator) Seek(key []byte) (r bool) {
	i.node = i.p.findGE(key, false)
	return i.Valid()
}

func (i *Iterator) Next() bool {
	if i.node == nil {
		return i.First()
	}
	i.node = i.node.getNext(0)
	res := i.Valid()
	if !res {
		i.onLast = true
	}
	return res
}

func (i *Iterator) Prev() bool {
	if i.node == nil {
		if i.onLast {
			return i.Last()
		}
		return false
	}
	i.node = i.p.findLT(i.node.key)
	if i.node == i.p.head {
		i.node = nil
	}
	return i.Valid()
}

func (i *Iterator) Key() []byte {
	if !i.Valid() {
		return nil
	}
	return i.node.key
}

func (i *Iterator) Value() []byte {
	if !i.Valid() {
		return nil
	}
	return i.node.value
}

func (i *Iterator) Error() error { return nil }
<|endoftext|>


<commit_before>package executor

import (
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"strconv"
	"syscall"

	"github.com/hashicorp/go-multierror"
	"github.com/opencontainers/runc/libcontainer/cgroups"
	cgroupFs "github.com/opencontainers/runc/libcontainer/cgroups/fs"
	"github.com/opencontainers/runc/libcontainer/cgroups/systemd"
	cgroupConfig "github.com/opencontainers/runc/libcontainer/configs"

	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/nomad/structs"
)

var (
	// A mapping of directories on the host OS to attempt to embed inside each
	// task's chroot.
	chrootEnv = map[string]string{
		"/bin":       "/bin",
		"/etc":       "/etc",
		"/lib":       "/lib",
		"/lib32":     "/lib32",
		"/lib64":     "/lib64",
		"/usr/bin":   "/usr/bin",
		"/usr/lib":   "/usr/lib",
		"/usr/share": "/usr/share",
	}
)

// configureIsolation configures chroot and creates cgroups
func (e *UniversalExecutor) configureIsolation() error {
	if e.ctx.FSIsolation {
		if err := e.configureChroot(); err != nil {
			return err
		}
	}

	if e.ctx.ResourceLimits {
		if err := e.configureCgroups(e.ctx.TaskResources); err != nil {
			return fmt.Errorf("error creating cgroups: %v", err)
		}
		if err := e.applyLimits(os.Getpid()); err != nil {
			if er := e.destroyCgroup(); er != nil {
				e.logger.Printf("[ERROR] error destroying cgroup: %v", er)
			}
			if er := e.removeChrootMounts(); er != nil {
				e.logger.Printf("[ERROR] error removing chroot: %v", er)
			}
			return fmt.Errorf("error entering the plugin process in the cgroup: %v:", err)
		}
	}
	return nil
}

// applyLimits puts a process in a pre-configured cgroup
func (e *UniversalExecutor) applyLimits(pid int) error {
	if !e.ctx.ResourceLimits {
		return nil
	}

	// Entering the process in the cgroup
	manager := e.getCgroupManager(e.groups)
	if err := manager.Apply(pid); err != nil {
		e.logger.Printf("[ERROR] unable to join cgroup: %v", err)
		if err := e.Exit(); err != nil {
			e.logger.Printf("[ERROR] unable to kill process: %v", err)
		}
		return err
	}

	return nil
}

// configureCgroups converts a Nomad Resources specification into the equivalent
// cgroup configuration. It returns an error if the resources are invalid.
func (e *UniversalExecutor) configureCgroups(resources *structs.Resources) error {
	e.groups = &cgroupConfig.Cgroup{}
	e.groups.Resources = &cgroupConfig.Resources{}
	e.groups.Name = structs.GenerateUUID()

	// TODO: verify this is needed for things like network access
	e.groups.Resources.AllowAllDevices = true

	if resources.MemoryMB > 0 {
		// Total amount of memory allowed to consume
		e.groups.Resources.Memory = int64(resources.MemoryMB * 1024 * 1024)
		// Disable swap to avoid issues on the machine
		e.groups.Resources.MemorySwap = int64(-1)
	}

	if resources.CPU < 2 {
		return fmt.Errorf("resources.CPU must be equal to or greater than 2: %v", resources.CPU)
	}

	// Set the relative CPU shares for this cgroup.
	e.groups.Resources.CpuShares = int64(resources.CPU)

	if resources.IOPS != 0 {
		// Validate it is in an acceptable range.
		if resources.IOPS < 10 || resources.IOPS > 1000 {
			return fmt.Errorf("resources.IOPS must be between 10 and 1000: %d", resources.IOPS)
		}

		e.groups.Resources.BlkioWeight = uint16(resources.IOPS)
	}

	return nil
}

// runAs takes a user id as a string and looks up the user, and sets the command
// to execute as that user.
func (e *UniversalExecutor) runAs(userid string) error {
	u, err := user.Lookup(userid)
	if err != nil {
		return fmt.Errorf("Failed to identify user %v: %v", userid, err)
	}

	// Convert the uid and gid
	uid, err := strconv.ParseUint(u.Uid, 10, 32)
	if err != nil {
		return fmt.Errorf("Unable to convert userid to uint32: %s", err)
	}
	gid, err := strconv.ParseUint(u.Gid, 10, 32)
	if err != nil {
		return fmt.Errorf("Unable to convert groupid to uint32: %s", err)
	}

	// Set the command to run as that user and group.
	if e.cmd.SysProcAttr == nil {
		e.cmd.SysProcAttr = &syscall.SysProcAttr{}
	}
	if e.cmd.SysProcAttr.Credential == nil {
		e.cmd.SysProcAttr.Credential = &syscall.Credential{}
	}
	e.cmd.SysProcAttr.Credential.Uid = uint32(uid)
	e.cmd.SysProcAttr.Credential.Gid = uint32(gid)

	return nil
}

// configureChroot configures a chroot
func (e *UniversalExecutor) configureChroot() error {
	allocDir := e.ctx.AllocDir
	if err := allocDir.MountSharedDir(e.ctx.TaskName); err != nil {
		return err
	}

	if err := allocDir.Embed(e.ctx.TaskName, chrootEnv); err != nil {
		return err
	}

	// Set the tasks AllocDir environment variable.
	e.ctx.TaskEnv.SetAllocDir(filepath.Join("/", allocdir.SharedAllocName)).SetTaskLocalDir(filepath.Join("/", allocdir.TaskLocal)).Build()

	if e.cmd.SysProcAttr == nil {
		e.cmd.SysProcAttr = &syscall.SysProcAttr{}
	}
	e.cmd.SysProcAttr.Chroot = e.taskDir
	e.cmd.Dir = "/"

	if err := allocDir.MountSpecialDirs(e.taskDir); err != nil {
		return err
	}

	return nil
}

// cleanTaskDir is an idempotent operation to clean the task directory and
// should be called when tearing down the task.
func (e *UniversalExecutor) removeChrootMounts() error {
	// Prevent a race between Wait/ForceStop
	e.lock.Lock()
	defer e.lock.Unlock()

	return e.ctx.AllocDir.UnmountSpecialDirs(e.taskDir)
}

// destroyCgroup kills all processes in the cgroup and removes the cgroup
// configuration from the host.
func (e *UniversalExecutor) destroyCgroup() error {
	if e.groups == nil {
		return fmt.Errorf("Can't destroy: cgroup configuration empty")
	}

	// Prevent a race between Wait/ForceStop
	e.lock.Lock()
	defer e.lock.Unlock()

	manager := e.getCgroupManager(e.groups)
	pids, err := manager.GetPids()
	if err != nil {
		return fmt.Errorf("Failed to get pids in the cgroup %v: %v", e.groups.Name, err)
	}

	errs := new(multierror.Error)
	for _, pid := range pids {
		process, err := os.FindProcess(pid)
		if err != nil {
			multierror.Append(errs, fmt.Errorf("Failed to find Pid %v: %v", pid, err))
			continue
		}

		if err := process.Kill(); err != nil && err.Error() != "os: process already finished" {
			multierror.Append(errs, fmt.Errorf("Failed to kill Pid %v: %v", pid, err))
			continue
		}
	}

	// Remove the cgroup.
	if err := manager.Destroy(); err != nil {
		multierror.Append(errs, fmt.Errorf("Failed to delete the cgroup directories: %v", err))
	}

	if len(errs.Errors) != 0 {
		return fmt.Errorf("Failed to destroy cgroup: %v", errs)
	}

	return nil
}

// getCgroupManager returns the correct libcontainer cgroup manager.
func (e *UniversalExecutor) getCgroupManager(groups *cgroupConfig.Cgroup) cgroups.Manager {
	var manager cgroups.Manager
	manager = &cgroupFs.Manager{Cgroups: groups}
	if systemd.UseSystemd() {
		manager = &systemd.Manager{Cgroups: groups}
	}
	return manager
}
<commit_msg>Appending names of sub-system before log lines<commit_after>package executor

import (
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"strconv"
	"syscall"

	"github.com/hashicorp/go-multierror"
	"github.com/opencontainers/runc/libcontainer/cgroups"
	cgroupFs "github.com/opencontainers/runc/libcontainer/cgroups/fs"
	"github.com/opencontainers/runc/libcontainer/cgroups/systemd"
	cgroupConfig "github.com/opencontainers/runc/libcontainer/configs"

	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/nomad/structs"
)

var (
	// A mapping of directories on the host OS to attempt to embed inside each
	// task's chroot.
	chrootEnv = map[string]string{
		"/bin":       "/bin",
		"/etc":       "/etc",
		"/lib":       "/lib",
		"/lib32":     "/lib32",
		"/lib64":     "/lib64",
		"/usr/bin":   "/usr/bin",
		"/usr/lib":   "/usr/lib",
		"/usr/share": "/usr/share",
	}
)

// configureIsolation configures chroot and creates cgroups
func (e *UniversalExecutor) configureIsolation() error {
	if e.ctx.FSIsolation {
		if err := e.configureChroot(); err != nil {
			return err
		}
	}

	if e.ctx.ResourceLimits {
		if err := e.configureCgroups(e.ctx.TaskResources); err != nil {
			return fmt.Errorf("error creating cgroups: %v", err)
		}
		if err := e.applyLimits(os.Getpid()); err != nil {
			if er := e.destroyCgroup(); er != nil {
				e.logger.Printf("[ERROR] executor: error destroying cgroup: %v", er)
			}
			if er := e.removeChrootMounts(); er != nil {
				e.logger.Printf("[ERROR] executor: error removing chroot: %v", er)
			}
			return fmt.Errorf("error entering the plugin process in the cgroup: %v:", err)
		}
	}
	return nil
}

// applyLimits puts a process in a pre-configured cgroup
func (e *UniversalExecutor) applyLimits(pid int) error {
	if !e.ctx.ResourceLimits {
		return nil
	}

	// Entering the process in the cgroup
	manager := e.getCgroupManager(e.groups)
	if err := manager.Apply(pid); err != nil {
		e.logger.Printf("[ERROR] executor: unable to join cgroup: %v", err)
		if err := e.Exit(); err != nil {
			e.logger.Printf("[ERROR] executor: unable to kill process: %v", err)
		}
		return err
	}

	return nil
}

// configureCgroups converts a Nomad Resources specification into the equivalent
// cgroup configuration. It returns an error if the resources are invalid.
func (e *UniversalExecutor) configureCgroups(resources *structs.Resources) error {
	e.groups = &cgroupConfig.Cgroup{}
	e.groups.Resources = &cgroupConfig.Resources{}
	e.groups.Name = structs.GenerateUUID()

	// TODO: verify this is needed for things like network access
	e.groups.Resources.AllowAllDevices = true

	if resources.MemoryMB > 0 {
		// Total amount of memory allowed to consume
		e.groups.Resources.Memory = int64(resources.MemoryMB * 1024 * 1024)
		// Disable swap to avoid issues on the machine
		e.groups.Resources.MemorySwap = int64(-1)
	}

	if resources.CPU < 2 {
		return fmt.Errorf("resources.CPU must be equal to or greater than 2: %v", resources.CPU)
	}

	// Set the relative CPU shares for this cgroup.
	e.groups.Resources.CpuShares = int64(resources.CPU)

	if resources.IOPS != 0 {
		// Validate it is in an acceptable range.
		if resources.IOPS < 10 || resources.IOPS > 1000 {
			return fmt.Errorf("resources.IOPS must be between 10 and 1000: %d", resources.IOPS)
		}

		e.groups.Resources.BlkioWeight = uint16(resources.IOPS)
	}

	return nil
}

// runAs takes a user id as a string and looks up the user, and sets the command
// to execute as that user.
func (e *UniversalExecutor) runAs(userid string) error {
	u, err := user.Lookup(userid)
	if err != nil {
		return fmt.Errorf("Failed to identify user %v: %v", userid, err)
	}

	// Convert the uid and gid
	uid, err := strconv.ParseUint(u.Uid, 10, 32)
	if err != nil {
		return fmt.Errorf("Unable to convert userid to uint32: %s", err)
	}
	gid, err := strconv.ParseUint(u.Gid, 10, 32)
	if err != nil {
		return fmt.Errorf("Unable to convert groupid to uint32: %s", err)
	}

	// Set the command to run as that user and group.
	if e.cmd.SysProcAttr == nil {
		e.cmd.SysProcAttr = &syscall.SysProcAttr{}
	}
	if e.cmd.SysProcAttr.Credential == nil {
		e.cmd.SysProcAttr.Credential = &syscall.Credential{}
	}
	e.cmd.SysProcAttr.Credential.Uid = uint32(uid)
	e.cmd.SysProcAttr.Credential.Gid = uint32(gid)

	return nil
}

// configureChroot configures a chroot
func (e *UniversalExecutor) configureChroot() error {
	allocDir := e.ctx.AllocDir
	if err := allocDir.MountSharedDir(e.ctx.TaskName); err != nil {
		return err
	}

	if err := allocDir.Embed(e.ctx.TaskName, chrootEnv); err != nil {
		return err
	}

	// Set the tasks AllocDir environment variable.
	e.ctx.TaskEnv.SetAllocDir(filepath.Join("/", allocdir.SharedAllocName)).SetTaskLocalDir(filepath.Join("/", allocdir.TaskLocal)).Build()

	if e.cmd.SysProcAttr == nil {
		e.cmd.SysProcAttr = &syscall.SysProcAttr{}
	}
	e.cmd.SysProcAttr.Chroot = e.taskDir
	e.cmd.Dir = "/"

	if err := allocDir.MountSpecialDirs(e.taskDir); err != nil {
		return err
	}

	return nil
}

// cleanTaskDir is an idempotent operation to clean the task directory and
// should be called when tearing down the task.
func (e *UniversalExecutor) removeChrootMounts() error {
	// Prevent a race between Wait/ForceStop
	e.lock.Lock()
	defer e.lock.Unlock()

	return e.ctx.AllocDir.UnmountSpecialDirs(e.taskDir)
}

// destroyCgroup kills all processes in the cgroup and removes the cgroup
// configuration from the host.
func (e *UniversalExecutor) destroyCgroup() error {
	if e.groups == nil {
		return fmt.Errorf("Can't destroy: cgroup configuration empty")
	}

	// Prevent a race between Wait/ForceStop
	e.lock.Lock()
	defer e.lock.Unlock()

	manager := e.getCgroupManager(e.groups)
	pids, err := manager.GetPids()
	if err != nil {
		return fmt.Errorf("Failed to get pids in the cgroup %v: %v", e.groups.Name, err)
	}

	errs := new(multierror.Error)
	for _, pid := range pids {
		process, err := os.FindProcess(pid)
		if err != nil {
			multierror.Append(errs, fmt.Errorf("Failed to find Pid %v: %v", pid, err))
			continue
		}

		if err := process.Kill(); err != nil && err.Error() != "os: process already finished" {
			multierror.Append(errs, fmt.Errorf("Failed to kill Pid %v: %v", pid, err))
			continue
		}
	}

	// Remove the cgroup.
	if err := manager.Destroy(); err != nil {
		multierror.Append(errs, fmt.Errorf("Failed to delete the cgroup directories: %v", err))
	}

	if len(errs.Errors) != 0 {
		return fmt.Errorf("Failed to destroy cgroup: %v", errs)
	}

	return nil
}

// getCgroupManager returns the correct libcontainer cgroup manager.
func (e *UniversalExecutor) getCgroupManager(groups *cgroupConfig.Cgroup) cgroups.Manager {
	var manager cgroups.Manager
	manager = &cgroupFs.Manager{Cgroups: groups}
	if systemd.UseSystemd() {
		manager = &systemd.Manager{Cgroups: groups}
	}
	return manager
}
<|endoftext|>


<commit_before>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * hprose/http_service.go                                 *
 *                                                        *
 * hprose http service for Go.                            *
 *                                                        *
 * LastModified: Oct 15, 2014                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package hprose

import (
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"reflect"
	"strconv"
	"strings"
	"time"
)

type HttpContext struct {
	*BaseContext
	Response http.ResponseWriter
	Request  *http.Request
}

type HttpServiceEvent interface {
	ServiceEvent
	OnSendHeader(context *HttpContext)
}

type HttpService struct {
	*BaseService
	P3PEnabled                   bool
	GetEnabled                   bool
	CrossDomainEnabled           bool
	accessControlAllowOrigins    map[string]bool
	lastModified                 string
	etag                         string
	crossDomainXmlFile           string
	crossDomainXmlContent        []byte
	clientAccessPolicyXmlFile    string
	clientAccessPolicyXmlContent []byte
}

type httpArgsFixer struct{}

func (httpArgsFixer) FixArgs(args []reflect.Value, lastParamType reflect.Type, context Context) []reflect.Value {
	if c, ok := context.(*HttpContext); ok {
		if lastParamType.String() == "*hprose.HttpContext" {
			return append(args, reflect.ValueOf(c))
		} else if lastParamType.String() == "*http.Request" {
			return append(args, reflect.ValueOf(c.Request))
		}
	}
	return fixArgs(args, lastParamType, context)
}

func NewHttpService() *HttpService {
	t := time.Now().UTC()
	rand.Seed(t.UnixNano())
	service := &HttpService{
		BaseService:               NewBaseService(),
		P3PEnabled:                true,
		GetEnabled:                true,
		CrossDomainEnabled:        true,
		accessControlAllowOrigins: make(map[string]bool),
		lastModified:              t.Format(time.RFC1123),
		etag:                      `"` + strconv.FormatInt(rand.Int63(), 16) + `"`,
	}
	service.argsfixer = httpArgsFixer{}
	return service
}

func (service *HttpService) crossDomainXmlHandler(response http.ResponseWriter, request *http.Request) bool {
	if strings.ToLower(request.URL.Path) == "/crossdomain.xml" {
		if request.Header.Get("if-modified-since") == service.lastModified &&
			request.Header.Get("if-none-match") == service.etag {
			response.WriteHeader(304)
		} else {
			response.Header().Set("Last-Modified", service.lastModified)
			response.Header().Set("Etag", service.etag)
			response.Header().Set("Content-Type", "text/xml")
			response.Header().Set("Content-Length", strconv.Itoa(len(service.crossDomainXmlContent)))
			response.Write(service.crossDomainXmlContent)
		}
		return true
	}
	return false
}

func (service *HttpService) clientAccessPolicyXmlHandler(response http.ResponseWriter, request *http.Request) bool {
	if strings.ToLower(request.URL.Path) == "/clientaccesspolicy.xml" {
		if request.Header.Get("if-modified-since") == service.lastModified &&
			request.Header.Get("if-none-match") == service.etag {
			response.WriteHeader(304)
		} else {
			response.Header().Set("Last-Modified", service.lastModified)
			response.Header().Set("Etag", service.etag)
			response.Header().Set("Content-Type", "text/xml")
			response.Header().Set("Content-Length", strconv.Itoa(len(service.clientAccessPolicyXmlContent)))
			response.Write(service.clientAccessPolicyXmlContent)
		}
		return true
	}
	return false
}

func (service *HttpService) sendHeader(context *HttpContext) {
	if service.ServiceEvent != nil {
		if event, ok := service.ServiceEvent.(HttpServiceEvent); ok {
			event.OnSendHeader(context)
		}
	}
	context.Response.Header().Set("Content-Type", "text/plain")
	if service.P3PEnabled {
		context.Response.Header().Set("P3P", `CP="CAO DSP COR CUR ADM DEV TAI PSA PSD IVAi IVDi `+
			`CONi TELo OTPi OUR DELi SAMi OTRi UNRi PUBi IND PHY ONL `+
			`UNI PUR FIN COM NAV INT DEM CNT STA POL HEA PRE GOV"`)
	}
	if service.CrossDomainEnabled {
		origin := context.Request.Header.Get("origin")
		if origin != "" && origin != "null" {
			if len(service.accessControlAllowOrigins) == 0 || service.accessControlAllowOrigins[origin] {
				context.Response.Header().Set("Access-Control-Allow-Origin", origin)
				context.Response.Header().Set("Access-Control-Allow-Credentials", "true")
			}
		} else {
			context.Response.Header().Set("Access-Control-Allow-Origin", "*")
		}
	}
}

func (service *HttpService) AddAccessControlAllowOrigin(origin string) {
	service.accessControlAllowOrigins[origin] = true
}

func (service *HttpService) RemoveAccessControlAllowOrigin(origin string) {
	delete(service.accessControlAllowOrigins, origin)
}

func (service *HttpService) CrossDomainXmlFile() string {
	return service.crossDomainXmlFile
}

func (service *HttpService) CrossDomainXmlContent() []byte {
	return service.crossDomainXmlContent
}

func (service *HttpService) ClientAccessPolicyXmlFile() string {
	return service.clientAccessPolicyXmlFile
}

func (service *HttpService) ClientAccessPolicyXmlContent() []byte {
	return service.clientAccessPolicyXmlContent
}

func (service *HttpService) SetCrossDomainXmlFile(filename string) {
	service.crossDomainXmlFile = filename
	service.crossDomainXmlContent, _ = ioutil.ReadFile(filename)
}

func (service *HttpService) SetClientAccessPolicyXmlFile(filename string) {
	service.clientAccessPolicyXmlFile = filename
	service.clientAccessPolicyXmlContent, _ = ioutil.ReadFile(filename)
}

func (service *HttpService) SetCrossDomainXmlContent(content []byte) {
	service.crossDomainXmlFile = ""
	service.crossDomainXmlContent = content
}

func (service *HttpService) SetClientAccessPolicyXmlContent(content []byte) {
	service.clientAccessPolicyXmlFile = ""
	service.clientAccessPolicyXmlContent = content
}

func (service *HttpService) readAll(request *http.Request) (data []byte, err error) {
	if request.ContentLength > 0 {
		data = make([]byte, request.ContentLength)
		_, err = io.ReadFull(request.Body, data)
		return data, err
	}
	if request.ContentLength < 0 {
		return ioutil.ReadAll(request.Body)
	}
	return make([]byte, 0), nil
}

func (service *HttpService) ServeHTTP(response http.ResponseWriter, request *http.Request) {
	if service.clientAccessPolicyXmlContent != nil && service.clientAccessPolicyXmlHandler(response, request) {
		return
	}
	if service.crossDomainXmlContent != nil && service.crossDomainXmlHandler(response, request) {
		return
	}
	context := &HttpContext{BaseContext: NewBaseContext(), Response: response, Request: request}
	service.sendHeader(context)
	switch request.Method {
	case "GET":
		if service.GetEnabled {
			response.Write(service.doFunctionList(context))
		} else {
			response.WriteHeader(403)
		}
	case "POST":
		data, err := service.readAll(request)
		request.Body.Close()
		if err != nil {
			response.Write(service.sendError(err, context))
		}
		response.Write(service.Handle(data, context))
	}
}
<commit_msg>Add Serve Method for user custom middleware.<commit_after>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * hprose/http_service.go                                 *
 *                                                        *
 * hprose http service for Go.                            *
 *                                                        *
 * LastModified: Dec 21, 2014                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package hprose

import (
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"reflect"
	"strconv"
	"strings"
	"time"
)

type HttpContext struct {
	*BaseContext
	Response http.ResponseWriter
	Request  *http.Request
}

type HttpServiceEvent interface {
	ServiceEvent
	OnSendHeader(context *HttpContext)
}

type HttpService struct {
	*BaseService
	P3PEnabled                   bool
	GetEnabled                   bool
	CrossDomainEnabled           bool
	accessControlAllowOrigins    map[string]bool
	lastModified                 string
	etag                         string
	crossDomainXmlFile           string
	crossDomainXmlContent        []byte
	clientAccessPolicyXmlFile    string
	clientAccessPolicyXmlContent []byte
}

type httpArgsFixer struct{}

func (httpArgsFixer) FixArgs(args []reflect.Value, lastParamType reflect.Type, context Context) []reflect.Value {
	if c, ok := context.(*HttpContext); ok {
		if lastParamType.String() == "*hprose.HttpContext" {
			return append(args, reflect.ValueOf(c))
		} else if lastParamType.String() == "*http.Request" {
			return append(args, reflect.ValueOf(c.Request))
		}
	}
	return fixArgs(args, lastParamType, context)
}

func NewHttpService() *HttpService {
	t := time.Now().UTC()
	rand.Seed(t.UnixNano())
	service := &HttpService{
		BaseService:               NewBaseService(),
		P3PEnabled:                true,
		GetEnabled:                true,
		CrossDomainEnabled:        true,
		accessControlAllowOrigins: make(map[string]bool),
		lastModified:              t.Format(time.RFC1123),
		etag:                      `"` + strconv.FormatInt(rand.Int63(), 16) + `"`,
	}
	service.argsfixer = httpArgsFixer{}
	return service
}

func (service *HttpService) crossDomainXmlHandler(response http.ResponseWriter, request *http.Request) bool {
	if strings.ToLower(request.URL.Path) == "/crossdomain.xml" {
		if request.Header.Get("if-modified-since") == service.lastModified &&
			request.Header.Get("if-none-match") == service.etag {
			response.WriteHeader(304)
		} else {
			response.Header().Set("Last-Modified", service.lastModified)
			response.Header().Set("Etag", service.etag)
			response.Header().Set("Content-Type", "text/xml")
			response.Header().Set("Content-Length", strconv.Itoa(len(service.crossDomainXmlContent)))
			response.Write(service.crossDomainXmlContent)
		}
		return true
	}
	return false
}

func (service *HttpService) clientAccessPolicyXmlHandler(response http.ResponseWriter, request *http.Request) bool {
	if strings.ToLower(request.URL.Path) == "/clientaccesspolicy.xml" {
		if request.Header.Get("if-modified-since") == service.lastModified &&
			request.Header.Get("if-none-match") == service.etag {
			response.WriteHeader(304)
		} else {
			response.Header().Set("Last-Modified", service.lastModified)
			response.Header().Set("Etag", service.etag)
			response.Header().Set("Content-Type", "text/xml")
			response.Header().Set("Content-Length", strconv.Itoa(len(service.clientAccessPolicyXmlContent)))
			response.Write(service.clientAccessPolicyXmlContent)
		}
		return true
	}
	return false
}

func (service *HttpService) sendHeader(context *HttpContext) {
	if service.ServiceEvent != nil {
		if event, ok := service.ServiceEvent.(HttpServiceEvent); ok {
			event.OnSendHeader(context)
		}
	}
	context.Response.Header().Set("Content-Type", "text/plain")
	if service.P3PEnabled {
		context.Response.Header().Set("P3P", `CP="CAO DSP COR CUR ADM DEV TAI PSA PSD IVAi IVDi `+
			`CONi TELo OTPi OUR DELi SAMi OTRi UNRi PUBi IND PHY ONL `+
			`UNI PUR FIN COM NAV INT DEM CNT STA POL HEA PRE GOV"`)
	}
	if service.CrossDomainEnabled {
		origin := context.Request.Header.Get("origin")
		if origin != "" && origin != "null" {
			if len(service.accessControlAllowOrigins) == 0 || service.accessControlAllowOrigins[origin] {
				context.Response.Header().Set("Access-Control-Allow-Origin", origin)
				context.Response.Header().Set("Access-Control-Allow-Credentials", "true")
			}
		} else {
			context.Response.Header().Set("Access-Control-Allow-Origin", "*")
		}
	}
}

func (service *HttpService) AddAccessControlAllowOrigin(origin string) {
	service.accessControlAllowOrigins[origin] = true
}

func (service *HttpService) RemoveAccessControlAllowOrigin(origin string) {
	delete(service.accessControlAllowOrigins, origin)
}

func (service *HttpService) CrossDomainXmlFile() string {
	return service.crossDomainXmlFile
}

func (service *HttpService) CrossDomainXmlContent() []byte {
	return service.crossDomainXmlContent
}

func (service *HttpService) ClientAccessPolicyXmlFile() string {
	return service.clientAccessPolicyXmlFile
}

func (service *HttpService) ClientAccessPolicyXmlContent() []byte {
	return service.clientAccessPolicyXmlContent
}

func (service *HttpService) SetCrossDomainXmlFile(filename string) {
	service.crossDomainXmlFile = filename
	service.crossDomainXmlContent, _ = ioutil.ReadFile(filename)
}

func (service *HttpService) SetClientAccessPolicyXmlFile(filename string) {
	service.clientAccessPolicyXmlFile = filename
	service.clientAccessPolicyXmlContent, _ = ioutil.ReadFile(filename)
}

func (service *HttpService) SetCrossDomainXmlContent(content []byte) {
	service.crossDomainXmlFile = ""
	service.crossDomainXmlContent = content
}

func (service *HttpService) SetClientAccessPolicyXmlContent(content []byte) {
	service.clientAccessPolicyXmlFile = ""
	service.clientAccessPolicyXmlContent = content
}

func (service *HttpService) readAll(request *http.Request) (data []byte, err error) {
	if request.ContentLength > 0 {
		data = make([]byte, request.ContentLength)
		_, err = io.ReadFull(request.Body, data)
		return data, err
	}
	if request.ContentLength < 0 {
		return ioutil.ReadAll(request.Body)
	}
	return make([]byte, 0), nil
}

func (service *HttpService) Serve(response http.ResponseWriter, request *http.Request, userData map[string]interface{}) {
	if service.clientAccessPolicyXmlContent != nil && service.clientAccessPolicyXmlHandler(response, request) {
		return
	}
	if service.crossDomainXmlContent != nil && service.crossDomainXmlHandler(response, request) {
		return
	}
	context := &HttpContext{BaseContext: NewBaseContext(), Response: response, Request: request}
	if userData != nil {
		for k, v := range userData {
			context.setInterface(k, v);
		}
	}
	service.sendHeader(context)
	switch request.Method {
	case "GET":
		if service.GetEnabled {
			response.Write(service.doFunctionList(context))
		} else {
			response.WriteHeader(403)
		}
	case "POST":
		data, err := service.readAll(request)
		request.Body.Close()
		if err != nil {
			response.Write(service.sendError(err, context))
		}
		response.Write(service.Handle(data, context))
	}
}

func (service *HttpService) ServeHTTP(response http.ResponseWriter, request *http.Request) {
	Serve(response, request, nil)
}
<|endoftext|>


<commit_before>// package bitswap implements the IPFS Exchange interface with the BitSwap
// bilateral exchange protocol.
package bitswap

import (
	"time"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	ds "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-datastore"

	blocks "github.com/jbenet/go-ipfs/blocks"
	blockstore "github.com/jbenet/go-ipfs/blockstore"
	exchange "github.com/jbenet/go-ipfs/exchange"
	bsmsg "github.com/jbenet/go-ipfs/exchange/bitswap/message"
	bsnet "github.com/jbenet/go-ipfs/exchange/bitswap/network"
	notifications "github.com/jbenet/go-ipfs/exchange/bitswap/notifications"
	strategy "github.com/jbenet/go-ipfs/exchange/bitswap/strategy"
	peer "github.com/jbenet/go-ipfs/peer"
	u "github.com/jbenet/go-ipfs/util"
	"github.com/jbenet/go-ipfs/util/eventlog"
)

var log = eventlog.Logger("bitswap")

// New initializes a BitSwap instance that communicates over the
// provided BitSwapNetwork. This function registers the returned instance as
// the network delegate.
// Runs until context is cancelled
func New(ctx context.Context, p peer.Peer,
	network bsnet.BitSwapNetwork, routing bsnet.Routing,
	d ds.ThreadSafeDatastore, nice bool) exchange.Interface {

	notif := notifications.New()
	go func() {
		<-ctx.Done()
		notif.Shutdown()
	}()

	bs := &bitswap{
		blockstore:    blockstore.NewBlockstore(d),
		notifications: notif,
		strategy:      strategy.New(nice),
		routing:       routing,
		sender:        network,
		wantlist:      u.NewKeySet(),
		blockReq:      make(chan u.Key, 32),
	}
	network.SetDelegate(bs)
	go bs.run(ctx)

	return bs
}

// bitswap instances implement the bitswap protocol.
type bitswap struct {

	// sender delivers messages on behalf of the session
	sender bsnet.BitSwapNetwork

	// blockstore is the local database
	// NB: ensure threadsafety
	blockstore blockstore.Blockstore

	// routing interface for communication
	routing bsnet.Routing

	notifications notifications.PubSub

	blockReq chan u.Key

	// strategy listens to network traffic and makes decisions about how to
	// interact with partners.
	// TODO(brian): save the strategy's state to the datastore
	strategy strategy.Strategy

	wantlist u.KeySet
}

// GetBlock attempts to retrieve a particular block from peers within the
// deadline enforced by the context
//
// TODO ensure only one active request per key
func (bs *bitswap) GetBlock(parent context.Context, k u.Key) (*blocks.Block, error) {

	// make sure to derive a new |ctx| and pass it to children. It's correct to
	// listen on |parent| here, but incorrect to pass |parent| to new async
	// functions. This is difficult to enforce. May this comment keep you safe.

	ctx, cancelFunc := context.WithCancel(parent)
	defer cancelFunc()

	ctx = eventlog.ContextWithMetadata(ctx, eventlog.Uuid("BitswapGetBlockRequest"))
	log.Event(ctx, "BitswapGetBlockRequestBegin", &k)

	defer func() {
		log.Event(ctx, "BitSwapGetBlockRequestEnd", &k)
	}()

	bs.wantlist.Add(k)
	promise := bs.notifications.Subscribe(ctx, k)

	select {
	case bs.blockReq <- k:
	case <-parent.Done():
		return nil, parent.Err()
	}

	select {
	case block := <-promise:
		bs.wantlist.Remove(k)
		return &block, nil
	case <-parent.Done():
		return nil, parent.Err()
	}
}

func (bs *bitswap) GetBlocks(parent context.Context, ks []u.Key) (*blocks.Block, error) {
	// TODO: something smart
	return nil, nil
}

func (bs *bitswap) sendWantListTo(ctx context.Context, peers <-chan peer.Peer) error {
	message := bsmsg.New()
	for _, wanted := range bs.wantlist.Keys() {
		message.AddWanted(wanted)
	}
	for peerToQuery := range peers {
		log.Debugf("bitswap got peersToQuery: %s", peerToQuery)
		go func(p peer.Peer) {

			log.Debugf("bitswap dialing peer: %s", p)
			err := bs.sender.DialPeer(ctx, p)
			if err != nil {
				log.Errorf("Error sender.DialPeer(%s)", p)
				return
			}

			response, err := bs.sender.SendRequest(ctx, p, message)
			if err != nil {
				log.Errorf("Error sender.SendRequest(%s) = %s", p, err)
				return
			}
			// FIXME ensure accounting is handled correctly when
			// communication fails. May require slightly different API to
			// get better guarantees. May need shared sequence numbers.
			bs.strategy.MessageSent(p, message)

			if response == nil {
				return
			}
			bs.ReceiveMessage(ctx, p, response)
		}(peerToQuery)
	}
	return nil
}

func (bs *bitswap) run(ctx context.Context) {
	var sendlist <-chan peer.Peer

	// Every so often, we should resend out our current want list
	rebroadcastTime := time.Second * 5

	// Time to wait before sending out wantlists to better batch up requests
	bufferTime := time.Millisecond * 3
	peersPerSend := 6

	timeout := time.After(rebroadcastTime)
	threshold := 10
	unsent := 0
	for {
		select {
		case <-timeout:
			wantlist := bs.wantlist.Keys()
			if len(wantlist) == 0 {
				continue
			}
			if sendlist == nil {
				// rely on semi randomness of maps
				firstKey := wantlist[0]
				sendlist = bs.routing.FindProvidersAsync(ctx, firstKey, 6)
			}
			err := bs.sendWantListTo(ctx, sendlist)
			if err != nil {
				log.Error("error sending wantlist: %s", err)
			}
			sendlist = nil
			timeout = time.After(rebroadcastTime)
		case k := <-bs.blockReq:
			if unsent == 0 {
				sendlist = bs.routing.FindProvidersAsync(ctx, k, peersPerSend)
			}
			unsent++

			if unsent >= threshold {
				// send wantlist to sendlist
				bs.sendWantListTo(ctx, sendlist)
				unsent = 0
				timeout = time.After(rebroadcastTime)
				sendlist = nil
			} else {
				// set a timeout to wait for more blocks or send current wantlist

				timeout = time.After(bufferTime)
			}
		case <-ctx.Done():
			return
		}
	}
}

// HasBlock announces the existance of a block to this bitswap service. The
// service will potentially notify its peers.
func (bs *bitswap) HasBlock(ctx context.Context, blk blocks.Block) error {
	log.Debugf("Has Block %v", blk.Key())
	bs.wantlist.Remove(blk.Key())
	bs.sendToPeersThatWant(ctx, blk)
	return bs.routing.Provide(ctx, blk.Key())
}

// TODO(brian): handle errors
func (bs *bitswap) ReceiveMessage(ctx context.Context, p peer.Peer, incoming bsmsg.BitSwapMessage) (
	peer.Peer, bsmsg.BitSwapMessage) {
	log.Debugf("ReceiveMessage from %v", p.Key())
	log.Debugf("Message wantlist: %v", incoming.Wantlist())

	if p == nil {
		log.Error("Received message from nil peer!")
		// TODO propagate the error upward
		return nil, nil
	}
	if incoming == nil {
		log.Error("Got nil bitswap message!")
		// TODO propagate the error upward
		return nil, nil
	}

	// Record message bytes in ledger
	// TODO: this is bad, and could be easily abused.
	// Should only track *useful* messages in ledger
	bs.strategy.MessageReceived(p, incoming) // FIRST

	for _, block := range incoming.Blocks() {
		// TODO verify blocks?
		if err := bs.blockstore.Put(&block); err != nil {
			continue // FIXME(brian): err ignored
		}
		bs.notifications.Publish(block)
		err := bs.HasBlock(ctx, block)
		if err != nil {
			log.Warningf("HasBlock errored: %s", err)
		}
	}

	message := bsmsg.New()
	for _, wanted := range bs.wantlist.Keys() {
		message.AddWanted(wanted)
	}
	for _, key := range incoming.Wantlist() {
		// TODO: might be better to check if we have the block before checking
		//			if we should send it to someone
		if bs.strategy.ShouldSendBlockToPeer(key, p) {
			if block, errBlockNotFound := bs.blockstore.Get(key); errBlockNotFound != nil {
				continue
			} else {
				message.AddBlock(*block)
			}
		}
	}

	bs.strategy.MessageSent(p, message)
	log.Debug("Returning message.")
	return p, message
}

func (bs *bitswap) ReceiveError(err error) {
	log.Errorf("Bitswap ReceiveError: %s", err)
	// TODO log the network error
	// TODO bubble the network error up to the parent context/error logger
}

// send strives to ensure that accounting is always performed when a message is
// sent
func (bs *bitswap) send(ctx context.Context, p peer.Peer, m bsmsg.BitSwapMessage) {
	bs.sender.SendMessage(ctx, p, m)
	bs.strategy.MessageSent(p, m)
}

func (bs *bitswap) sendToPeersThatWant(ctx context.Context, block blocks.Block) {
	log.Debugf("Sending %v to peers that want it", block.Key())

	for _, p := range bs.strategy.Peers() {
		if bs.strategy.BlockIsWantedByPeer(block.Key(), p) {
			log.Debugf("%v wants %v", p, block.Key())
			if bs.strategy.ShouldSendBlockToPeer(block.Key(), p) {
				message := bsmsg.New()
				message.AddBlock(block)
				for _, wanted := range bs.wantlist.Keys() {
					message.AddWanted(wanted)
				}
				bs.send(ctx, p, message)
			}
		}
	}
}
<commit_msg>fix(bitswap) handle error<commit_after>// package bitswap implements the IPFS Exchange interface with the BitSwap
// bilateral exchange protocol.
package bitswap

import (
	"time"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	ds "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-datastore"

	blocks "github.com/jbenet/go-ipfs/blocks"
	blockstore "github.com/jbenet/go-ipfs/blockstore"
	exchange "github.com/jbenet/go-ipfs/exchange"
	bsmsg "github.com/jbenet/go-ipfs/exchange/bitswap/message"
	bsnet "github.com/jbenet/go-ipfs/exchange/bitswap/network"
	notifications "github.com/jbenet/go-ipfs/exchange/bitswap/notifications"
	strategy "github.com/jbenet/go-ipfs/exchange/bitswap/strategy"
	peer "github.com/jbenet/go-ipfs/peer"
	u "github.com/jbenet/go-ipfs/util"
	"github.com/jbenet/go-ipfs/util/eventlog"
)

var log = eventlog.Logger("bitswap")

// New initializes a BitSwap instance that communicates over the
// provided BitSwapNetwork. This function registers the returned instance as
// the network delegate.
// Runs until context is cancelled
func New(ctx context.Context, p peer.Peer,
	network bsnet.BitSwapNetwork, routing bsnet.Routing,
	d ds.ThreadSafeDatastore, nice bool) exchange.Interface {

	notif := notifications.New()
	go func() {
		<-ctx.Done()
		notif.Shutdown()
	}()

	bs := &bitswap{
		blockstore:    blockstore.NewBlockstore(d),
		notifications: notif,
		strategy:      strategy.New(nice),
		routing:       routing,
		sender:        network,
		wantlist:      u.NewKeySet(),
		blockReq:      make(chan u.Key, 32),
	}
	network.SetDelegate(bs)
	go bs.run(ctx)

	return bs
}

// bitswap instances implement the bitswap protocol.
type bitswap struct {

	// sender delivers messages on behalf of the session
	sender bsnet.BitSwapNetwork

	// blockstore is the local database
	// NB: ensure threadsafety
	blockstore blockstore.Blockstore

	// routing interface for communication
	routing bsnet.Routing

	notifications notifications.PubSub

	blockReq chan u.Key

	// strategy listens to network traffic and makes decisions about how to
	// interact with partners.
	// TODO(brian): save the strategy's state to the datastore
	strategy strategy.Strategy

	wantlist u.KeySet
}

// GetBlock attempts to retrieve a particular block from peers within the
// deadline enforced by the context
//
// TODO ensure only one active request per key
func (bs *bitswap) GetBlock(parent context.Context, k u.Key) (*blocks.Block, error) {

	// make sure to derive a new |ctx| and pass it to children. It's correct to
	// listen on |parent| here, but incorrect to pass |parent| to new async
	// functions. This is difficult to enforce. May this comment keep you safe.

	ctx, cancelFunc := context.WithCancel(parent)
	defer cancelFunc()

	ctx = eventlog.ContextWithMetadata(ctx, eventlog.Uuid("BitswapGetBlockRequest"))
	log.Event(ctx, "BitswapGetBlockRequestBegin", &k)

	defer func() {
		log.Event(ctx, "BitSwapGetBlockRequestEnd", &k)
	}()

	bs.wantlist.Add(k)
	promise := bs.notifications.Subscribe(ctx, k)

	select {
	case bs.blockReq <- k:
	case <-parent.Done():
		return nil, parent.Err()
	}

	select {
	case block := <-promise:
		bs.wantlist.Remove(k)
		return &block, nil
	case <-parent.Done():
		return nil, parent.Err()
	}
}

func (bs *bitswap) GetBlocks(parent context.Context, ks []u.Key) (*blocks.Block, error) {
	// TODO: something smart
	return nil, nil
}

func (bs *bitswap) sendWantListTo(ctx context.Context, peers <-chan peer.Peer) error {
	message := bsmsg.New()
	for _, wanted := range bs.wantlist.Keys() {
		message.AddWanted(wanted)
	}
	for peerToQuery := range peers {
		log.Debugf("bitswap got peersToQuery: %s", peerToQuery)
		go func(p peer.Peer) {

			log.Debugf("bitswap dialing peer: %s", p)
			err := bs.sender.DialPeer(ctx, p)
			if err != nil {
				log.Errorf("Error sender.DialPeer(%s)", p)
				return
			}

			response, err := bs.sender.SendRequest(ctx, p, message)
			if err != nil {
				log.Errorf("Error sender.SendRequest(%s) = %s", p, err)
				return
			}
			// FIXME ensure accounting is handled correctly when
			// communication fails. May require slightly different API to
			// get better guarantees. May need shared sequence numbers.
			bs.strategy.MessageSent(p, message)

			if response == nil {
				return
			}
			bs.ReceiveMessage(ctx, p, response)
		}(peerToQuery)
	}
	return nil
}

func (bs *bitswap) run(ctx context.Context) {
	var sendlist <-chan peer.Peer

	// Every so often, we should resend out our current want list
	rebroadcastTime := time.Second * 5

	// Time to wait before sending out wantlists to better batch up requests
	bufferTime := time.Millisecond * 3
	peersPerSend := 6

	timeout := time.After(rebroadcastTime)
	threshold := 10
	unsent := 0
	for {
		select {
		case <-timeout:
			wantlist := bs.wantlist.Keys()
			if len(wantlist) == 0 {
				continue
			}
			if sendlist == nil {
				// rely on semi randomness of maps
				firstKey := wantlist[0]
				sendlist = bs.routing.FindProvidersAsync(ctx, firstKey, 6)
			}
			err := bs.sendWantListTo(ctx, sendlist)
			if err != nil {
				log.Error("error sending wantlist: %s", err)
			}
			sendlist = nil
			timeout = time.After(rebroadcastTime)
		case k := <-bs.blockReq:
			if unsent == 0 {
				sendlist = bs.routing.FindProvidersAsync(ctx, k, peersPerSend)
			}
			unsent++

			if unsent >= threshold {
				// send wantlist to sendlist
				err := bs.sendWantListTo(ctx, sendlist)
				if err != nil {
					log.Errorf("error sending wantlist: %s", err)
				}
				unsent = 0
				timeout = time.After(rebroadcastTime)
				sendlist = nil
			} else {
				// set a timeout to wait for more blocks or send current wantlist

				timeout = time.After(bufferTime)
			}
		case <-ctx.Done():
			return
		}
	}
}

// HasBlock announces the existance of a block to this bitswap service. The
// service will potentially notify its peers.
func (bs *bitswap) HasBlock(ctx context.Context, blk blocks.Block) error {
	log.Debugf("Has Block %v", blk.Key())
	bs.wantlist.Remove(blk.Key())
	bs.sendToPeersThatWant(ctx, blk)
	return bs.routing.Provide(ctx, blk.Key())
}

// TODO(brian): handle errors
func (bs *bitswap) ReceiveMessage(ctx context.Context, p peer.Peer, incoming bsmsg.BitSwapMessage) (
	peer.Peer, bsmsg.BitSwapMessage) {
	log.Debugf("ReceiveMessage from %v", p.Key())
	log.Debugf("Message wantlist: %v", incoming.Wantlist())

	if p == nil {
		log.Error("Received message from nil peer!")
		// TODO propagate the error upward
		return nil, nil
	}
	if incoming == nil {
		log.Error("Got nil bitswap message!")
		// TODO propagate the error upward
		return nil, nil
	}

	// Record message bytes in ledger
	// TODO: this is bad, and could be easily abused.
	// Should only track *useful* messages in ledger
	bs.strategy.MessageReceived(p, incoming) // FIRST

	for _, block := range incoming.Blocks() {
		// TODO verify blocks?
		if err := bs.blockstore.Put(&block); err != nil {
			continue // FIXME(brian): err ignored
		}
		bs.notifications.Publish(block)
		err := bs.HasBlock(ctx, block)
		if err != nil {
			log.Warningf("HasBlock errored: %s", err)
		}
	}

	message := bsmsg.New()
	for _, wanted := range bs.wantlist.Keys() {
		message.AddWanted(wanted)
	}
	for _, key := range incoming.Wantlist() {
		// TODO: might be better to check if we have the block before checking
		//			if we should send it to someone
		if bs.strategy.ShouldSendBlockToPeer(key, p) {
			if block, errBlockNotFound := bs.blockstore.Get(key); errBlockNotFound != nil {
				continue
			} else {
				message.AddBlock(*block)
			}
		}
	}

	bs.strategy.MessageSent(p, message)
	log.Debug("Returning message.")
	return p, message
}

func (bs *bitswap) ReceiveError(err error) {
	log.Errorf("Bitswap ReceiveError: %s", err)
	// TODO log the network error
	// TODO bubble the network error up to the parent context/error logger
}

// send strives to ensure that accounting is always performed when a message is
// sent
func (bs *bitswap) send(ctx context.Context, p peer.Peer, m bsmsg.BitSwapMessage) {
	bs.sender.SendMessage(ctx, p, m)
	bs.strategy.MessageSent(p, m)
}

func (bs *bitswap) sendToPeersThatWant(ctx context.Context, block blocks.Block) {
	log.Debugf("Sending %v to peers that want it", block.Key())

	for _, p := range bs.strategy.Peers() {
		if bs.strategy.BlockIsWantedByPeer(block.Key(), p) {
			log.Debugf("%v wants %v", p, block.Key())
			if bs.strategy.ShouldSendBlockToPeer(block.Key(), p) {
				message := bsmsg.New()
				message.AddBlock(block)
				for _, wanted := range bs.wantlist.Keys() {
					message.AddWanted(wanted)
				}
				bs.send(ctx, p, message)
			}
		}
	}
}
<|endoftext|>


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

package integration

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"path/filepath"
	"testing"
	"time"

	"go.uber.org/zap"
	"google.golang.org/grpc"

	"go.etcd.io/etcd/clientv3"
	"go.etcd.io/etcd/etcdserver/api/v3rpc/rpctypes"
	"go.etcd.io/etcd/integration"
	"go.etcd.io/etcd/lease"
	"go.etcd.io/etcd/mvcc"
	"go.etcd.io/etcd/mvcc/backend"
	"go.etcd.io/etcd/pkg/testutil"
)

func TestMaintenanceHashKV(t *testing.T) {
	defer testutil.AfterTest(t)

	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 3})
	defer clus.Terminate(t)

	for i := 0; i < 3; i++ {
		if _, err := clus.RandClient().Put(context.Background(), "foo", "bar"); err != nil {
			t.Fatal(err)
		}
	}

	var hv uint32
	for i := 0; i < 3; i++ {
		cli := clus.Client(i)
		// ensure writes are replicated
		if _, err := cli.Get(context.TODO(), "foo"); err != nil {
			t.Fatal(err)
		}
		hresp, err := cli.HashKV(context.Background(), clus.Members[i].GRPCAddr(), 0)
		if err != nil {
			t.Fatal(err)
		}
		if hv == 0 {
			hv = hresp.Hash
			continue
		}
		if hv != hresp.Hash {
			t.Fatalf("#%d: hash expected %d, got %d", i, hv, hresp.Hash)
		}
	}
}

func TestMaintenanceMoveLeader(t *testing.T) {
	defer testutil.AfterTest(t)

	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 3})
	defer clus.Terminate(t)

	oldLeadIdx := clus.WaitLeader(t)
	targetIdx := (oldLeadIdx + 1) % 3
	target := uint64(clus.Members[targetIdx].ID())

	cli := clus.Client(targetIdx)
	_, err := cli.MoveLeader(context.Background(), target)
	if err != rpctypes.ErrNotLeader {
		t.Fatalf("error expected %v, got %v", rpctypes.ErrNotLeader, err)
	}

	cli = clus.Client(oldLeadIdx)
	_, err = cli.MoveLeader(context.Background(), target)
	if err != nil {
		t.Fatal(err)
	}

	leadIdx := clus.WaitLeader(t)
	lead := uint64(clus.Members[leadIdx].ID())
	if target != lead {
		t.Fatalf("new leader expected %d, got %d", target, lead)
	}
}

// TestMaintenanceSnapshotError ensures that context cancel/timeout
// before snapshot reading returns corresponding context errors.
func TestMaintenanceSnapshotError(t *testing.T) {
	defer testutil.AfterTest(t)

	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	// reading snapshot with canceled context should error out
	ctx, cancel := context.WithCancel(context.Background())
	rc1, err := clus.RandClient().Snapshot(ctx)
	if err != nil {
		t.Fatal(err)
	}
	defer rc1.Close()

	cancel()
	_, err = io.Copy(ioutil.Discard, rc1)
	if err != context.Canceled {
		t.Errorf("expected %v, got %v", context.Canceled, err)
	}

	// reading snapshot with deadline exceeded should error out
	ctx, cancel = context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	rc2, err := clus.RandClient().Snapshot(ctx)
	if err != nil {
		t.Fatal(err)
	}
	defer rc2.Close()

	time.Sleep(2 * time.Second)

	_, err = io.Copy(ioutil.Discard, rc2)
	if err != nil && !isClientTimeout(err) {
		t.Errorf("expected client timeout, got %v", err)
	}
}

// TestMaintenanceSnapshotErrorInflight ensures that inflight context cancel/timeout
// fails snapshot reading with corresponding context errors.
func TestMaintenanceSnapshotErrorInflight(t *testing.T) {
	defer testutil.AfterTest(t)

	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	// take about 1-second to read snapshot
	clus.Members[0].Stop(t)
	dpath := filepath.Join(clus.Members[0].DataDir, "member", "snap", "db")
	b := backend.NewDefaultBackend(dpath)
	s := mvcc.NewStore(zap.NewExample(), b, &lease.FakeLessor{}, nil)
	rev := 100000
	for i := 2; i <= rev; i++ {
		s.Put([]byte(fmt.Sprintf("%10d", i)), bytes.Repeat([]byte("a"), 1024), lease.NoLease)
	}
	s.Close()
	b.Close()
	clus.Members[0].Restart(t)

	cli := clus.RandClient()
	// reading snapshot with canceled context should error out
	ctx, cancel := context.WithCancel(context.Background())
	rc1, err := cli.Snapshot(ctx)
	if err != nil {
		t.Fatal(err)
	}
	defer rc1.Close()

	donec := make(chan struct{})
	go func() {
		time.Sleep(300 * time.Millisecond)
		cancel()
		close(donec)
	}()
	_, err = io.Copy(ioutil.Discard, rc1)
	if err != nil && err != context.Canceled {
		t.Errorf("expected %v, got %v", context.Canceled, err)
	}
	<-donec

	// reading snapshot with deadline exceeded should error out
	ctx, cancel = context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	rc2, err := clus.RandClient().Snapshot(ctx)
	if err != nil {
		t.Fatal(err)
	}
	defer rc2.Close()

	// 300ms left and expect timeout while snapshot reading is in progress
	time.Sleep(700 * time.Millisecond)
	_, err = io.Copy(ioutil.Discard, rc2)
	if err != nil && !isClientTimeout(err) {
		t.Errorf("expected client timeout, got %v", err)
	}
}

func TestMaintenanceStatus(t *testing.T) {
	defer testutil.AfterTest(t)

	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 3})
	defer clus.Terminate(t)

	clus.WaitLeader(t)

	eps := make([]string, 3)
	for i := 0; i < 3; i++ {
		eps[i] = clus.Members[i].GRPCAddr()
	}

	cli, err := clientv3.New(clientv3.Config{Endpoints: eps, DialOptions: []grpc.DialOption{grpc.WithBlock()}})
	if err != nil {
		t.Fatal(err)
	}
	defer cli.Close()

	prevID, leaderFound := uint64(0), false
	for i := 0; i < 3; i++ {
		resp, err := cli.Status(context.TODO(), eps[i])
		if err != nil {
			t.Fatal(err)
		}
		if prevID == 0 {
			prevID, leaderFound = resp.Header.MemberId, resp.Header.MemberId == resp.Leader
			continue
		}
		if prevID == resp.Header.MemberId {
			t.Errorf("#%d: status returned duplicate member ID with %016x", i, prevID)
		}
		if leaderFound && resp.Header.MemberId == resp.Leader {
			t.Errorf("#%d: leader already found, but found another %016x", i, resp.Header.MemberId)
		}
		if !leaderFound {
			leaderFound = resp.Header.MemberId == resp.Leader
		}
	}
	if !leaderFound {
		t.Fatal("no leader found")
	}
}
<commit_msg>clientv3/integration: fix "mvcc.NewStore" call<commit_after>// Copyright 2017 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package integration

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"path/filepath"
	"testing"
	"time"

	"go.uber.org/zap"
	"google.golang.org/grpc"

	"go.etcd.io/etcd/clientv3"
	"go.etcd.io/etcd/etcdserver/api/v3rpc/rpctypes"
	"go.etcd.io/etcd/integration"
	"go.etcd.io/etcd/lease"
	"go.etcd.io/etcd/mvcc"
	"go.etcd.io/etcd/mvcc/backend"
	"go.etcd.io/etcd/pkg/testutil"
)

func TestMaintenanceHashKV(t *testing.T) {
	defer testutil.AfterTest(t)

	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 3})
	defer clus.Terminate(t)

	for i := 0; i < 3; i++ {
		if _, err := clus.RandClient().Put(context.Background(), "foo", "bar"); err != nil {
			t.Fatal(err)
		}
	}

	var hv uint32
	for i := 0; i < 3; i++ {
		cli := clus.Client(i)
		// ensure writes are replicated
		if _, err := cli.Get(context.TODO(), "foo"); err != nil {
			t.Fatal(err)
		}
		hresp, err := cli.HashKV(context.Background(), clus.Members[i].GRPCAddr(), 0)
		if err != nil {
			t.Fatal(err)
		}
		if hv == 0 {
			hv = hresp.Hash
			continue
		}
		if hv != hresp.Hash {
			t.Fatalf("#%d: hash expected %d, got %d", i, hv, hresp.Hash)
		}
	}
}

func TestMaintenanceMoveLeader(t *testing.T) {
	defer testutil.AfterTest(t)

	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 3})
	defer clus.Terminate(t)

	oldLeadIdx := clus.WaitLeader(t)
	targetIdx := (oldLeadIdx + 1) % 3
	target := uint64(clus.Members[targetIdx].ID())

	cli := clus.Client(targetIdx)
	_, err := cli.MoveLeader(context.Background(), target)
	if err != rpctypes.ErrNotLeader {
		t.Fatalf("error expected %v, got %v", rpctypes.ErrNotLeader, err)
	}

	cli = clus.Client(oldLeadIdx)
	_, err = cli.MoveLeader(context.Background(), target)
	if err != nil {
		t.Fatal(err)
	}

	leadIdx := clus.WaitLeader(t)
	lead := uint64(clus.Members[leadIdx].ID())
	if target != lead {
		t.Fatalf("new leader expected %d, got %d", target, lead)
	}
}

// TestMaintenanceSnapshotError ensures that context cancel/timeout
// before snapshot reading returns corresponding context errors.
func TestMaintenanceSnapshotError(t *testing.T) {
	defer testutil.AfterTest(t)

	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	// reading snapshot with canceled context should error out
	ctx, cancel := context.WithCancel(context.Background())
	rc1, err := clus.RandClient().Snapshot(ctx)
	if err != nil {
		t.Fatal(err)
	}
	defer rc1.Close()

	cancel()
	_, err = io.Copy(ioutil.Discard, rc1)
	if err != context.Canceled {
		t.Errorf("expected %v, got %v", context.Canceled, err)
	}

	// reading snapshot with deadline exceeded should error out
	ctx, cancel = context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	rc2, err := clus.RandClient().Snapshot(ctx)
	if err != nil {
		t.Fatal(err)
	}
	defer rc2.Close()

	time.Sleep(2 * time.Second)

	_, err = io.Copy(ioutil.Discard, rc2)
	if err != nil && !isClientTimeout(err) {
		t.Errorf("expected client timeout, got %v", err)
	}
}

// TestMaintenanceSnapshotErrorInflight ensures that inflight context cancel/timeout
// fails snapshot reading with corresponding context errors.
func TestMaintenanceSnapshotErrorInflight(t *testing.T) {
	defer testutil.AfterTest(t)

	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	// take about 1-second to read snapshot
	clus.Members[0].Stop(t)
	dpath := filepath.Join(clus.Members[0].DataDir, "member", "snap", "db")
	b := backend.NewDefaultBackend(dpath)
	s := mvcc.NewStore(zap.NewExample(), b, &lease.FakeLessor{}, nil, mvcc.StoreConfig{CompactionBatchLimit: math.MaxInt32})
	rev := 100000
	for i := 2; i <= rev; i++ {
		s.Put([]byte(fmt.Sprintf("%10d", i)), bytes.Repeat([]byte("a"), 1024), lease.NoLease)
	}
	s.Close()
	b.Close()
	clus.Members[0].Restart(t)

	cli := clus.RandClient()
	// reading snapshot with canceled context should error out
	ctx, cancel := context.WithCancel(context.Background())
	rc1, err := cli.Snapshot(ctx)
	if err != nil {
		t.Fatal(err)
	}
	defer rc1.Close()

	donec := make(chan struct{})
	go func() {
		time.Sleep(300 * time.Millisecond)
		cancel()
		close(donec)
	}()
	_, err = io.Copy(ioutil.Discard, rc1)
	if err != nil && err != context.Canceled {
		t.Errorf("expected %v, got %v", context.Canceled, err)
	}
	<-donec

	// reading snapshot with deadline exceeded should error out
	ctx, cancel = context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	rc2, err := clus.RandClient().Snapshot(ctx)
	if err != nil {
		t.Fatal(err)
	}
	defer rc2.Close()

	// 300ms left and expect timeout while snapshot reading is in progress
	time.Sleep(700 * time.Millisecond)
	_, err = io.Copy(ioutil.Discard, rc2)
	if err != nil && !isClientTimeout(err) {
		t.Errorf("expected client timeout, got %v", err)
	}
}

func TestMaintenanceStatus(t *testing.T) {
	defer testutil.AfterTest(t)

	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 3})
	defer clus.Terminate(t)

	clus.WaitLeader(t)

	eps := make([]string, 3)
	for i := 0; i < 3; i++ {
		eps[i] = clus.Members[i].GRPCAddr()
	}

	cli, err := clientv3.New(clientv3.Config{Endpoints: eps, DialOptions: []grpc.DialOption{grpc.WithBlock()}})
	if err != nil {
		t.Fatal(err)
	}
	defer cli.Close()

	prevID, leaderFound := uint64(0), false
	for i := 0; i < 3; i++ {
		resp, err := cli.Status(context.TODO(), eps[i])
		if err != nil {
			t.Fatal(err)
		}
		if prevID == 0 {
			prevID, leaderFound = resp.Header.MemberId, resp.Header.MemberId == resp.Leader
			continue
		}
		if prevID == resp.Header.MemberId {
			t.Errorf("#%d: status returned duplicate member ID with %016x", i, prevID)
		}
		if leaderFound && resp.Header.MemberId == resp.Leader {
			t.Errorf("#%d: leader already found, but found another %016x", i, resp.Header.MemberId)
		}
		if !leaderFound {
			leaderFound = resp.Header.MemberId == resp.Leader
		}
	}
	if !leaderFound {
		t.Fatal("no leader found")
	}
}
<|endoftext|>