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

package certs

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/x509"
	"encoding/asn1"
	"encoding/pem"
	"net"
	"reflect"
	"testing"

	"github.com/01org/ciao/ssntp"
)

func TestGenerateKey(t *testing.T) {
	key, err := generatePrivateKey(false)

	_, ok := key.(*rsa.PrivateKey)
	if !ok || err != nil {
		t.Errorf("RSA key expected from generatePrivateKey: %v", err)
	}

	pKey := publicKey(key)
	_, ok = pKey.(*rsa.PublicKey)
	if !ok {
		t.Error("RSA public key expected")
	}

	key, err = generatePrivateKey(true)
	_, ok = key.(*ecdsa.PrivateKey)

	if !ok || err != nil {
		t.Errorf("ECDSA key expected from generatePrivateKey: %v", err)
	}

	pKey = publicKey(key)
	_, ok = pKey.(*ecdsa.PublicKey)
	if !ok {
		t.Error("ECDSA public key expected")
	}
}

func TestPemBlockParsing(t *testing.T) {
	origKey, _ := generatePrivateKey(false)

	block, err := pemBlockForKey(origKey)

	if block.Type != "RSA PRIVATE KEY" || err != nil {
		t.Errorf("Unexpected PEM block type: %v: err: %v", block.Type, err)
	}

	parsedKey, err := keyFromPemBlock(block)
	_, ok := parsedKey.(*rsa.PrivateKey)

	if !ok || err != nil {
		t.Errorf("Expected RSA private key: %v", err)
	}

	origKey, _ = generatePrivateKey(true)

	block, err = pemBlockForKey(origKey)

	if block.Type != "EC PRIVATE KEY" || err != nil {
		t.Errorf("Unexpected PEM block type: %v: err: %v", block.Type, err)
	}

	parsedKey, err = keyFromPemBlock(block)
	_, ok = parsedKey.(*ecdsa.PrivateKey)

	if !ok || err != nil {
		t.Errorf("Expected ECDSA private key: %v", err)
	}
}

func TestAddOIDs(t *testing.T) {
	mapping := map[ssntp.Role]asn1.ObjectIdentifier{
		ssntp.AGENT:      ssntp.RoleAgentOID,
		ssntp.SCHEDULER:  ssntp.RoleSchedulerOID,
		ssntp.Controller: ssntp.RoleControllerOID,
		ssntp.NETAGENT:   ssntp.RoleNetAgentOID,
		ssntp.SERVER:     ssntp.RoleServerOID,
		ssntp.CNCIAGENT:  ssntp.RoleCNCIAgentOID,
	}

	// Check all
	for role, oid := range mapping {
		oids := addOIDs(role, []asn1.ObjectIdentifier{})
		if len(oids) != 1 {
			t.Errorf("Expected only one OID found %d", len(oids))
		}

		if !reflect.DeepEqual(oids[0], oid) {
			t.Errorf("Unexpected role to OID mapping: %v -> %v", role, oid)
		}
	}

	// Check common pairing
	oids := addOIDs(ssntp.AGENT|ssntp.NETAGENT, []asn1.ObjectIdentifier{})
	if len(oids) != 2 {
		t.Errorf("Expected two OIDS found %d", len(oids))
	}

	if !reflect.DeepEqual(oids[0], ssntp.RoleAgentOID) || !reflect.DeepEqual(oids[1], ssntp.RoleNetAgentOID) {
		t.Errorf("Unexpected OIDS in list: %v and %v", oids[0], oids[1])
	}
}

func TestCreateCertTemplateHosts(t *testing.T) {
	hosts := []string{"test.example.com", "test2.example.com"}
	mgmtIPs := []string{}
	cert, err := CreateCertTemplate(ssntp.AGENT, "ACME Corp", "test@example.com", hosts, mgmtIPs)

	if err != nil {
		t.Errorf("Unexpected error when creating cert template: %v", err)
	}

	if !reflect.DeepEqual(cert.DNSNames, hosts) {
		t.Errorf("Hosts in certificate don't match: %v %v", hosts, cert.DNSNames)
	}
}

func TestCreateCertTemplateIPs(t *testing.T) {
	hosts := []string{}
	mgmtIPs := []string{"127.0.0.1", "127.0.0.2"}
	cert, err := CreateCertTemplate(ssntp.AGENT, "ACME Corp", "test@example.com", hosts, mgmtIPs)

	if err != nil {
		t.Errorf("Unexpected error when creating cert template: %v", err)
	}

	IPs := []net.IP{net.IPv4(127, 0, 0, 1), net.IPv4(127, 0, 0, 2)}
	if !reflect.DeepEqual(cert.IPAddresses, IPs) {
		t.Errorf("IPs in certificate don't match: %v %v", hosts, cert.DNSNames)
	}
}

func TestCreateCertTemplateRoles(t *testing.T) {
	hosts := []string{"test.example.com", "test2.example.com"}
	mgmtIPs := []string{}
	cert, err := CreateCertTemplate(ssntp.AGENT|ssntp.NETAGENT, "ACME Corp", "test@example.com", hosts, mgmtIPs)

	if err != nil {
		t.Errorf("Unexpected error when creating cert template: %v", err)
	}

	if ssntp.GetRoleFromOIDs(cert.UnknownExtKeyUsage) != ssntp.AGENT|ssntp.NETAGENT {
		t.Error("Roles on created cert do not match those requested")
	}
}

func TestCreateAnchorCert(t *testing.T) {
	var certOutput, caCertOutput bytes.Buffer

	hosts := []string{"test.example.com", "test2.example.com"}
	mgmtIPs := []string{}

	template, err := CreateCertTemplate(ssntp.AGENT, "ACME Corp", "test@example.com", hosts, mgmtIPs)
	if err != nil {
		t.Errorf("Unexpected error when creating cert template: %v", err)
	}

	err = CreateAnchorCert(template, false, &certOutput, &caCertOutput)
	if err != nil {
		t.Errorf("Unexpected error when creating anchor cert: %v", err)
	}

	// Decode server cert & private key
	certBlock, rest := pem.Decode(certOutput.Bytes())
	cert, err := x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		t.Errorf("Failed to parse certificate: %v", err)
	}

	if cert.IsCA != true {
		t.Errorf("Expected certificate to be a CA")
	}

	privKeyBlock, _ := pem.Decode(rest)
	if privKeyBlock == nil {
		t.Errorf("Unable to extract private key from anchor cert")
	}

	anchorPrivKey, err := keyFromPemBlock(privKeyBlock)
	if err != nil {
		t.Errorf("Unable to parse private key from anchor cert: %v", err)
	}

	_, ok := anchorPrivKey.(*rsa.PrivateKey)
	if !ok || err != nil {
		t.Errorf("Expected RSA private key: %v", err)
	}

	// Decode CA
	certBlock, rest = pem.Decode(caCertOutput.Bytes())
	if len(rest) > 0 {
		t.Error("Unexpected data after certificate PEM block")
	}

	cert, err = x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		t.Errorf("Failed to parse certificate: %v", err)
	}

	if cert.IsCA != true {
		t.Errorf("Expected certificate to be a CA")
	}
}

func TestCreateCert(t *testing.T) {
	var anchorCertOutput, caCertOutput, certOutput bytes.Buffer

	hosts := []string{"test.example.com", "test2.example.com"}
	mgmtIPs := []string{}

	template, err := CreateCertTemplate(ssntp.AGENT, "ACME Corp", "test@example.com", hosts, mgmtIPs)
	if err != nil {
		t.Errorf("Unexpected error when creating cert template: %v", err)
	}

	err = CreateAnchorCert(template, false, &anchorCertOutput, &caCertOutput)
	if err != nil {
		t.Errorf("Unexpected error when creating anchor cert: %v", err)
	}

	err = CreateCert(template, false, anchorCertOutput.Bytes(), &certOutput)
	if err != nil {
		t.Errorf("Unexpected error when creating signed cert: %v", err)
	}

	// Decode signed cert & private key
	certBlock, rest := pem.Decode(certOutput.Bytes())
	privKeyBlock, _ := pem.Decode(rest)
	if privKeyBlock == nil {
		t.Errorf("Unable to extract private key from anchor cert")
	}

	anchorPrivKey, err := keyFromPemBlock(privKeyBlock)
	if err != nil {
		t.Errorf("Unable to parse private key from anchor cert: %v", err)
	}

	_, ok := anchorPrivKey.(*rsa.PrivateKey)
	if !ok || err != nil {
		t.Errorf("Expected RSA private key: %v", err)
	}

	// Verify cert is signed with CA cert
	cert, err := x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		t.Errorf("Failed to parse certificate: %v", err)
	}

	if cert.IsCA != false {
		t.Errorf("Expected certificate not to be a CA")
	}

	roots := x509.NewCertPool()
	ok = roots.AppendCertsFromPEM(caCertOutput.Bytes())
	if !ok {
		t.Errorf("Could not add CA cert to pool")
	}

	opts := x509.VerifyOptions{
		Roots: roots,
	}

	if _, err = cert.Verify(opts); err != nil {
		t.Errorf("Failed to verify certificate: %v", err)
	}
}
<commit_msg>ssntp/certs: Add testing for CSR flow<commit_after>//
// Copyright (c) 2016 Intel Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

package certs

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/x509"
	"encoding/asn1"
	"encoding/pem"
	"net"
	"reflect"
	"testing"

	"crypto/tls"

	"github.com/01org/ciao/ssntp"
)

func TestGenerateKey(t *testing.T) {
	key, err := generatePrivateKey(false)

	_, ok := key.(*rsa.PrivateKey)
	if !ok || err != nil {
		t.Errorf("RSA key expected from generatePrivateKey: %v", err)
	}

	pKey := publicKey(key)
	_, ok = pKey.(*rsa.PublicKey)
	if !ok {
		t.Error("RSA public key expected")
	}

	key, err = generatePrivateKey(true)
	_, ok = key.(*ecdsa.PrivateKey)

	if !ok || err != nil {
		t.Errorf("ECDSA key expected from generatePrivateKey: %v", err)
	}

	pKey = publicKey(key)
	_, ok = pKey.(*ecdsa.PublicKey)
	if !ok {
		t.Error("ECDSA public key expected")
	}
}

func TestPemBlockParsing(t *testing.T) {
	origKey, _ := generatePrivateKey(false)

	block, err := pemBlockForKey(origKey)

	if block.Type != "RSA PRIVATE KEY" || err != nil {
		t.Errorf("Unexpected PEM block type: %v: err: %v", block.Type, err)
	}

	parsedKey, err := keyFromPemBlock(block)
	_, ok := parsedKey.(*rsa.PrivateKey)

	if !ok || err != nil {
		t.Errorf("Expected RSA private key: %v", err)
	}

	origKey, _ = generatePrivateKey(true)

	block, err = pemBlockForKey(origKey)

	if block.Type != "EC PRIVATE KEY" || err != nil {
		t.Errorf("Unexpected PEM block type: %v: err: %v", block.Type, err)
	}

	parsedKey, err = keyFromPemBlock(block)
	_, ok = parsedKey.(*ecdsa.PrivateKey)

	if !ok || err != nil {
		t.Errorf("Expected ECDSA private key: %v", err)
	}
}

func TestAddOIDs(t *testing.T) {
	mapping := map[ssntp.Role]asn1.ObjectIdentifier{
		ssntp.AGENT:      ssntp.RoleAgentOID,
		ssntp.SCHEDULER:  ssntp.RoleSchedulerOID,
		ssntp.Controller: ssntp.RoleControllerOID,
		ssntp.NETAGENT:   ssntp.RoleNetAgentOID,
		ssntp.SERVER:     ssntp.RoleServerOID,
		ssntp.CNCIAGENT:  ssntp.RoleCNCIAgentOID,
	}

	// Check all
	for role, oid := range mapping {
		oids := addOIDs(role, []asn1.ObjectIdentifier{})
		if len(oids) != 1 {
			t.Errorf("Expected only one OID found %d", len(oids))
		}

		if !reflect.DeepEqual(oids[0], oid) {
			t.Errorf("Unexpected role to OID mapping: %v -> %v", role, oid)
		}
	}

	// Check common pairing
	oids := addOIDs(ssntp.AGENT|ssntp.NETAGENT, []asn1.ObjectIdentifier{})
	if len(oids) != 2 {
		t.Errorf("Expected two OIDS found %d", len(oids))
	}

	if !reflect.DeepEqual(oids[0], ssntp.RoleAgentOID) || !reflect.DeepEqual(oids[1], ssntp.RoleNetAgentOID) {
		t.Errorf("Unexpected OIDS in list: %v and %v", oids[0], oids[1])
	}
}

func TestCreateCertTemplateHosts(t *testing.T) {
	hosts := []string{"test.example.com", "test2.example.com"}
	mgmtIPs := []string{}
	cert, err := CreateCertTemplate(ssntp.AGENT, "ACME Corp", "test@example.com", hosts, mgmtIPs)

	if err != nil {
		t.Errorf("Unexpected error when creating cert template: %v", err)
	}

	if !reflect.DeepEqual(cert.DNSNames, hosts) {
		t.Errorf("Hosts in certificate don't match: %v %v", hosts, cert.DNSNames)
	}
}

func TestCreateCertTemplateIPs(t *testing.T) {
	hosts := []string{}
	mgmtIPs := []string{"127.0.0.1", "127.0.0.2"}
	cert, err := CreateCertTemplate(ssntp.AGENT, "ACME Corp", "test@example.com", hosts, mgmtIPs)

	if err != nil {
		t.Errorf("Unexpected error when creating cert template: %v", err)
	}

	IPs := []net.IP{net.IPv4(127, 0, 0, 1), net.IPv4(127, 0, 0, 2)}
	if !reflect.DeepEqual(cert.IPAddresses, IPs) {
		t.Errorf("IPs in certificate don't match: %v %v", hosts, cert.DNSNames)
	}
}

func TestCreateCertTemplateRoles(t *testing.T) {
	hosts := []string{"test.example.com", "test2.example.com"}
	mgmtIPs := []string{}
	cert, err := CreateCertTemplate(ssntp.AGENT|ssntp.NETAGENT, "ACME Corp", "test@example.com", hosts, mgmtIPs)

	if err != nil {
		t.Errorf("Unexpected error when creating cert template: %v", err)
	}

	if ssntp.GetRoleFromOIDs(cert.UnknownExtKeyUsage) != ssntp.AGENT|ssntp.NETAGENT {
		t.Error("Roles on created cert do not match those requested")
	}
}

func TestCreateAnchorCert(t *testing.T) {
	var certOutput, caCertOutput bytes.Buffer

	hosts := []string{"test.example.com", "test2.example.com"}
	mgmtIPs := []string{}

	template, err := CreateCertTemplate(ssntp.AGENT, "ACME Corp", "test@example.com", hosts, mgmtIPs)
	if err != nil {
		t.Errorf("Unexpected error when creating cert template: %v", err)
	}

	err = CreateAnchorCert(template, false, &certOutput, &caCertOutput)
	if err != nil {
		t.Errorf("Unexpected error when creating anchor cert: %v", err)
	}

	// Decode server cert & private key
	certBlock, rest := pem.Decode(certOutput.Bytes())
	cert, err := x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		t.Errorf("Failed to parse certificate: %v", err)
	}

	if cert.IsCA != true {
		t.Errorf("Expected certificate to be a CA")
	}

	privKeyBlock, _ := pem.Decode(rest)
	if privKeyBlock == nil {
		t.Errorf("Unable to extract private key from anchor cert")
	}

	anchorPrivKey, err := keyFromPemBlock(privKeyBlock)
	if err != nil {
		t.Errorf("Unable to parse private key from anchor cert: %v", err)
	}

	_, ok := anchorPrivKey.(*rsa.PrivateKey)
	if !ok || err != nil {
		t.Errorf("Expected RSA private key: %v", err)
	}

	// Decode CA
	certBlock, rest = pem.Decode(caCertOutput.Bytes())
	if len(rest) > 0 {
		t.Error("Unexpected data after certificate PEM block")
	}

	cert, err = x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		t.Errorf("Failed to parse certificate: %v", err)
	}

	if cert.IsCA != true {
		t.Errorf("Expected certificate to be a CA")
	}
}

func TestCreateCert(t *testing.T) {
	var anchorCertOutput, caCertOutput, certOutput bytes.Buffer

	hosts := []string{"test.example.com", "test2.example.com"}
	mgmtIPs := []string{}

	template, err := CreateCertTemplate(ssntp.AGENT, "ACME Corp", "test@example.com", hosts, mgmtIPs)
	if err != nil {
		t.Errorf("Unexpected error when creating cert template: %v", err)
	}

	err = CreateAnchorCert(template, false, &anchorCertOutput, &caCertOutput)
	if err != nil {
		t.Errorf("Unexpected error when creating anchor cert: %v", err)
	}

	err = CreateCert(template, false, anchorCertOutput.Bytes(), &certOutput)
	if err != nil {
		t.Errorf("Unexpected error when creating signed cert: %v", err)
	}

	// Decode signed cert & private key
	certBlock, rest := pem.Decode(certOutput.Bytes())
	privKeyBlock, _ := pem.Decode(rest)
	if privKeyBlock == nil {
		t.Errorf("Unable to extract private key from anchor cert")
	}

	anchorPrivKey, err := keyFromPemBlock(privKeyBlock)
	if err != nil {
		t.Errorf("Unable to parse private key from anchor cert: %v", err)
	}

	_, ok := anchorPrivKey.(*rsa.PrivateKey)
	if !ok || err != nil {
		t.Errorf("Expected RSA private key: %v", err)
	}

	// Verify cert is signed with CA cert
	cert, err := x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		t.Errorf("Failed to parse certificate: %v", err)
	}

	if cert.IsCA != false {
		t.Errorf("Expected certificate not to be a CA")
	}

	roots := x509.NewCertPool()
	ok = roots.AppendCertsFromPEM(caCertOutput.Bytes())
	if !ok {
		t.Errorf("Could not add CA cert to pool")
	}

	opts := x509.VerifyOptions{
		Roots: roots,
	}

	if _, err = cert.Verify(opts); err != nil {
		t.Errorf("Failed to verify certificate: %v", err)
	}
}

func TestCSRFlow(t *testing.T) {
	var anchorCertOutput, caCertOutput, certOutput, csrOutput, privKeyOutput, combinedOutput bytes.Buffer

	hosts := []string{"test.example.com", "test2.example.com"}
	mgmtIPs := []string{}

	template, err := CreateCertTemplate(ssntp.AGENT, "ACME Corp", "test@example.com", hosts, mgmtIPs)
	if err != nil {
		t.Fatalf("Unexpected error when creating cert template: %v", err)
	}

	err = CreateAnchorCert(template, true, &anchorCertOutput, &caCertOutput)
	if err != nil {
		t.Fatalf("Unexpected error when creating anchor cert: %v", err)
	}

	request := CreateCertificateRequest(true, "ACME Corp", "test@example.com", hosts, mgmtIPs)

	err = CreateCSR(request, true, &csrOutput, &privKeyOutput)
	if err != nil {
		t.Fatalf("Unexpected error when creating CSR: %v", err)
	}

	err = CreateCertFromCSR(ssntp.AGENT, csrOutput.Bytes(), anchorCertOutput.Bytes(), &certOutput)
	if err != nil {
		t.Fatalf("Unexpected error when creating cert from CSR: %v", err)
	}

	err = AddPrivateKeyToCert(&certOutput, &privKeyOutput, &combinedOutput)
	if err != nil {
		t.Fatalf("Unexpected error when merging cert and private key")
	}

	err = VerifyCert(anchorCertOutput.Bytes(), combinedOutput.Bytes())
	if err != nil {
		t.Fatalf("Unexpected error when verifying merged cert: %v", err)
	}

	// Test that this cert is useful for ssntp TLS
	_, err = tls.X509KeyPair(combinedOutput.Bytes(), combinedOutput.Bytes())
	if err != nil {
		t.Fatalf("Unexpected error when checking merged cert: %v", err)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2014 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"runtime"
	"runtime/pprof"

	"github.com/btcsuite/btcd/limits"
)

var (
	cfg             *config
	shutdownChannel = make(chan struct{})
)

// winServiceMain is only invoked on Windows.  It detects when btcd is running
// as a service and reacts accordingly.
var winServiceMain func() (bool, error)

// btcdMain is the real main function for btcd.  It is necessary to work around
// the fact that deferred functions do not run when os.Exit() is called.  The
// optional serverChan parameter is mainly used by the service code to be
// notified with the server once it is setup so it can gracefully stop it when
// requested from the service control manager.
func btcdMain(serverChan chan<- *server) error {
	// Load configuration and parse command line.  This function also
	// initializes logging and configures it accordingly.
	tcfg, _, err := loadConfig()
	if err != nil {
		return err
	}
	cfg = tcfg
	defer backendLog.Flush()

	// Show version at startup.
	btcdLog.Infof("Version %s", version())

	// Enable http profiling server if requested.
	if cfg.Profile != "" {
		go func() {
			listenAddr := net.JoinHostPort("", cfg.Profile)
			btcdLog.Infof("Profile server listening on %s", listenAddr)
			profileRedirect := http.RedirectHandler("/debug/pprof",
				http.StatusSeeOther)
			http.Handle("/", profileRedirect)
			btcdLog.Errorf("%v", http.ListenAndServe(listenAddr, nil))
		}()
	}

	// Write cpu profile if requested.
	if cfg.CPUProfile != "" {
		f, err := os.Create(cfg.CPUProfile)
		if err != nil {
			btcdLog.Errorf("Unable to create cpu profile: %v", err)
			return err
		}
		pprof.StartCPUProfile(f)
		defer f.Close()
		defer pprof.StopCPUProfile()
	}

	// Perform upgrades to btcd as new versions require it.
	if err := doUpgrades(); err != nil {
		btcdLog.Errorf("%v", err)
		return err
	}

	// Load the block database.
	db, err := loadBlockDB()
	if err != nil {
		btcdLog.Errorf("%v", err)
		return err
	}
	defer db.Close()

	if cfg.DropAddrIndex {
		btcdLog.Info("Deleting entire addrindex.")
		err := db.DeleteAddrIndex()
		if err != nil {
			btcdLog.Errorf("Unable to delete the addrindex: %v", err)
			return err
		}
		btcdLog.Info("Successfully deleted addrindex, exiting")
		return nil
	}

	// Ensure the database is sync'd and closed on Ctrl+C.
	addInterruptHandler(func() {
		btcdLog.Infof("Gracefully shutting down the database...")
		db.RollbackClose()
	})

	// Create server and start it.
	server, err := newServer(cfg.Listeners, db, activeNetParams.Params)
	if err != nil {
		// TODO(oga) this logging could do with some beautifying.
		btcdLog.Errorf("Unable to start server on %v: %v",
			cfg.Listeners, err)
		return err
	}
	addInterruptHandler(func() {
		btcdLog.Infof("Gracefully shutting down the server...")
		server.Stop()
		server.WaitForShutdown()
	})
	server.Start()
	if serverChan != nil {
		serverChan <- server
	}

	// Monitor for graceful server shutdown and signal the main goroutine
	// when done.  This is done in a separate goroutine rather than waiting
	// directly so the main goroutine can be signaled for shutdown by either
	// a graceful shutdown or from the main interrupt handler.  This is
	// necessary since the main goroutine must be kept running long enough
	// for the interrupt handler goroutine to finish.
	go func() {
		server.WaitForShutdown()
		srvrLog.Infof("Server shutdown complete")
		shutdownChannel <- struct{}{}
	}()

	// Wait for shutdown signal from either a graceful server stop or from
	// the interrupt handler.
	<-shutdownChannel
	btcdLog.Info("Shutdown complete")
	return nil
}

func main() {
	// Use all processor cores.
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Up some limits.
	if err := limits.SetLimits(); err != nil {
		os.Exit(1)
	}

	// Call serviceMain on Windows to handle running as a service.  When
	// the return isService flag is true, exit now since we ran as a
	// service.  Otherwise, just fall through to normal operation.
	if runtime.GOOS == "windows" {
		isService, err := winServiceMain()
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		if isService {
			os.Exit(0)
		}
	}

	// Work around defer not working after os.Exit()
	if err := btcdMain(nil); err != nil {
		os.Exit(1)
	}
}
<commit_msg>Output error to stderr if the limits can't be set.<commit_after>// Copyright (c) 2013-2014 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"runtime"
	"runtime/pprof"

	"github.com/btcsuite/btcd/limits"
)

var (
	cfg             *config
	shutdownChannel = make(chan struct{})
)

// winServiceMain is only invoked on Windows.  It detects when btcd is running
// as a service and reacts accordingly.
var winServiceMain func() (bool, error)

// btcdMain is the real main function for btcd.  It is necessary to work around
// the fact that deferred functions do not run when os.Exit() is called.  The
// optional serverChan parameter is mainly used by the service code to be
// notified with the server once it is setup so it can gracefully stop it when
// requested from the service control manager.
func btcdMain(serverChan chan<- *server) error {
	// Load configuration and parse command line.  This function also
	// initializes logging and configures it accordingly.
	tcfg, _, err := loadConfig()
	if err != nil {
		return err
	}
	cfg = tcfg
	defer backendLog.Flush()

	// Show version at startup.
	btcdLog.Infof("Version %s", version())

	// Enable http profiling server if requested.
	if cfg.Profile != "" {
		go func() {
			listenAddr := net.JoinHostPort("", cfg.Profile)
			btcdLog.Infof("Profile server listening on %s", listenAddr)
			profileRedirect := http.RedirectHandler("/debug/pprof",
				http.StatusSeeOther)
			http.Handle("/", profileRedirect)
			btcdLog.Errorf("%v", http.ListenAndServe(listenAddr, nil))
		}()
	}

	// Write cpu profile if requested.
	if cfg.CPUProfile != "" {
		f, err := os.Create(cfg.CPUProfile)
		if err != nil {
			btcdLog.Errorf("Unable to create cpu profile: %v", err)
			return err
		}
		pprof.StartCPUProfile(f)
		defer f.Close()
		defer pprof.StopCPUProfile()
	}

	// Perform upgrades to btcd as new versions require it.
	if err := doUpgrades(); err != nil {
		btcdLog.Errorf("%v", err)
		return err
	}

	// Load the block database.
	db, err := loadBlockDB()
	if err != nil {
		btcdLog.Errorf("%v", err)
		return err
	}
	defer db.Close()

	if cfg.DropAddrIndex {
		btcdLog.Info("Deleting entire addrindex.")
		err := db.DeleteAddrIndex()
		if err != nil {
			btcdLog.Errorf("Unable to delete the addrindex: %v", err)
			return err
		}
		btcdLog.Info("Successfully deleted addrindex, exiting")
		return nil
	}

	// Ensure the database is sync'd and closed on Ctrl+C.
	addInterruptHandler(func() {
		btcdLog.Infof("Gracefully shutting down the database...")
		db.RollbackClose()
	})

	// Create server and start it.
	server, err := newServer(cfg.Listeners, db, activeNetParams.Params)
	if err != nil {
		// TODO(oga) this logging could do with some beautifying.
		btcdLog.Errorf("Unable to start server on %v: %v",
			cfg.Listeners, err)
		return err
	}
	addInterruptHandler(func() {
		btcdLog.Infof("Gracefully shutting down the server...")
		server.Stop()
		server.WaitForShutdown()
	})
	server.Start()
	if serverChan != nil {
		serverChan <- server
	}

	// Monitor for graceful server shutdown and signal the main goroutine
	// when done.  This is done in a separate goroutine rather than waiting
	// directly so the main goroutine can be signaled for shutdown by either
	// a graceful shutdown or from the main interrupt handler.  This is
	// necessary since the main goroutine must be kept running long enough
	// for the interrupt handler goroutine to finish.
	go func() {
		server.WaitForShutdown()
		srvrLog.Infof("Server shutdown complete")
		shutdownChannel <- struct{}{}
	}()

	// Wait for shutdown signal from either a graceful server stop or from
	// the interrupt handler.
	<-shutdownChannel
	btcdLog.Info("Shutdown complete")
	return nil
}

func main() {
	// Use all processor cores.
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Up some limits.
	if err := limits.SetLimits(); err != nil {
		fmt.Fprintf(os.Stderr, "failed to set limits: %v\n", err)
		os.Exit(1)
	}

	// Call serviceMain on Windows to handle running as a service.  When
	// the return isService flag is true, exit now since we ran as a
	// service.  Otherwise, just fall through to normal operation.
	if runtime.GOOS == "windows" {
		isService, err := winServiceMain()
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		if isService {
			os.Exit(0)
		}
	}

	// Work around defer not working after os.Exit()
	if err := btcdMain(nil); err != nil {
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package mgo

import (
	"bytes"
	"gopkg.in/mgo.v2-unstable/bson"
	"sort"
)

// Bulk represents an operation that can be prepared with several
// orthogonal changes before being delivered to the server.
//
// MongoDB servers older than version 2.6 do not have proper support for bulk
// operations, so the driver attempts to map its API as much as possible into
// the functionality that works. In particular, in those releases updates and
// removals are sent individually, and inserts are sent in bulk but have
// suboptimal error reporting compared to more recent versions of the server.
// See the documentation of BulkErrorCase for details on that.
//
// Relevant documentation:
//
//   http://blog.mongodb.org/post/84922794768/mongodbs-new-bulk-api
//
type Bulk struct {
	c       *Collection
	opcount int
	actions []bulkAction
	ordered bool
}

type bulkOp int

const (
	bulkInsert bulkOp = iota + 1
	bulkUpdate
	bulkUpdateAll
	bulkRemove
)

type bulkAction struct {
	op   bulkOp
	docs []interface{}
	idxs []int
}

type bulkUpdateOp []interface{}
type bulkDeleteOp []interface{}

// BulkResult holds the results for a bulk operation.
type BulkResult struct {
	Matched  int
	Modified int // Available only for MongoDB 2.6+

	// Be conservative while we understand exactly how to report these
	// results in a useful and convenient way, and also how to emulate
	// them with prior servers.
	private bool
}

// BulkError holds an error returned from running a Bulk operation.
// Individual errors may be obtained and inspected via the Cases method.
type BulkError struct {
	ecases []BulkErrorCase
}

func (e *BulkError) Error() string {
	if len(e.ecases) == 0 {
		return "invalid BulkError instance: no errors"
	}
	if len(e.ecases) == 1 {
		return e.ecases[0].Err.Error()
	}
	msgs := make([]string, 0, len(e.ecases))
	seen := make(map[string]bool)
	for _, ecase := range e.ecases {
		msg := ecase.Err.Error()
		if !seen[msg] {
			seen[msg] = true
			msgs = append(msgs, msg)
		}
	}
	if len(msgs) == 1 {
		return msgs[0]
	}
	var buf bytes.Buffer
	buf.WriteString("multiple errors in bulk operation:\n")
	for _, msg := range msgs {
		buf.WriteString("  - ")
		buf.WriteString(msg)
		buf.WriteByte('\n')
	}
	return buf.String()
}

type bulkErrorCases []BulkErrorCase

func (slice bulkErrorCases) Len() int           { return len(slice) }
func (slice bulkErrorCases) Less(i, j int) bool { return slice[i].Index < slice[j].Index }
func (slice bulkErrorCases) Swap(i, j int)      { slice[i], slice[j] = slice[j], slice[i] }

// BulkErrorCase holds an individual error found while attempting a single change
// within a bulk operation, and the position in which it was enqueued.
//
// MongoDB servers older than version 2.6 do not have proper support for bulk
// operations, so the driver attempts to map its API as much as possible into
// the functionality that works. In particular, only the last error is reported
// for bulk inserts and without any positional information, so the Index
// field is set to -1 in these cases.
type BulkErrorCase struct {
	Index int   // Position of operation that failed, or -1 if unknown.
	Err   error
}

// Cases returns all individual errors found while attempting the requested changes.
//
// See the documentation of BulkErrorCase for limitations in older MongoDB releases.
func (e *BulkError) Cases() []BulkErrorCase {
	return e.ecases
}

// Bulk returns a value to prepare the execution of a bulk operation.
func (c *Collection) Bulk() *Bulk {
	return &Bulk{c: c, ordered: true}
}

// Unordered puts the bulk operation in unordered mode.
//
// In unordered mode the indvidual operations may be sent
// out of order, which means latter operations may proceed
// even if prior ones have failed.
func (b *Bulk) Unordered() {
	b.ordered = false
}

func (b *Bulk) action(op bulkOp, opcount int) *bulkAction {
	var action *bulkAction
	if len(b.actions) > 0 && b.actions[len(b.actions)-1].op == op {
		action = &b.actions[len(b.actions)-1]
	} else if !b.ordered {
		for i := range b.actions {
			if b.actions[i].op == op {
				action = &b.actions[i]
				break
			}
		}
	}
	if action == nil {
		b.actions = append(b.actions, bulkAction{op: op})
		action = &b.actions[len(b.actions)-1]
	}
	for i := 0; i < opcount; i++ {
		action.idxs = append(action.idxs, b.opcount)
		b.opcount++
	}
	return action
}

// Insert queues up the provided documents for insertion.
func (b *Bulk) Insert(docs ...interface{}) {
	action := b.action(bulkInsert, len(docs))
	action.docs = append(action.docs, docs...)
}

// Remove queues up the provided selectors for removing matching documents.
// Each selector will remove only a single matching document.
func (b *Bulk) Remove(selectors ...interface{}) {
	action := b.action(bulkRemove, len(selectors))
	for _, selector := range selectors {
		if selector == nil {
			selector = bson.D{}
		}
		action.docs = append(action.docs, &deleteOp{
			Collection: b.c.FullName,
			Selector:   selector,
			Flags:      1,
			Limit:      1,
		})
	}
}

// RemoveAll queues up the provided selectors for removing all matching documents.
// Each selector will remove all matching documents.
func (b *Bulk) RemoveAll(selectors ...interface{}) {
	action := b.action(bulkRemove, len(selectors))
	for _, selector := range selectors {
		if selector == nil {
			selector = bson.D{}
		}
		action.docs = append(action.docs, &deleteOp{
			Collection: b.c.FullName,
			Selector:   selector,
			Flags:      0,
			Limit:      0,
		})
	}
}

// Update queues up the provided pairs of updating instructions.
// The first element of each pair selects which documents must be
// updated, and the second element defines how to update it.
// Each pair matches exactly one document for updating at most.
func (b *Bulk) Update(pairs ...interface{}) {
	if len(pairs)%2 != 0 {
		panic("Bulk.Update requires an even number of parameters")
	}
	action := b.action(bulkUpdate, len(pairs)/2)
	for i := 0; i < len(pairs); i += 2 {
		selector := pairs[i]
		if selector == nil {
			selector = bson.D{}
		}
		action.docs = append(action.docs, &updateOp{
			Collection: b.c.FullName,
			Selector:   selector,
			Update:     pairs[i+1],
		})
	}
}

// UpdateAll queues up the provided pairs of updating instructions.
// The first element of each pair selects which documents must be
// updated, and the second element defines how to update it.
// Each pair updates all documents matching the selector.
func (b *Bulk) UpdateAll(pairs ...interface{}) {
	if len(pairs)%2 != 0 {
		panic("Bulk.UpdateAll requires an even number of parameters")
	}
	action := b.action(bulkUpdate, len(pairs)/2)
	for i := 0; i < len(pairs); i += 2 {
		selector := pairs[i]
		if selector == nil {
			selector = bson.D{}
		}
		action.docs = append(action.docs, &updateOp{
			Collection: b.c.FullName,
			Selector:   selector,
			Update:     pairs[i+1],
			Flags:      2,
			Multi:      true,
		})
	}
}

// Upsert queues up the provided pairs of upserting instructions.
// The first element of each pair selects which documents must be
// updated, and the second element defines how to update it.
// Each pair matches exactly one document for updating at most.
func (b *Bulk) Upsert(pairs ...interface{}) {
	if len(pairs)%2 != 0 {
		panic("Bulk.Update requires an even number of parameters")
	}
	action := b.action(bulkUpdate, len(pairs)/2)
	for i := 0; i < len(pairs); i += 2 {
		selector := pairs[i]
		if selector == nil {
			selector = bson.D{}
		}
		action.docs = append(action.docs, &updateOp{
			Collection: b.c.FullName,
			Selector:   selector,
			Update:     pairs[i+1],
			Flags:      1,
			Upsert:     true,
		})
	}
}

// Run runs all the operations queued up.
//
// If an error is reported on an unordered bulk operation, the error value may
// be an aggregation of all issues observed. As an exception to that, Insert
// operations running on MongoDB versions prior to 2.6 will report the last
// error only due to a limitation in the wire protocol.
func (b *Bulk) Run() (*BulkResult, error) {
	var result BulkResult
	var berr BulkError
	var failed bool
	for i := range b.actions {
		action := &b.actions[i]
		var ok bool
		switch action.op {
		case bulkInsert:
			ok = b.runInsert(action, &result, &berr)
		case bulkUpdate:
			ok = b.runUpdate(action, &result, &berr)
		case bulkRemove:
			ok = b.runRemove(action, &result, &berr)
		default:
			panic("unknown bulk operation")
		}
		if !ok {
			failed = true
			if b.ordered {
				break
			}
		}
	}
	if failed {
		sort.Sort(bulkErrorCases(berr.ecases))
		return nil, &berr
	}
	return &result, nil
}

func (b *Bulk) runInsert(action *bulkAction, result *BulkResult, berr *BulkError) bool {
	op := &insertOp{b.c.FullName, action.docs, 0}
	if !b.ordered {
		op.flags = 1 // ContinueOnError
	}
	lerr, err := b.c.writeOp(op, b.ordered)
	return b.checkSuccess(action, berr, lerr, err)
}

func (b *Bulk) runUpdate(action *bulkAction, result *BulkResult, berr *BulkError) bool {
	lerr, err := b.c.writeOp(bulkUpdateOp(action.docs), b.ordered)
	result.Matched += lerr.N
	result.Modified += lerr.modified
	return b.checkSuccess(action, berr, lerr, err)
}

func (b *Bulk) runRemove(action *bulkAction, result *BulkResult, berr *BulkError) bool {
	lerr, err := b.c.writeOp(bulkDeleteOp(action.docs), b.ordered)
	result.Matched += lerr.N
	result.Modified += lerr.modified
	return b.checkSuccess(action, berr, lerr, err)
}

func (b *Bulk) checkSuccess(action *bulkAction, berr *BulkError, lerr *LastError, err error) bool {
	if lerr != nil && len(lerr.ecases) > 0 {
		for i := 0; i < len(lerr.ecases); i++ {
			// Map back from the local error index into the visible one.
			ecase := lerr.ecases[i]
			idx := ecase.Index
			if idx >= 0 {
				idx = action.idxs[idx]
			}
			berr.ecases = append(berr.ecases, BulkErrorCase{idx, ecase.Err})
		}
		return false
	} else if err != nil {
		for i := 0; i < len(action.idxs); i++ {
			berr.ecases = append(berr.ecases, BulkErrorCase{action.idxs[i], err})
		}
		return false
	}
	return true
}
<commit_msg>Fix crash on bulk update and delete returned error.<commit_after>package mgo

import (
	"bytes"
	"gopkg.in/mgo.v2-unstable/bson"
	"sort"
)

// Bulk represents an operation that can be prepared with several
// orthogonal changes before being delivered to the server.
//
// MongoDB servers older than version 2.6 do not have proper support for bulk
// operations, so the driver attempts to map its API as much as possible into
// the functionality that works. In particular, in those releases updates and
// removals are sent individually, and inserts are sent in bulk but have
// suboptimal error reporting compared to more recent versions of the server.
// See the documentation of BulkErrorCase for details on that.
//
// Relevant documentation:
//
//   http://blog.mongodb.org/post/84922794768/mongodbs-new-bulk-api
//
type Bulk struct {
	c       *Collection
	opcount int
	actions []bulkAction
	ordered bool
}

type bulkOp int

const (
	bulkInsert bulkOp = iota + 1
	bulkUpdate
	bulkUpdateAll
	bulkRemove
)

type bulkAction struct {
	op   bulkOp
	docs []interface{}
	idxs []int
}

type bulkUpdateOp []interface{}
type bulkDeleteOp []interface{}

// BulkResult holds the results for a bulk operation.
type BulkResult struct {
	Matched  int
	Modified int // Available only for MongoDB 2.6+

	// Be conservative while we understand exactly how to report these
	// results in a useful and convenient way, and also how to emulate
	// them with prior servers.
	private bool
}

// BulkError holds an error returned from running a Bulk operation.
// Individual errors may be obtained and inspected via the Cases method.
type BulkError struct {
	ecases []BulkErrorCase
}

func (e *BulkError) Error() string {
	if len(e.ecases) == 0 {
		return "invalid BulkError instance: no errors"
	}
	if len(e.ecases) == 1 {
		return e.ecases[0].Err.Error()
	}
	msgs := make([]string, 0, len(e.ecases))
	seen := make(map[string]bool)
	for _, ecase := range e.ecases {
		msg := ecase.Err.Error()
		if !seen[msg] {
			seen[msg] = true
			msgs = append(msgs, msg)
		}
	}
	if len(msgs) == 1 {
		return msgs[0]
	}
	var buf bytes.Buffer
	buf.WriteString("multiple errors in bulk operation:\n")
	for _, msg := range msgs {
		buf.WriteString("  - ")
		buf.WriteString(msg)
		buf.WriteByte('\n')
	}
	return buf.String()
}

type bulkErrorCases []BulkErrorCase

func (slice bulkErrorCases) Len() int           { return len(slice) }
func (slice bulkErrorCases) Less(i, j int) bool { return slice[i].Index < slice[j].Index }
func (slice bulkErrorCases) Swap(i, j int)      { slice[i], slice[j] = slice[j], slice[i] }

// BulkErrorCase holds an individual error found while attempting a single change
// within a bulk operation, and the position in which it was enqueued.
//
// MongoDB servers older than version 2.6 do not have proper support for bulk
// operations, so the driver attempts to map its API as much as possible into
// the functionality that works. In particular, only the last error is reported
// for bulk inserts and without any positional information, so the Index
// field is set to -1 in these cases.
type BulkErrorCase struct {
	Index int   // Position of operation that failed, or -1 if unknown.
	Err   error
}

// Cases returns all individual errors found while attempting the requested changes.
//
// See the documentation of BulkErrorCase for limitations in older MongoDB releases.
func (e *BulkError) Cases() []BulkErrorCase {
	return e.ecases
}

// Bulk returns a value to prepare the execution of a bulk operation.
func (c *Collection) Bulk() *Bulk {
	return &Bulk{c: c, ordered: true}
}

// Unordered puts the bulk operation in unordered mode.
//
// In unordered mode the indvidual operations may be sent
// out of order, which means latter operations may proceed
// even if prior ones have failed.
func (b *Bulk) Unordered() {
	b.ordered = false
}

func (b *Bulk) action(op bulkOp, opcount int) *bulkAction {
	var action *bulkAction
	if len(b.actions) > 0 && b.actions[len(b.actions)-1].op == op {
		action = &b.actions[len(b.actions)-1]
	} else if !b.ordered {
		for i := range b.actions {
			if b.actions[i].op == op {
				action = &b.actions[i]
				break
			}
		}
	}
	if action == nil {
		b.actions = append(b.actions, bulkAction{op: op})
		action = &b.actions[len(b.actions)-1]
	}
	for i := 0; i < opcount; i++ {
		action.idxs = append(action.idxs, b.opcount)
		b.opcount++
	}
	return action
}

// Insert queues up the provided documents for insertion.
func (b *Bulk) Insert(docs ...interface{}) {
	action := b.action(bulkInsert, len(docs))
	action.docs = append(action.docs, docs...)
}

// Remove queues up the provided selectors for removing matching documents.
// Each selector will remove only a single matching document.
func (b *Bulk) Remove(selectors ...interface{}) {
	action := b.action(bulkRemove, len(selectors))
	for _, selector := range selectors {
		if selector == nil {
			selector = bson.D{}
		}
		action.docs = append(action.docs, &deleteOp{
			Collection: b.c.FullName,
			Selector:   selector,
			Flags:      1,
			Limit:      1,
		})
	}
}

// RemoveAll queues up the provided selectors for removing all matching documents.
// Each selector will remove all matching documents.
func (b *Bulk) RemoveAll(selectors ...interface{}) {
	action := b.action(bulkRemove, len(selectors))
	for _, selector := range selectors {
		if selector == nil {
			selector = bson.D{}
		}
		action.docs = append(action.docs, &deleteOp{
			Collection: b.c.FullName,
			Selector:   selector,
			Flags:      0,
			Limit:      0,
		})
	}
}

// Update queues up the provided pairs of updating instructions.
// The first element of each pair selects which documents must be
// updated, and the second element defines how to update it.
// Each pair matches exactly one document for updating at most.
func (b *Bulk) Update(pairs ...interface{}) {
	if len(pairs)%2 != 0 {
		panic("Bulk.Update requires an even number of parameters")
	}
	action := b.action(bulkUpdate, len(pairs)/2)
	for i := 0; i < len(pairs); i += 2 {
		selector := pairs[i]
		if selector == nil {
			selector = bson.D{}
		}
		action.docs = append(action.docs, &updateOp{
			Collection: b.c.FullName,
			Selector:   selector,
			Update:     pairs[i+1],
		})
	}
}

// UpdateAll queues up the provided pairs of updating instructions.
// The first element of each pair selects which documents must be
// updated, and the second element defines how to update it.
// Each pair updates all documents matching the selector.
func (b *Bulk) UpdateAll(pairs ...interface{}) {
	if len(pairs)%2 != 0 {
		panic("Bulk.UpdateAll requires an even number of parameters")
	}
	action := b.action(bulkUpdate, len(pairs)/2)
	for i := 0; i < len(pairs); i += 2 {
		selector := pairs[i]
		if selector == nil {
			selector = bson.D{}
		}
		action.docs = append(action.docs, &updateOp{
			Collection: b.c.FullName,
			Selector:   selector,
			Update:     pairs[i+1],
			Flags:      2,
			Multi:      true,
		})
	}
}

// Upsert queues up the provided pairs of upserting instructions.
// The first element of each pair selects which documents must be
// updated, and the second element defines how to update it.
// Each pair matches exactly one document for updating at most.
func (b *Bulk) Upsert(pairs ...interface{}) {
	if len(pairs)%2 != 0 {
		panic("Bulk.Update requires an even number of parameters")
	}
	action := b.action(bulkUpdate, len(pairs)/2)
	for i := 0; i < len(pairs); i += 2 {
		selector := pairs[i]
		if selector == nil {
			selector = bson.D{}
		}
		action.docs = append(action.docs, &updateOp{
			Collection: b.c.FullName,
			Selector:   selector,
			Update:     pairs[i+1],
			Flags:      1,
			Upsert:     true,
		})
	}
}

// Run runs all the operations queued up.
//
// If an error is reported on an unordered bulk operation, the error value may
// be an aggregation of all issues observed. As an exception to that, Insert
// operations running on MongoDB versions prior to 2.6 will report the last
// error only due to a limitation in the wire protocol.
func (b *Bulk) Run() (*BulkResult, error) {
	var result BulkResult
	var berr BulkError
	var failed bool
	for i := range b.actions {
		action := &b.actions[i]
		var ok bool
		switch action.op {
		case bulkInsert:
			ok = b.runInsert(action, &result, &berr)
		case bulkUpdate:
			ok = b.runUpdate(action, &result, &berr)
		case bulkRemove:
			ok = b.runRemove(action, &result, &berr)
		default:
			panic("unknown bulk operation")
		}
		if !ok {
			failed = true
			if b.ordered {
				break
			}
		}
	}
	if failed {
		sort.Sort(bulkErrorCases(berr.ecases))
		return nil, &berr
	}
	return &result, nil
}

func (b *Bulk) runInsert(action *bulkAction, result *BulkResult, berr *BulkError) bool {
	op := &insertOp{b.c.FullName, action.docs, 0}
	if !b.ordered {
		op.flags = 1 // ContinueOnError
	}
	lerr, err := b.c.writeOp(op, b.ordered)
	return b.checkSuccess(action, berr, lerr, err)
}

func (b *Bulk) runUpdate(action *bulkAction, result *BulkResult, berr *BulkError) bool {
	lerr, err := b.c.writeOp(bulkUpdateOp(action.docs), b.ordered)
	if lerr != nil {
		result.Matched += lerr.N
		result.Modified += lerr.modified
	}
	return b.checkSuccess(action, berr, lerr, err)
}

func (b *Bulk) runRemove(action *bulkAction, result *BulkResult, berr *BulkError) bool {
	lerr, err := b.c.writeOp(bulkDeleteOp(action.docs), b.ordered)
	if lerr != nil {
		result.Matched += lerr.N
		result.Modified += lerr.modified
	}
	return b.checkSuccess(action, berr, lerr, err)
}

func (b *Bulk) checkSuccess(action *bulkAction, berr *BulkError, lerr *LastError, err error) bool {
	if lerr != nil && len(lerr.ecases) > 0 {
		for i := 0; i < len(lerr.ecases); i++ {
			// Map back from the local error index into the visible one.
			ecase := lerr.ecases[i]
			idx := ecase.Index
			if idx >= 0 {
				idx = action.idxs[idx]
			}
			berr.ecases = append(berr.ecases, BulkErrorCase{idx, ecase.Err})
		}
		return false
	} else if err != nil {
		for i := 0; i < len(action.idxs); i++ {
			berr.ecases = append(berr.ecases, BulkErrorCase{action.idxs[i], err})
		}
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package apns

import (
	"crypto/tls"
	"errors"
	"io"
	"log"
	"net"
	"strings"
	"sync"
	"time"

	"code.google.com/p/go.net/context"
)

type client struct {
	gateway           string
	certificateFile   string
	certificateBase64 string
	keyFile           string
	keyBase64         string
}

// PersistentClient opens a persistent connexion with the gateway
type PersistentClient struct {
	client *client
	conn   net.Conn
	ip     string

	mu          sync.Mutex
	isConnected bool
}

// NewPersistentClient creates a new persistent connection to the APNs servers
func NewPersistentClient(gateway, ip, certificateFile, keyFile string) (*PersistentClient, error) {

	var c *PersistentClient = &PersistentClient{}
	c.client = &client{gateway: gateway, certificateFile: certificateFile, keyFile: keyFile}
	c.ip = ip
	err := c.Connect()
	if err != nil {
		return nil, err
	}
	return c, err
}

// Connect connects the persistent client to one of the APNs server
// If the connection is already established and was not closed, it does nothing.
func (c *PersistentClient) Connect() error {

	// Check if there is already a connection
	if c.isConnected == true {
		// If connection is not nil it should be ok
		// c.conn is set to nil when there is an error on read or write
		// because the gateway close it anyway in this case
		return nil
	}
	return c.Reconnect()
}

// Reconnect forces a new connection to the gateway
// If a connection exists it is closed before the creation of a new one
func (c *PersistentClient) Reconnect() error {

	var cert tls.Certificate
	var err error

	if c.isConnected == true {
		c.Close()
	}

	if len(c.client.certificateBase64) == 0 && len(c.client.keyBase64) == 0 {
		// The user did not specify raw block contents, so check the filesystem.
		cert, err = tls.LoadX509KeyPair(c.client.certificateFile, c.client.keyFile)
	} else {
		// The user provided the raw block contents, so use that.
		cert, err = tls.X509KeyPair([]byte(c.client.certificateBase64), []byte(c.client.keyBase64))
	}
	if err != nil {
		return err
	}
	gatewayParts := strings.Split(c.client.gateway, ":")
	conf := &tls.Config{
		Certificates: []tls.Certificate{cert},
		ServerName:   gatewayParts[0],
	}
	if c.ip == "" { // If the ip is not provided pick one
		ip, err := pickGatewayIP(gatewayParts[0])
		if err != nil {
			return err
		}
		c.ip = ip
	}
	conn, err := net.Dial("tcp", c.ip+":"+gatewayParts[1])
	if err != nil {
		return err
	}
	tlsConn := tls.Client(conn, conf)
	err = tlsConn.Handshake()
	if err != nil {
		conn.Close()
		return err
	}
	c.mu.Lock()
	c.conn = net.Conn(tlsConn)
	c.isConnected = true
	c.mu.Unlock()
	log.Printf("Address of %s is %s", c.client.gateway, c.conn.RemoteAddr().String())
	return nil
}

// Send sends push notification to the APNs.
func (c *PersistentClient) Send(ctx context.Context, pn *PushNotification) *PushNotificationResponse {

	resp := NewPushNotificationResponse(pn)
	payload, err := pn.ToBytes()
	if err != nil {
		resp.Success = false
		resp.Error = err
		return resp
	}

	_, err = c.Write(payload)
	if err != nil {
		resp.Success = false
		resp.ResponseCommand = LocalResponseCommand
		resp.ResponseStatus = RetryPushNotificationStatus
		resp.Error = err
		return resp
	}
	log.Println("Sending push notification with ID", pn.Identifier)

	// This channel will contain the raw response
	// from Apple in the event of a failure.
	responseChannel := make(chan []byte, 1)
	go func() {
		buffer := make([]byte, 6)
		n, err := c.Read(buffer)
		if n != 6 && err != nil {
			buffer[0] = LocalResponseCommand
			e, ok := err.(net.Error)
			switch {
			case err == io.EOF: // Socket has been closed
				buffer[1] = RetryPushNotificationStatus
			case ok && e.Timeout(): // There is an error and it is a timeout
				buffer[1] = NoErrorsStatus
			default:
				buffer[1] = UnknownErrorStatus
			}
		}
		responseChannel <- buffer
	}()

	select {
	case <-ctx.Done():
		<-responseChannel // Wait for the read to end.
		resp.Success = false
		resp.ResponseCommand = LocalResponseCommand
		resp.ResponseStatus = CanceledPushNotificationStatus
		resp.Error = ctx.Err()
	case r := <-responseChannel:
		resp.FromRawAppleResponse(r)
	}
	return resp
}

// Close closes the persistent client
func (c *PersistentClient) Close() {
	c.closeAndSetDisconnected()
}

// closeAndSetDisconnected closes a persistent connection and set the isConnected flag to false
func (c *PersistentClient) closeAndSetDisconnected() {
	log.Printf("Closing %s at address %s", c.client.gateway, c.conn.RemoteAddr().String())
	c.mu.Lock()
	defer c.mu.Unlock()
	c.conn.Close()
	c.isConnected = false
}

func (c *PersistentClient) Write(b []byte) (n int, err error) {

	if !c.isConnected {
		return 0, errors.New("persistentclient: write to closed conn")
	}
	return c.conn.Write(b)
}

func (c *PersistentClient) Read(b []byte) (n int, err error) {

	if !c.isConnected {
		return 0, errors.New("persistentclient: read to closed conn")
	}
	c.conn.SetReadDeadline(time.Now().Add(time.Second * TimeoutSeconds))
	return c.conn.Read(b)
}
<commit_msg>deleted persistentclient.go<commit_after><|endoftext|>


<commit_before>//  Copyright (c) 2014 Couchbase, Inc.
//  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 main

import (
	"fmt"

	log "github.com/couchbaselabs/clog"
)

// A janitor maintains feeds, creating and deleting as necessary.
func (mgr *Manager) JanitorLoop() {
	for reason := range mgr.janitorCh {
		log.Printf("janitor awakes, reason: %s", reason)

		if mgr.cfg == nil { // Can occur during testing.
			log.Printf("janitor skipped due to nil cfg")
			continue
		}

		planPIndexes, _, err := CfgGetPlanPIndexes(mgr.cfg)
		if err != nil {
			log.Printf("janitor skipped due to CfgGetPlanPIndexes err: %v", err)
			continue
		}
		if planPIndexes == nil {
			log.Printf("janitor skipped due to nil planPIndexes")
			continue
		}

		currFeeds, currPIndexes := mgr.CurrentMaps()

		addPlanPIndexes, removePIndexes :=
			CalcPIndexesDelta(currPIndexes, planPIndexes)
		log.Printf("janitor pindexes add: %v, remove: %v",
			addPlanPIndexes, removePIndexes)

		// Create pindexes that we're missing.
		for _, addPlanPIndex := range addPlanPIndexes {
			mgr.StartPIndex(addPlanPIndex)
		}

		// Teardown pindexes that need to be removed.
		for _, removePIndex := range removePIndexes {
			mgr.StopPIndex(removePIndex)
		}

		addFeeds, removeFeeds :=
			CalcFeedsDelta(currFeeds, currPIndexes)
		log.Printf("janitor feeds add: %v, remove: %v",
			addFeeds, removeFeeds)

		// Create feeds that we're missing.
		for _, targetPindexes := range addFeeds {
			mgr.StartFeed(targetPindexes)
		}

		// Teardown feeds that need to be removed.
		for _, removeFeed := range removeFeeds {
			mgr.StopFeed(removeFeed)
		}
	}
}

// Functionally determine the delta of which pindexes need creation
// and which should be shut down.
func CalcPIndexesDelta(currPIndexes map[string]*PIndex,
	wantedPlanPIndexes *PlanPIndexes) (
	addPlanPIndexes []*PlanPIndex,
	removePIndexes []*PIndex) {
	// Allocate our return arrays.
	addPlanPIndexes = make([]*PlanPIndex, 0)
	removePIndexes = make([]*PIndex, 0)

	// Just for fast transient lookups.
	mapWantedPlanPIndex := make(map[string]*PlanPIndex)
	mapRemovePIndex := make(map[string]*PIndex)

	// For each wanted plan pindex, if a pindex does not exist or is
	// different, then schedule to add.
	for _, wantedPlanPIndex := range wantedPlanPIndexes.PlanPIndexes {
		mapWantedPlanPIndex[wantedPlanPIndex.Name] = wantedPlanPIndex

		currPIndex, exists := currPIndexes[wantedPlanPIndex.Name]
		if !exists {
			addPlanPIndexes = append(addPlanPIndexes, wantedPlanPIndex)
		} else if PIndexMatchesPlan(currPIndex, wantedPlanPIndex) == false {
			addPlanPIndexes = append(addPlanPIndexes, wantedPlanPIndex)
			removePIndexes = append(removePIndexes, currPIndex)
			mapRemovePIndex[currPIndex.Name] = currPIndex
		}
	}

	// For each existing pindex, if not part of wanted plan pindex,
	// then schedule for removal.
	for _, currPIndex := range currPIndexes {
		if _, exists := mapWantedPlanPIndex[currPIndex.Name]; !exists {
			if _, exists = mapRemovePIndex[currPIndex.Name]; !exists {
				removePIndexes = append(removePIndexes, currPIndex)
				mapRemovePIndex[currPIndex.Name] = currPIndex
			}
		}
	}

	return addPlanPIndexes, removePIndexes
}

// Functionally determine the delta of which feeds need creation and
// which should be shut down.
func CalcFeedsDelta(currFeeds map[string]Feed, pindexes map[string]*PIndex) (
	addFeeds [][]*PIndex, removeFeeds []Feed) {
	addFeeds = make([][]*PIndex, 0)
	removeFeeds = make([]Feed, 0)

	for _, pindex := range pindexes {
		addFeedName := FeedName("default", pindex.Name, "")
		if _, ok := currFeeds[addFeedName]; !ok {
			addFeeds = append(addFeeds, []*PIndex{pindex})
		}
	}

	return addFeeds, removeFeeds
}

// --------------------------------------------------------

func (mgr *Manager) StartPIndex(planPIndex *PlanPIndex) error {
	// TODO.
	return nil
}

func (mgr *Manager) StopPIndex(pindex *PIndex) error {
	// TODO.
	return nil
}

// --------------------------------------------------------

func (mgr *Manager) StartFeed(pindexes []*PIndex) error {
	// TODO: Need to create a fan-out feed.
	for _, pindex := range pindexes {
		// TODO: Need bucket UUID.
		err := mgr.StartSimpleFeed(pindex) // TODO: err handling.
		if err != nil {
			log.Printf("error: could not start feed for pindex: %s, err: %v",
				pindex.Name, err)
		}
	}
	return nil
}

func (mgr *Manager) StopFeed(feed Feed) error {
	// TODO.
	return nil
}

// --------------------------------------------------------

func (mgr *Manager) StartSimpleFeed(pindex *PIndex) error {
	indexName := pindex.Name // TODO: bad assumption of 1-to-1 pindex.name to indexName

	bucketName := indexName // TODO: read bucketName out of bleve storage.
	bucketUUID := ""        // TODO: read bucketUUID & vbucket list from bleve storage.
	feed, err := NewTAPFeed(mgr.server, "default", bucketName, bucketUUID,
		pindex.Stream)
	if err != nil {
		return fmt.Errorf("error: could not prepare TAP stream to server: %s,"+
			" bucketName: %s, indexName: %s, err: %v",
			mgr.server, bucketName, indexName, err)
		// TODO: need a way to collect these errors so REST api
		// can show them to user ("hey, perhaps you deleted a bucket
		// and should delete these related full-text indexes?
		// or the couchbase cluster is just down.");
		// perhaps as specialized clog writer?
		// TODO: cleanup on error?
	}

	if err = feed.Start(); err != nil {
		// TODO: need way to track dead cows (non-beef)
		// TODO: cleanup?
		return fmt.Errorf("error: could not start feed, server: %s, err: %v",
			mgr.server, err)
	}

	if err = mgr.RegisterFeed(feed); err != nil {
		// TODO: cleanup?
		return err
	}

	return nil
}

// --------------------------------------------------------

func PIndexMatchesPlan(pindex *PIndex, planPIndex *PlanPIndex) bool {
	same :=
		pindex.Name == planPIndex.Name &&
			pindex.UUID == planPIndex.UUID &&
			pindex.IndexName == planPIndex.IndexName &&
			pindex.IndexUUID == planPIndex.IndexUUID &&
			pindex.IndexMapping == planPIndex.IndexMapping &&
			pindex.SourceType == planPIndex.SourceType &&
			pindex.SourceName == planPIndex.SourceName &&
			pindex.SourceUUID == planPIndex.SourceUUID &&
			pindex.SourcePartitions == planPIndex.SourcePartitions
	return same
}
<commit_msg>CalcPIndexesDelta() now only focuses on the subset of plan for the local node<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  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 main

import (
	"fmt"

	log "github.com/couchbaselabs/clog"
)

// A janitor maintains feeds, creating and deleting as necessary.
func (mgr *Manager) JanitorLoop() {
	for reason := range mgr.janitorCh {
		log.Printf("janitor awakes, reason: %s", reason)

		if mgr.cfg == nil { // Can occur during testing.
			log.Printf("janitor skipped due to nil cfg")
			continue
		}

		planPIndexes, _, err := CfgGetPlanPIndexes(mgr.cfg)
		if err != nil {
			log.Printf("janitor skipped due to CfgGetPlanPIndexes err: %v", err)
			continue
		}
		if planPIndexes == nil {
			log.Printf("janitor skipped due to nil planPIndexes")
			continue
		}

		currFeeds, currPIndexes := mgr.CurrentMaps()

		addPlanPIndexes, removePIndexes :=
			CalcPIndexesDelta(mgr.uuid, currPIndexes, planPIndexes)
		log.Printf("janitor pindexes add: %v, remove: %v",
			addPlanPIndexes, removePIndexes)

		// Create pindexes that we're missing.
		for _, addPlanPIndex := range addPlanPIndexes {
			mgr.StartPIndex(addPlanPIndex)
		}

		// Teardown pindexes that need to be removed.
		for _, removePIndex := range removePIndexes {
			mgr.StopPIndex(removePIndex)
		}

		currFeeds, currPIndexes = mgr.CurrentMaps()

		addFeeds, removeFeeds :=
			CalcFeedsDelta(currFeeds, currPIndexes)
		log.Printf("janitor feeds add: %v, remove: %v",
			addFeeds, removeFeeds)

		// Create feeds that we're missing.
		for _, targetPindexes := range addFeeds {
			mgr.StartFeed(targetPindexes)
		}

		// Teardown feeds that need to be removed.
		for _, removeFeed := range removeFeeds {
			mgr.StopFeed(removeFeed)
		}
	}
}

// Functionally determine the delta of which pindexes need creation
// and which should be shut down on our local node (mgrUUID).
func CalcPIndexesDelta(mgrUUID string,
	currPIndexes map[string]*PIndex,
	wantedPlanPIndexes *PlanPIndexes) (
	addPlanPIndexes []*PlanPIndex,
	removePIndexes []*PIndex) {
	// Allocate our return arrays.
	addPlanPIndexes = make([]*PlanPIndex, 0)
	removePIndexes = make([]*PIndex, 0)

	// Just for fast transient lookups.
	mapWantedPlanPIndex := make(map[string]*PlanPIndex)
	mapRemovePIndex := make(map[string]*PIndex)

	// For each wanted plan pindex, if a pindex does not exist or is
	// different, then schedule to add.
	for _, wantedPlanPIndex := range wantedPlanPIndexes.PlanPIndexes {
	nodeUUIDs:
		for _, nodeUUID := range wantedPlanPIndex.NodeUUIDs {
			if nodeUUID == mgrUUID {
				mapWantedPlanPIndex[wantedPlanPIndex.Name] = wantedPlanPIndex

				currPIndex, exists := currPIndexes[wantedPlanPIndex.Name]
				if !exists {
					addPlanPIndexes = append(addPlanPIndexes, wantedPlanPIndex)
				} else if PIndexMatchesPlan(currPIndex, wantedPlanPIndex) == false {
					addPlanPIndexes = append(addPlanPIndexes, wantedPlanPIndex)
					removePIndexes = append(removePIndexes, currPIndex)
					mapRemovePIndex[currPIndex.Name] = currPIndex
				}

				break nodeUUIDs
			}
		}
	}

	// For each existing pindex, if not part of wanted plan pindex,
	// then schedule for removal.
	for _, currPIndex := range currPIndexes {
		if _, exists := mapWantedPlanPIndex[currPIndex.Name]; !exists {
			if _, exists = mapRemovePIndex[currPIndex.Name]; !exists {
				removePIndexes = append(removePIndexes, currPIndex)
				mapRemovePIndex[currPIndex.Name] = currPIndex
			}
		}
	}

	return addPlanPIndexes, removePIndexes
}

// Functionally determine the delta of which feeds need creation and
// which should be shut down.
func CalcFeedsDelta(currFeeds map[string]Feed, pindexes map[string]*PIndex) (
	addFeeds [][]*PIndex, removeFeeds []Feed) {
	addFeeds = make([][]*PIndex, 0)
	removeFeeds = make([]Feed, 0)

	for _, pindex := range pindexes {
		addFeedName := FeedName("default", pindex.SourceName, "")
		if _, ok := currFeeds[addFeedName]; !ok {
			addFeeds = append(addFeeds, []*PIndex{pindex})
		}
	}

	return addFeeds, removeFeeds
}

// --------------------------------------------------------

func (mgr *Manager) StartPIndex(planPIndex *PlanPIndex) error {
	// TODO.
	return nil
}

func (mgr *Manager) StopPIndex(pindex *PIndex) error {
	// TODO.
	return nil
}

// --------------------------------------------------------

func (mgr *Manager) StartFeed(pindexes []*PIndex) error {
	// TODO: Need to create a fan-out feed.
	for _, pindex := range pindexes {
		// TODO: Need bucket UUID.
		err := mgr.StartSimpleFeed(pindex) // TODO: err handling.
		if err != nil {
			log.Printf("error: could not start feed for pindex: %s, err: %v",
				pindex.Name, err)
		}
	}
	return nil
}

func (mgr *Manager) StopFeed(feed Feed) error {
	// TODO.
	return nil
}

// --------------------------------------------------------

func (mgr *Manager) StartSimpleFeed(pindex *PIndex) error {
	indexName := pindex.Name // TODO: bad assumption of 1-to-1 pindex.name to indexName

	bucketName := indexName // TODO: read bucketName out of bleve storage.
	bucketUUID := ""        // TODO: read bucketUUID & vbucket list from bleve storage.
	feed, err := NewTAPFeed(mgr.server, "default", bucketName, bucketUUID,
		pindex.Stream)
	if err != nil {
		return fmt.Errorf("error: could not prepare TAP stream to server: %s,"+
			" bucketName: %s, indexName: %s, err: %v",
			mgr.server, bucketName, indexName, err)
		// TODO: need a way to collect these errors so REST api
		// can show them to user ("hey, perhaps you deleted a bucket
		// and should delete these related full-text indexes?
		// or the couchbase cluster is just down.");
		// perhaps as specialized clog writer?
		// TODO: cleanup on error?
	}

	if err = feed.Start(); err != nil {
		// TODO: need way to track dead cows (non-beef)
		// TODO: cleanup?
		return fmt.Errorf("error: could not start feed, server: %s, err: %v",
			mgr.server, err)
	}

	if err = mgr.RegisterFeed(feed); err != nil {
		// TODO: cleanup?
		return err
	}

	return nil
}

// --------------------------------------------------------

func PIndexMatchesPlan(pindex *PIndex, planPIndex *PlanPIndex) bool {
	same :=
		pindex.Name == planPIndex.Name &&
			pindex.UUID == planPIndex.UUID &&
			pindex.IndexName == planPIndex.IndexName &&
			pindex.IndexUUID == planPIndex.IndexUUID &&
			pindex.IndexMapping == planPIndex.IndexMapping &&
			pindex.SourceType == planPIndex.SourceType &&
			pindex.SourceName == planPIndex.SourceName &&
			pindex.SourceUUID == planPIndex.SourceUUID &&
			pindex.SourcePartitions == planPIndex.SourcePartitions
	return same
}
<|endoftext|>


<commit_before>// Package scanner provides a recursive file scanner that is useful for
// efficiently processing relatively static datasets.
package scanner

import (
	"io/ioutil"
	"os"
	"strings"
	"sync"
)

// Scanner recursively scans a directory for files.
type Scanner struct {

	// Buffer is length of the files channel buffer.
	Buffer int

	// errs is the channel that errors are sent to.
	errs chan error

	// files is the challen that discovered files are sent to.
	files chan string

	// path is the directory being scanned.
	path string

	// separator is the OS's path separator represented as a string. We
	// store this value in a struct so we don;t have to repeatedly convert
	// the rune to a string. This is a micro-optimization, but an
	// optimization none the less.
	separator string
}

// New returns a new Scanner instance.
func New(path string) *Scanner {
	return &Scanner{
		Buffer:    1,
		errs:      make(chan error),
		path:      path,
		separator: string(os.PathSeparator),
	}
}

// Scan recursively scans the directory and sends the files and errors to
// the passed Handler in goroutines.
func (s *Scanner) Scan(h Handler) {
	var wg sync.WaitGroup
	s.files = make(chan string, s.Buffer)

	go func() {
		s.scan(s.path)
		close(s.files)
		close(s.errs)
	}()

	if h == nil {
		return
	}

	wg.Add(2)

	go func() {
		defer wg.Done()
		for err := range s.errs {
			h.HandleError(err)
		}
	}()

	go func() {
		defer wg.Done()
		for f := range s.files {
			h.Handle(f)
		}
	}()

	wg.Wait()
}

// scan recursively scans path and sends the discovered files and errors to
// the built-in channels.
func (s *Scanner) scan(path string) {
	files, err := ioutil.ReadDir(path)
	if err != nil {
		s.errs <- err
		return
	}

	basedir := strings.TrimRight(path, s.separator)
	for _, f := range files {
		file := basedir + s.separator + f.Name()
		if f.IsDir() {
			s.scan(file)
		} else {
			s.files <- file
		}
	}
}

// Handler is an interface for handlers that process scanned files.
type Handler interface {
	Handle(string)
	HandleError(error)
}

// MemoryHandler returns a new memoryHandler instance.
func MemoryHandler() *memoryHandler {
	return &memoryHandler{}
}

// memoryHandler stores the scanned files and errors in slices.
type memoryHandler struct {
	Files  []string
	Errors []error
}

// Handle stores file in a struct.
func (h *memoryHandler) Handle(file string) {
	h.Files = append(h.Files, file)
}

// HandleError stores err in a struct.
func (h *memoryHandler) HandleError(err error) {
	h.Errors = append(h.Errors, err)
}
<commit_msg>docs update<commit_after>// Package scanner provides a recursive file scanner that is useful for
// efficiently processing relatively static datasets.
package scanner

import (
	"io/ioutil"
	"os"
	"strings"
	"sync"
)

// Scanner recursively scans a directory for files.
type Scanner struct {

	// errs is the channel that errors are sent to.
	errs chan error

	// files is the challen that discovered files are sent to.
	files chan string

	// path is the directory being scanned.
	path string

	// separator is the OS's path separator represented as a string. We
	// store this value in a struct so we don;t have to repeatedly convert
	// the rune to a string. This is a micro-optimization, but an
	// optimization none the less.
	separator string

	// Buffer is length of the files channel buffer.
	Buffer int
}

// New returns a new Scanner instance.
func New(path string) *Scanner {
	return &Scanner{
		Buffer:    1,
		errs:      make(chan error),
		path:      path,
		separator: string(os.PathSeparator),
	}
}

// Scan recursively scans the directory and sends the files and errors to
// the passed Handler in goroutines.
func (s *Scanner) Scan(h Handler) {
	var wg sync.WaitGroup
	s.files = make(chan string, s.Buffer)

	go func() {
		s.scan(s.path)
		close(s.files)
		close(s.errs)
	}()

	if h == nil {
		return
	}

	wg.Add(2)

	go func() {
		defer wg.Done()
		for err := range s.errs {
			h.HandleError(err)
		}
	}()

	go func() {
		defer wg.Done()
		for f := range s.files {
			h.Handle(f)
		}
	}()

	wg.Wait()
}

// scan recursively scans path and sends the discovered files and errors to
// the built-in channels.
func (s *Scanner) scan(path string) {
	files, err := ioutil.ReadDir(path)
	if err != nil {
		s.errs <- err
		return
	}

	basedir := strings.TrimRight(path, s.separator)
	for _, f := range files {
		file := basedir + s.separator + f.Name()
		if f.IsDir() {
			s.scan(file)
		} else {
			s.files <- file
		}
	}
}

// Handler is an interface for handlers that process scanned files.
type Handler interface {

	// Handle is passed the path to the file that was discovered during the
	// scan operation. Files are passed in a single-threaded manner, so
	// concurrency is the responsibility of the Handler implementation.
	Handle(string)

	// HandleError is passed the errors that occur during scan operations.
	// Like the Handle method, errors are passed in a single-threaded
	// manner.
	HandleError(error)
}

// MemoryHandler returns a new memoryHandler instance, which stores the
// scanned files and error in local slices.
func MemoryHandler() *memoryHandler {
	return &memoryHandler{}
}

// memoryHandler stores the scanned files and errors in slices.
type memoryHandler struct {
	Files  []string
	Errors []error
}

// Handle stores file in a struct.
func (h *memoryHandler) Handle(file string) {
	h.Files = append(h.Files, file)
}

// HandleError stores err in a struct.
func (h *memoryHandler) HandleError(err error) {
	h.Errors = append(h.Errors, err)
}
<|endoftext|>


<commit_before>package mongoproto

import (
	"bytes"
	"fmt"
	"log"

	"gopkg.in/mgo.v2/bson"
)

const (
	_ OpQueryFlags = 1 << iota

	OpQueryTailableCursor  // Tailable means cursor is not closed when the last data is retrieved. Rather, the cursor marks the final object’s position. You can resume using the cursor later, from where it was located, if more data were received. Like any “latent cursor”, the cursor may become invalid at some point (CursorNotFound) – for example if the final object it references were deleted.
	OpQuerySlaveOk         // Allow query of replica slave. Normally these return an error except for namespace “local”.
	OpQueryOplogReplay     // Internal replication use only - driver should not set
	OpQueryNoCursorTimeout // The server normally times out idle cursors after an inactivity period (10 minutes) to prevent excess memory use. Set this option to prevent that.
	OpQueryAwaitData       // Use with TailableCursor. If we are at the end of the data, block for a while rather than returning no data. After a timeout period, we do return as normal.
	OpQueryExhaust         // Stream the data down full blast in multiple “more” packages, on the assumption that the client will fully read all data queried. Faster when you are pulling a lot of data and know you want to pull it all down. Note: the client is not allowed to not read all the data unless it closes the connection.
	OpQueryPartial         // Get partial results from a mongos if some shards are down (instead of throwing an error)
)

type OpQueryFlags int32

// OpQuery is used to query the database for documents in a collection.
// http://docs.mongodb.org/meta-driver/latest/legacy/mongodb-wire-protocol/#op-query
type OpQuery struct {
	Header               MsgHeader
	Flags                OpQueryFlags
	FullCollectionName   string // "dbname.collectionname"
	NumberToSkip         int32  // number of documents to skip
	NumberToReturn       int32  // number of documents to return
	Query                []byte // query object
	ReturnFieldsSelector []byte // Optional. Selector indicating the fields to return
}

func (op *OpQuery) String() string {
	var query interface{}
	bson.Unmarshal(op.Query, &query)
	return fmt.Sprintf("%#v - %v", op, query)
}

func (op *OpQuery) OpCode() OpCode {
	return OpCodeQuery
}

func (op *OpQuery) FromWire(b []byte) {
	if len(b) < 16 {
		return
	}
	op.Flags = OpQueryFlags(getInt32(b, 0))
	op.FullCollectionName = readCString(b[4:])

	b = b[4+len(op.FullCollectionName)+1:]
	op.NumberToSkip = getInt32(b, 0)
	op.NumberToReturn = getInt32(b, 4)

	b = b[8:]
	query, err := ReadDocument(bytes.NewReader(b))
	log.Println(err)
	op.Query = query
	b = b[len(query):]

	returnFields, err := ReadDocument(bytes.NewReader(b))
	if err == nil {
		op.ReturnFieldsSelector = returnFields
	}
}

func (op *OpQuery) ToWire() []byte {
	return nil
}
<commit_msg>Add JSON representation<commit_after>package mongoproto

import (
	"bytes"
	"encoding/json"
	"fmt"

	"github.com/mongodb/mongo-tools/common/bsonutil"
	"gopkg.in/mgo.v2/bson"
)

const (
	_ OpQueryFlags = 1 << iota

	OpQueryTailableCursor  // Tailable means cursor is not closed when the last data is retrieved. Rather, the cursor marks the final object’s position. You can resume using the cursor later, from where it was located, if more data were received. Like any “latent cursor”, the cursor may become invalid at some point (CursorNotFound) – for example if the final object it references were deleted.
	OpQuerySlaveOk         // Allow query of replica slave. Normally these return an error except for namespace “local”.
	OpQueryOplogReplay     // Internal replication use only - driver should not set
	OpQueryNoCursorTimeout // The server normally times out idle cursors after an inactivity period (10 minutes) to prevent excess memory use. Set this option to prevent that.
	OpQueryAwaitData       // Use with TailableCursor. If we are at the end of the data, block for a while rather than returning no data. After a timeout period, we do return as normal.
	OpQueryExhaust         // Stream the data down full blast in multiple “more” packages, on the assumption that the client will fully read all data queried. Faster when you are pulling a lot of data and know you want to pull it all down. Note: the client is not allowed to not read all the data unless it closes the connection.
	OpQueryPartial         // Get partial results from a mongos if some shards are down (instead of throwing an error)
)

type OpQueryFlags int32

// OpQuery is used to query the database for documents in a collection.
// http://docs.mongodb.org/meta-driver/latest/legacy/mongodb-wire-protocol/#op-query
type OpQuery struct {
	Header               MsgHeader
	Flags                OpQueryFlags
	FullCollectionName   string // "dbname.collectionname"
	NumberToSkip         int32  // number of documents to skip
	NumberToReturn       int32  // number of documents to return
	Query                []byte // query object
	ReturnFieldsSelector []byte // Optional. Selector indicating the fields to return
}

func (op *OpQuery) String() string {
	var query interface{}
	bson.Unmarshal(op.Query, &query)
	queryAsJSON, err := bsonutil.ConvertBSONValueToJSON(query)
	if err != nil {
		return fmt.Sprintf("%#v - %v", op, err)
	}
	asJSON, err := json.Marshal(queryAsJSON)
	if err != nil {
		return fmt.Sprintf("%#v - %v", op, err)
	}
	return fmt.Sprintf("OpQuery %v %v", op.FullCollectionName, string(asJSON))
}

func (op *OpQuery) OpCode() OpCode {
	return OpCodeQuery
}

func (op *OpQuery) FromWire(b []byte) {
	if len(b) < 16 {
		return
	}
	op.Flags = OpQueryFlags(getInt32(b, 0))
	op.FullCollectionName = readCString(b[4:])

	b = b[4+len(op.FullCollectionName)+1:]
	op.NumberToSkip = getInt32(b, 0)
	op.NumberToReturn = getInt32(b, 4)

	b = b[8:]
	op.Query, _ = ReadDocument(bytes.NewReader(b))
	b = b[len(op.Query):]

	if len(b) < 4 {
		return
	}
	returnFields, err := ReadDocument(bytes.NewReader(b))
	if err == nil {
		op.ReturnFieldsSelector = returnFields
	}
}

func (op *OpQuery) ToWire() []byte {
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"strings"

	trans "github.com/ayiga/go-kit-middlewarer/examples/stringsvc/transport/http"
)

// StringService represents an object that will implement the StringService
// interface
type StringService struct{}

// Uppercase implements StringService
func (StringService) Uppercase(str string) (string, error) {
	return strings.ToUpper(str), nil
}

// Count implements StringService
func (StringService) Count(str string) int {
	return len(str)
}

func main() {
	var svc StringService

	trans.HTTPServersForEndpoints(svc)
	http.ListenAndServe(":9000", nil)
}
<commit_msg>add logging to the default example<commit_after>package main

import (
	"net/http"
	"os"
	"strings"

	"github.com/go-kit/kit/log"

	"github.com/ayiga/go-kit-middlewarer/examples/stringsvc"
	"github.com/ayiga/go-kit-middlewarer/examples/stringsvc/logging"
	trans "github.com/ayiga/go-kit-middlewarer/examples/stringsvc/transport/http"
)

// StringService represents an object that will implement the StringService
// interface
type StringService struct{}

// Uppercase implements StringService
func (StringService) Uppercase(str string) (string, error) {
	return strings.ToUpper(str), nil
}

// Count implements StringService
func (StringService) Count(str string) int {
	return len(str)
}

func main() {
	var svc stringsvc.StringService = StringService{}
	l := log.NewLogfmtLogger(os.Stderr)
	svc = logging.Middleware(l, svc)(svc)

	trans.HTTPServersForEndpoints(svc)
	http.ListenAndServe(":9000", nil)
}
<|endoftext|>


<commit_before>package crypto

import (
	"bytes"
	"crypto/sha256"
	"io"

	"github.com/lucas-clemente/quic-go/protocol"
	"github.com/lucas-clemente/quic-go/utils"

	"golang.org/x/crypto/hkdf"
)

// DeriveKeysChacha20 derives the client and server keys and creates a matching chacha20poly1305 AEAD instance
// func DeriveKeysChacha20(version protocol.VersionNumber, forwardSecure bool, sharedSecret, nonces []byte, connID protocol.ConnectionID, chlo []byte, scfg []byte, cert []byte, divNonce []byte) (AEAD, error) {
// 	otherKey, myKey, otherIV, myIV, err := deriveKeys(version, forwardSecure, sharedSecret, nonces, connID, chlo, scfg, cert, divNonce, 32)
// 	if err != nil {
// 		return nil, err
// 	}
// 	return NewAEADChacha20Poly1305(otherKey, myKey, otherIV, myIV)
// }

// DeriveKeysAESGCM derives the client and server keys and creates a matching AES-GCM AEAD instance
func DeriveKeysAESGCM(version protocol.VersionNumber, forwardSecure bool, sharedSecret, nonces []byte, connID protocol.ConnectionID, chlo []byte, scfg []byte, cert []byte, divNonce []byte) (AEAD, error) {
	otherKey, myKey, otherIV, myIV, err := deriveKeys(version, forwardSecure, sharedSecret, nonces, connID, chlo, scfg, cert, divNonce, 16)
	if err != nil {
		return nil, err
	}
	return NewAEADAESGCM(otherKey, myKey, otherIV, myIV)
}

func deriveKeys(version protocol.VersionNumber, forwardSecure bool, sharedSecret, nonces []byte, connID protocol.ConnectionID, chlo, scfg, cert, divNonce []byte, keyLen int) ([]byte, []byte, []byte, []byte, error) {
	var info bytes.Buffer
	if forwardSecure {
		info.Write([]byte("QUIC forward secure key expansion\x00"))
	} else {
		info.Write([]byte("QUIC key expansion\x00"))
	}
	utils.WriteUint64(&info, uint64(connID))
	info.Write(chlo)
	info.Write(scfg)
	info.Write(cert)

	r := hkdf.New(sha256.New, sharedSecret, nonces, info.Bytes())

	otherKey := make([]byte, keyLen)
	myKey := make([]byte, keyLen)
	otherIV := make([]byte, 4)
	myIV := make([]byte, 4)

	if _, err := io.ReadFull(r, otherKey); err != nil {
		return nil, nil, nil, nil, err
	}
	if _, err := io.ReadFull(r, myKey); err != nil {
		return nil, nil, nil, nil, err
	}
	if _, err := io.ReadFull(r, otherIV); err != nil {
		return nil, nil, nil, nil, err
	}
	if _, err := io.ReadFull(r, myIV); err != nil {
		return nil, nil, nil, nil, err
	}

	if !forwardSecure && version >= protocol.Version33 {
		if err := diversify(myKey, myIV, divNonce); err != nil {
			return nil, nil, nil, nil, err
		}
	}

	return otherKey, myKey, otherIV, myIV, nil
}

func diversify(key, iv, divNonce []byte) error {
	secret := make([]byte, len(key)+len(iv))
	copy(secret, key)
	copy(secret[len(key):], iv)

	r := hkdf.New(sha256.New, secret, divNonce, []byte("QUIC key diversification"))

	if _, err := io.ReadFull(r, key); err != nil {
		return err
	}
	if _, err := io.ReadFull(r, iv); err != nil {
		return err
	}

	return nil
}
<commit_msg>simplify key derivation<commit_after>package crypto

import (
	"bytes"
	"crypto/sha256"
	"io"

	"github.com/lucas-clemente/quic-go/protocol"
	"github.com/lucas-clemente/quic-go/utils"

	"golang.org/x/crypto/hkdf"
)

// DeriveKeysChacha20 derives the client and server keys and creates a matching chacha20poly1305 AEAD instance
// func DeriveKeysChacha20(version protocol.VersionNumber, forwardSecure bool, sharedSecret, nonces []byte, connID protocol.ConnectionID, chlo []byte, scfg []byte, cert []byte, divNonce []byte) (AEAD, error) {
// 	otherKey, myKey, otherIV, myIV, err := deriveKeys(version, forwardSecure, sharedSecret, nonces, connID, chlo, scfg, cert, divNonce, 32)
// 	if err != nil {
// 		return nil, err
// 	}
// 	return NewAEADChacha20Poly1305(otherKey, myKey, otherIV, myIV)
// }

// DeriveKeysAESGCM derives the client and server keys and creates a matching AES-GCM AEAD instance
func DeriveKeysAESGCM(version protocol.VersionNumber, forwardSecure bool, sharedSecret, nonces []byte, connID protocol.ConnectionID, chlo []byte, scfg []byte, cert []byte, divNonce []byte) (AEAD, error) {
	otherKey, myKey, otherIV, myIV, err := deriveKeys(version, forwardSecure, sharedSecret, nonces, connID, chlo, scfg, cert, divNonce, 16)
	if err != nil {
		return nil, err
	}
	return NewAEADAESGCM(otherKey, myKey, otherIV, myIV)
}

func deriveKeys(version protocol.VersionNumber, forwardSecure bool, sharedSecret, nonces []byte, connID protocol.ConnectionID, chlo, scfg, cert, divNonce []byte, keyLen int) ([]byte, []byte, []byte, []byte, error) {
	var info bytes.Buffer
	if forwardSecure {
		info.Write([]byte("QUIC forward secure key expansion\x00"))
	} else {
		info.Write([]byte("QUIC key expansion\x00"))
	}
	utils.WriteUint64(&info, uint64(connID))
	info.Write(chlo)
	info.Write(scfg)
	info.Write(cert)

	r := hkdf.New(sha256.New, sharedSecret, nonces, info.Bytes())

	s := make([]byte, 2*keyLen+2*4)
	if _, err := io.ReadFull(r, s); err != nil {
		return nil, nil, nil, nil, err
	}
	otherKey := s[:keyLen]
	myKey := s[keyLen : 2*keyLen]
	otherIV := s[2*keyLen : 2*keyLen+4]
	myIV := s[2*keyLen+4:]

	if !forwardSecure && version >= protocol.Version33 {
		if err := diversify(myKey, myIV, divNonce); err != nil {
			return nil, nil, nil, nil, err
		}
	}

	return otherKey, myKey, otherIV, myIV, nil
}

func diversify(key, iv, divNonce []byte) error {
	secret := make([]byte, len(key)+len(iv))
	copy(secret, key)
	copy(secret[len(key):], iv)

	r := hkdf.New(sha256.New, secret, divNonce, []byte("QUIC key diversification"))

	if _, err := io.ReadFull(r, key); err != nil {
		return err
	}
	if _, err := io.ReadFull(r, iv); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before><commit_msg>fix (multipart): remove multipart form data - #401<commit_after><|endoftext|>


<commit_before>package ssh

const (
	// This is a constant representing a script to install and uninstall public
	// key in remote hosts.
	DefaultPublicKeyInstallScript = `
#!/bin/bash
#
# This is a default script which installs or uninstalls an RSA public key to/from
# authoried_keys file in a typical linux machine.
#
# If the platform differs or if the binaries used in this script are not available
# in targer machine, use the 'install_script' parameter with 'roles/' endpoint to
# register a custom script (applicable for Dynamic type only).
#
# Vault server runs this script on the target machine with the following params:
#
# $1:INSTALL_OPTION: "install" or "uninstall"
#
# $2:PUBLIC_KEY_FILE: File name containing public key to be installed. Vault server
# uses UUID as name to avoid collisions with public keys generated for other requests.
#
# $3:AUTH_KEYS_FILE: Absolute path of the authorized_keys file.
# Currently, vault uses /home/<username>/.ssh/authorized_keys as the path.
#
# [Note: This script will be run by Vault using the registered admin username.
# Notice that some commands below are run as 'sudo'. For graceful execution of
# this script there should not be any password prompts. So, disable password
# prompt for the admin username registered with Vault.

set -e

# Storing arguments into variables, to increase readability of the script.
INSTALL_OPTION=$1
PUBLIC_KEY_FILE=$2
AUTH_KEYS_FILE=$3

# Delete the public key file and the temporary file
function cleanup
{
	rm -f "$PUBLIC_KEY_FILE" temp_$PUBLIC_KEY_FILE
}

# 'cleanup' will be called if the script ends or if any command fails.
trap cleanup EXIT

# Return if the option is anything other than 'install' or 'uninstall'.
if [ "$INSTALL_OPTION" != "install" ] && [ "$INSTALL_OPTION" != "uninstall" ]; then
	exit 1
fi

# Create the .ssh directory and authorized_keys file if it does not exist
SSH_DIR=$(dirname $AUTH_KEYS_FILE)
sudo mkdir -p "$SSH_DIR"
sudo touch "$AUTH_KEYS_FILE"

# Remove the key from authorized_keys file if it is already present.
# This step is common for both install and uninstall.
grep -vFf "$PUBLIC_KEY_FILE" "$AUTH_KEYS_FILE" > temp_$PUBLIC_KEY_FILE || true
cat temp_$PUBLIC_KEY_FILE | sudo tee "$AUTH_KEYS_FILE"

# Append the new public key to authorized_keys file
if [ "$INSTALL_OPTION" == "install" ]; then
	cat "$PUBLIC_KEY_FILE" | sudo tee --append "$AUTH_KEYS_FILE"
fi
`
)
<commit_msg>Ensure authorized_keys file is readable when uninstalling an ssh key<commit_after>package ssh

const (
	// This is a constant representing a script to install and uninstall public
	// key in remote hosts.
	DefaultPublicKeyInstallScript = `
#!/bin/bash
#
# This is a default script which installs or uninstalls an RSA public key to/from
# authoried_keys file in a typical linux machine.
#
# If the platform differs or if the binaries used in this script are not available
# in targer machine, use the 'install_script' parameter with 'roles/' endpoint to
# register a custom script (applicable for Dynamic type only).
#
# Vault server runs this script on the target machine with the following params:
#
# $1:INSTALL_OPTION: "install" or "uninstall"
#
# $2:PUBLIC_KEY_FILE: File name containing public key to be installed. Vault server
# uses UUID as name to avoid collisions with public keys generated for other requests.
#
# $3:AUTH_KEYS_FILE: Absolute path of the authorized_keys file.
# Currently, vault uses /home/<username>/.ssh/authorized_keys as the path.
#
# [Note: This script will be run by Vault using the registered admin username.
# Notice that some commands below are run as 'sudo'. For graceful execution of
# this script there should not be any password prompts. So, disable password
# prompt for the admin username registered with Vault.

set -e

# Storing arguments into variables, to increase readability of the script.
INSTALL_OPTION=$1
PUBLIC_KEY_FILE=$2
AUTH_KEYS_FILE=$3

# Delete the public key file and the temporary file
function cleanup
{
	rm -f "$PUBLIC_KEY_FILE" temp_$PUBLIC_KEY_FILE
}

# 'cleanup' will be called if the script ends or if any command fails.
trap cleanup EXIT

# Return if the option is anything other than 'install' or 'uninstall'.
if [ "$INSTALL_OPTION" != "install" ] && [ "$INSTALL_OPTION" != "uninstall" ]; then
	exit 1
fi

# Create the .ssh directory and authorized_keys file if it does not exist
SSH_DIR=$(dirname $AUTH_KEYS_FILE)
sudo mkdir -p "$SSH_DIR"
sudo touch "$AUTH_KEYS_FILE"

# Remove the key from authorized_keys file if it is already present.
# This step is common for both install and uninstall.  Note that grep's
# return code is ignored, thus if grep fails all keys will be removed
# rather than none and it fails secure
sudo grep -vFf "$PUBLIC_KEY_FILE" "$AUTH_KEYS_FILE" > temp_$PUBLIC_KEY_FILE || true
cat temp_$PUBLIC_KEY_FILE | sudo tee "$AUTH_KEYS_FILE"

# Append the new public key to authorized_keys file
if [ "$INSTALL_OPTION" == "install" ]; then
	cat "$PUBLIC_KEY_FILE" | sudo tee --append "$AUTH_KEYS_FILE"
fi
`
)
<|endoftext|>


<commit_before>// Copyright (c) 2014 The SkyDNS Authors. All rights reserved.
// Copyright (c) 2015 Jan Broer
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package server

import (
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/miekg/dns"
)

// ServeDNSForward forwards a request to a nameservers and returns the response.
func (s *server) ServeDNSForward(w dns.ResponseWriter, req *dns.Msg) *dns.Msg {
	if s.config.NoRec || len(s.config.Nameservers) == 0 {
		m := new(dns.Msg)
		m.SetReply(req)
		m.SetRcode(req, dns.RcodeServerFailure)
		m.Authoritative = false
		m.RecursionAvailable = false
		if len(s.config.Nameservers) == 0 {
			log.Debug("Can not forward query, no nameservers defined")
			m.RecursionAvailable = true
		} else {
			m.RecursionAvailable = false
		}

		w.WriteMsg(m)
		return m
	}

	name := req.Question[0].Name

	if dns.CountLabel(name) < 2 || dns.CountLabel(name) < s.config.Ndots {
		// Don't process single-label queries when searching is not enabled
		if !s.config.AppendDomain || len(s.config.SearchDomains) == 0 {
			log.Debugf("Can not forward query, name too short: `%s'", name)
			m := new(dns.Msg)
			m.SetReply(req)
			m.SetRcode(req, dns.RcodeServerFailure)
			m.Authoritative = false
			m.RecursionAvailable = true
			w.WriteMsg(m)
			return m
		}
	}

	StatsForwardCount.Inc(1)

	var (
		r       *dns.Msg
		err     error
		nsIndex int // nameserver list index
		sdIndex int // search list index
		sdName  string // QNAME with search path
		sdCname = new(dns.CNAME) // CNAME record returned when query resolved by searching
	)

	tcp := isTCP(w)
	reqCopy := req.Copy()
	canSearch := false
	doingSearch := false

	if s.config.AppendDomain && len(s.config.SearchDomains) > 0 {
		canSearch = true
	}

Redo:
	if dns.CountLabel(name) < 2 {
		// always qualify single-label names
		if !doingSearch && canSearch {
			doingSearch = true
			sdIndex = 0
		}
	}
	if doingSearch {
		sdName = strings.ToLower(appendDomain(name, s.config.SearchDomains[sdIndex]))
		sdCname.Hdr = dns.RR_Header{Name: name, Rrtype: dns.TypeCNAME, Class: dns.ClassINET, Ttl: 360}
		sdCname.Target = sdName
		req.Question[0] = dns.Question{sdName, req.Question[0].Qtype, req.Question[0].Qclass}
	}

	switch tcp {
	case false:
		r, _, err = s.dnsUDPclient.Exchange(req, s.config.Nameservers[nsIndex])
	case true:
		r, _, err = s.dnsTCPclient.Exchange(req, s.config.Nameservers[nsIndex])
	}
	if err == nil {
		if canSearch {
			// replicate libc's getaddrinfo.c search logic
			switch {
			case r.Rcode == dns.RcodeSuccess && len(r.Answer) == 0: // NODATA
				fallthrough
			case r.Rcode == dns.RcodeNameError: // NXDOMAIN
				fallthrough
			case r.Rcode == dns.RcodeServerFailure: // SERVFAIL
				if doingSearch && (sdIndex + 1) < len(s.config.SearchDomains) {
					// continue searching
					sdIndex++
					goto Redo
				}
				if !doingSearch {
					// start searching
					doingSearch = true
					sdIndex = 0
					goto Redo
				}
			}
		}

		if r.Rcode == dns.RcodeServerFailure || r.Rcode == dns.RcodeRefused {
			// continue with next available nameserver
			if (nsIndex + 1) < len(s.config.Nameservers) {
				nsIndex++
				doingSearch = false
				goto Redo
			}
		}

		// We are done querying. Process the reply to return to the client.

		if doingSearch {
			// Insert cname record pointing name to name.searchdomain
			if len(r.Answer) > 0 {
				answers := []dns.RR{sdCname}
				for _, rr := range r.Answer {
					answers = append(answers, rr)
				}
				r.Answer = answers
			}
			// Restore original question
			r.Question[0] = reqCopy.Question[0]
		}

		r.Compress = true
		r.Id = req.Id
		w.WriteMsg(r)
		return r
	} else {
		log.Debugf("Error querying nameserver %s: %q", s.config.Nameservers[nsIndex], err)
		// Got an error, this usually means the server did not respond
		// Continue with next available nameserver
		if (nsIndex + 1) < len(s.config.Nameservers) {
			nsIndex++
			doingSearch = false
			goto Redo
		}
	}

	// If we got here it means forwarding failed
	log.Errorf("Failure forwarding request %q", err)
	m := new(dns.Msg)
	m.SetReply(reqCopy)
	m.SetRcode(reqCopy, dns.RcodeServerFailure)
	w.WriteMsg(m)
	return m
}

// ServeDNSReverse is the handler for DNS requests for the reverse zone. If nothing is found
// locally the request is forwarded to the forwarder for resolution.
func (s *server) ServeDNSReverse(w dns.ResponseWriter, req *dns.Msg) *dns.Msg {
	m := new(dns.Msg)
	m.SetReply(req)
	m.Compress = true
	m.Authoritative = false
	m.RecursionAvailable = true
	if records, err := s.PTRRecords(req.Question[0]); err == nil && len(records) > 0 {
		m.Answer = records
		if err := w.WriteMsg(m); err != nil {
			log.Errorf("Failure returning reply %q", err)
		}
		return m
	}
	// Always forward if not found locally.
	return s.ServeDNSForward(w, req)
}
<commit_msg>Don't search when empty response is truncated<commit_after>// Copyright (c) 2014 The SkyDNS Authors. All rights reserved.
// Copyright (c) 2015 Jan Broer
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package server

import (
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/miekg/dns"
)

// ServeDNSForward forwards a request to a nameservers and returns the response.
func (s *server) ServeDNSForward(w dns.ResponseWriter, req *dns.Msg) *dns.Msg {
	if s.config.NoRec || len(s.config.Nameservers) == 0 {
		m := new(dns.Msg)
		m.SetReply(req)
		m.SetRcode(req, dns.RcodeServerFailure)
		m.Authoritative = false
		m.RecursionAvailable = false
		if len(s.config.Nameservers) == 0 {
			log.Debug("Can not forward query, no nameservers defined")
			m.RecursionAvailable = true
		} else {
			m.RecursionAvailable = false
		}

		w.WriteMsg(m)
		return m
	}

	name := req.Question[0].Name

	if dns.CountLabel(name) < 2 || dns.CountLabel(name) < s.config.Ndots {
		// Don't process single-label queries when searching is not enabled
		if !s.config.AppendDomain || len(s.config.SearchDomains) == 0 {
			log.Debugf("Can not forward query, name too short: `%s'", name)
			m := new(dns.Msg)
			m.SetReply(req)
			m.SetRcode(req, dns.RcodeServerFailure)
			m.Authoritative = false
			m.RecursionAvailable = true
			w.WriteMsg(m)
			return m
		}
	}

	StatsForwardCount.Inc(1)

	var (
		r       *dns.Msg
		err     error
		nsIndex int // nameserver list index
		sdIndex int // search list index
		sdName  string // QNAME with search path
		sdCname = new(dns.CNAME) // CNAME record returned when query resolved by searching
	)

	tcp := isTCP(w)
	reqCopy := req.Copy()
	canSearch := false
	doingSearch := false

	if s.config.AppendDomain && len(s.config.SearchDomains) > 0 {
		canSearch = true
	}

Redo:
	if dns.CountLabel(name) < 2 {
		// always qualify single-label names
		if !doingSearch && canSearch {
			doingSearch = true
			sdIndex = 0
		}
	}
	if doingSearch {
		sdName = strings.ToLower(appendDomain(name, s.config.SearchDomains[sdIndex]))
		sdCname.Hdr = dns.RR_Header{Name: name, Rrtype: dns.TypeCNAME, Class: dns.ClassINET, Ttl: 360}
		sdCname.Target = sdName
		req.Question[0] = dns.Question{sdName, req.Question[0].Qtype, req.Question[0].Qclass}
	}

	switch tcp {
	case false:
		r, _, err = s.dnsUDPclient.Exchange(req, s.config.Nameservers[nsIndex])
	case true:
		r, _, err = s.dnsTCPclient.Exchange(req, s.config.Nameservers[nsIndex])
	}
	if err == nil {
		if canSearch {
			// replicate libc's getaddrinfo.c search logic
			switch {
			case r.Rcode == dns.RcodeSuccess && len(r.Answer) == 0 && !r.MsgHdr.Truncated: // NODATA !Truncated
				fallthrough
			case r.Rcode == dns.RcodeNameError: // NXDOMAIN
				fallthrough
			case r.Rcode == dns.RcodeServerFailure: // SERVFAIL
				if doingSearch && (sdIndex + 1) < len(s.config.SearchDomains) {
					// continue searching
					sdIndex++
					goto Redo
				}
				if !doingSearch {
					// start searching
					doingSearch = true
					sdIndex = 0
					goto Redo
				}
			}
		}

		if r.Rcode == dns.RcodeServerFailure || r.Rcode == dns.RcodeRefused {
			// continue with next available nameserver
			if (nsIndex + 1) < len(s.config.Nameservers) {
				nsIndex++
				doingSearch = false
				goto Redo
			}
		}

		// We are done querying. Process the reply to return to the client.

		if doingSearch {
			// Insert cname record pointing name to name.searchdomain
			if len(r.Answer) > 0 {
				answers := []dns.RR{sdCname}
				for _, rr := range r.Answer {
					answers = append(answers, rr)
				}
				r.Answer = answers
			}
			// Restore original question
			r.Question[0] = reqCopy.Question[0]
		}

		r.Compress = true
		r.Id = req.Id
		w.WriteMsg(r)
		return r
	} else {
		log.Debugf("Error querying nameserver %s: %q", s.config.Nameservers[nsIndex], err)
		// Got an error, this usually means the server did not respond
		// Continue with next available nameserver
		if (nsIndex + 1) < len(s.config.Nameservers) {
			nsIndex++
			doingSearch = false
			goto Redo
		}
	}

	// If we got here it means forwarding failed
	log.Errorf("Failure forwarding request %q", err)
	m := new(dns.Msg)
	m.SetReply(reqCopy)
	m.SetRcode(reqCopy, dns.RcodeServerFailure)
	w.WriteMsg(m)
	return m
}

// ServeDNSReverse is the handler for DNS requests for the reverse zone. If nothing is found
// locally the request is forwarded to the forwarder for resolution.
func (s *server) ServeDNSReverse(w dns.ResponseWriter, req *dns.Msg) *dns.Msg {
	m := new(dns.Msg)
	m.SetReply(req)
	m.Compress = true
	m.Authoritative = false
	m.RecursionAvailable = true
	if records, err := s.PTRRecords(req.Question[0]); err == nil && len(records) > 0 {
		m.Answer = records
		if err := w.WriteMsg(m); err != nil {
			log.Errorf("Failure returning reply %q", err)
		}
		return m
	}
	// Always forward if not found locally.
	return s.ServeDNSForward(w, req)
}
<|endoftext|>