<commit_before>package local

import (
	"bufio"
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"os"
	"os/exec"
	"path"
	"strings"
	"text/template"
	"time"

	"github.com/rockbears/log"

	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/hatchery"
)

type localWorkerRunner struct{}

func (localWorkerRunner) NewCmd(ctx context.Context, command string, args ...string) *exec.Cmd {
	var cmd = exec.CommandContext(ctx, command, args...)
	cmd.Env = []string{}
	return cmd
}

type localWorkerLogger struct {
	name string
}

func (l localWorkerLogger) Logf(fmt string, values ...interface{}) {
	fmt = strings.TrimSuffix(fmt, "\n")
	log.Info(context.Background(), "hatchery> local> worker> %s> "+fmt, l.name)
}

func (l localWorkerLogger) Errorf(fmt string, values ...interface{}) {
	fmt = strings.TrimSuffix(fmt, "\n")
	log.Error(context.Background(), "hatchery> local> worker> %s> "+fmt, l.name)
}

func (l localWorkerLogger) Fatalf(fmt string, values ...interface{}) {
	fmt = strings.TrimSuffix(fmt, "\n")
	log.Fatal(context.TODO(), "hatchery> local> worker> %s> "+fmt, l.name)
}

const workerCmdTmpl = "{{.WorkerBinary}} --api={{.API}} --token={{.Token}} --log-level=debug --basedir={{.BaseDir}} --name={{.Name}} --hatchery-name={{.HatcheryName}} --insecure={{.HTTPInsecure}} --graylog-extra-key={{.GraylogExtraKey}} --graylog-extra-value={{.GraylogExtraValue}} --graylog-host={{.GraylogHost}} --graylog-port={{.GraylogPort}} --booked-workflow-job-id={{.WorkflowJobID}}"

// SpawnWorker starts a new worker process
func (h *HatcheryLocal) SpawnWorker(ctx context.Context, spawnArgs hatchery.SpawnArguments) error {
	log.Debug(ctx, "HatcheryLocal.SpawnWorker> %s want to spawn a worker named %s (jobID = %d)", spawnArgs.HatcheryName, spawnArgs.WorkerName, spawnArgs.JobID)

	if spawnArgs.JobID == 0 && !spawnArgs.RegisterOnly {
		return sdk.WithStack(fmt.Errorf("no job ID and no register"))
	}

	// Generate a random string 16 chars length
	bs := make([]byte, 16)
	if _, err := rand.Read(bs); err != nil {
		return err
	}
	rndstr := hex.EncodeToString(bs)[0:16]
	basedir := path.Join(h.Config.Basedir, rndstr)
	// Create the directory
	if err := os.MkdirAll(basedir, os.FileMode(0755)); err != nil {
		return err
	}

	log.Info(ctx, "HatcheryLocal.SpawnWorker> basedir: %s", basedir)

	udataParam := h.GenerateWorkerArgs(ctx, h, spawnArgs)
	udataParam.BaseDir = basedir
	udataParam.WorkerBinary = path.Join(h.BasedirDedicated, h.getWorkerBinaryName())
	udataParam.WorkflowJobID = spawnArgs.JobID

	tmpl, errt := template.New("cmd").Parse(workerCmdTmpl)
	if errt != nil {
		return errt
	}
	var buffer bytes.Buffer
	if errTmpl := tmpl.Execute(&buffer, udataParam); errTmpl != nil {
		return errTmpl
	}

	cmdSplitted := strings.Split(buffer.String(), " -")
	for i := range cmdSplitted[1:] {
		cmdSplitted[i+1] = "-" + strings.Trim(cmdSplitted[i+1], " ")
	}

	// Prefix the command with the directory where the worker binary has been downloaded
	log.Debug(ctx, "Command exec: %v", cmdSplitted)
	var cmd *exec.Cmd
	if spawnArgs.RegisterOnly {
		cmdSplitted[0] = "register"
		cmd = h.LocalWorkerRunner.NewCmd(ctx, cmdSplitted[0], cmdSplitted...)
	} else {
		cmd = h.LocalWorkerRunner.NewCmd(ctx, cmdSplitted[0], cmdSplitted[1:]...)
	}

	cmd.Dir = udataParam.BaseDir

	// Clearenv
	env := os.Environ()
	for _, e := range env {
		if !strings.HasPrefix(e, "CDS") && !strings.HasPrefix(e, "HATCHERY") {
			cmd.Env = append(cmd.Env, e)
		}
	}
	for k, v := range udataParam.InjectEnvVars {
		cmd.Env = append(cmd.Env, k+"="+v)
	}

	// Wait in a goroutine so that when process exits, Wait() update cmd.ProcessState
	go func() {
		log.Debug(ctx, "hatchery> local> starting worker: %s", spawnArgs.WorkerName)
		if err := h.startCmd(spawnArgs.WorkerName, cmd, localWorkerLogger{spawnArgs.WorkerName}); err != nil {
			log.Error(ctx, "hatchery> local> %v", err)
		}
	}()

	return nil
}

type Logger interface {
	Logf(fmt string, values ...interface{})
	Errorf(fmt string, values ...interface{})
	Fatalf(fmt string, values ...interface{})
}

func (h *HatcheryLocal) startCmd(name string, cmd *exec.Cmd, logger Logger) error {
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("Failure due to internal error: unable to capture stdout: %v", err)
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("Failure due to internal error: unable to capture stderr: %v", err)
	}

	stdoutreader := bufio.NewReader(stdout)
	stderrreader := bufio.NewReader(stderr)

	outchan := make(chan bool)
	go func() {
		for {
			line, err := stdoutreader.ReadString('\n')
			if line != "" {
				logger.Logf(line)
			}
			if err != nil {
				stdout.Close()
				close(outchan)
				return
			}
		}
	}()

	errchan := make(chan bool)
	go func() {
		for {
			line, err := stderrreader.ReadString('\n')
			if line != "" {
				logger.Logf(line)
			}
			if err != nil {
				stderr.Close()
				close(errchan)
				return
			}
		}
	}()

	if err := cmd.Start(); err != nil {
		return fmt.Errorf("unable to start command: %v", err)
	}

	h.Lock()
	h.workers[name] = workerCmd{cmd: cmd, created: time.Now()}
	h.Unlock()

	<-outchan
	<-errchan
	if err := cmd.Wait(); err != nil {
		return fmt.Errorf("command failure: %v", err)
	}

	return nil
}
<commit_msg>fix(hatchery:local): unable to run worker cmd: context cancelled (#5958)<commit_after>package local

import (
	"bufio"
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"os"
	"os/exec"
	"path"
	"strings"
	"text/template"
	"time"

	"github.com/rockbears/log"

	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/hatchery"
)

type localWorkerRunner struct{}

func (localWorkerRunner) NewCmd(ctx context.Context, command string, args ...string) *exec.Cmd {
	var cmd = exec.CommandContext(ctx, command, args...)
	cmd.Env = []string{}
	return cmd
}

type localWorkerLogger struct {
	name string
}

func (l localWorkerLogger) Logf(fmt string, values ...interface{}) {
	fmt = strings.TrimSuffix(fmt, "\n")
	log.Info(context.Background(), "hatchery> local> worker> %s> "+fmt, l.name)
}

func (l localWorkerLogger) Errorf(fmt string, values ...interface{}) {
	fmt = strings.TrimSuffix(fmt, "\n")
	log.Error(context.Background(), "hatchery> local> worker> %s> "+fmt, l.name)
}

func (l localWorkerLogger) Fatalf(fmt string, values ...interface{}) {
	fmt = strings.TrimSuffix(fmt, "\n")
	log.Fatal(context.TODO(), "hatchery> local> worker> %s> "+fmt, l.name)
}

const workerCmdTmpl = "{{.WorkerBinary}} --api={{.API}} --token={{.Token}} --log-level=debug --basedir={{.BaseDir}} --name={{.Name}} --hatchery-name={{.HatcheryName}} --insecure={{.HTTPInsecure}} --graylog-extra-key={{.GraylogExtraKey}} --graylog-extra-value={{.GraylogExtraValue}} --graylog-host={{.GraylogHost}} --graylog-port={{.GraylogPort}} --booked-workflow-job-id={{.WorkflowJobID}}"

// SpawnWorker starts a new worker process
func (h *HatcheryLocal) SpawnWorker(ctx context.Context, spawnArgs hatchery.SpawnArguments) error {
	log.Debug(ctx, "HatcheryLocal.SpawnWorker> %s want to spawn a worker named %s (jobID = %d)", spawnArgs.HatcheryName, spawnArgs.WorkerName, spawnArgs.JobID)

	if spawnArgs.JobID == 0 && !spawnArgs.RegisterOnly {
		return sdk.WithStack(fmt.Errorf("no job ID and no register"))
	}

	// Generate a random string 16 chars length
	bs := make([]byte, 16)
	if _, err := rand.Read(bs); err != nil {
		return err
	}
	rndstr := hex.EncodeToString(bs)[0:16]
	basedir := path.Join(h.Config.Basedir, rndstr)
	// Create the directory
	if err := os.MkdirAll(basedir, os.FileMode(0755)); err != nil {
		return err
	}

	log.Info(ctx, "HatcheryLocal.SpawnWorker> basedir: %s", basedir)

	udataParam := h.GenerateWorkerArgs(ctx, h, spawnArgs)
	udataParam.BaseDir = basedir
	udataParam.WorkerBinary = path.Join(h.BasedirDedicated, h.getWorkerBinaryName())
	udataParam.WorkflowJobID = spawnArgs.JobID

	tmpl, errt := template.New("cmd").Parse(workerCmdTmpl)
	if errt != nil {
		return errt
	}
	var buffer bytes.Buffer
	if errTmpl := tmpl.Execute(&buffer, udataParam); errTmpl != nil {
		return errTmpl
	}

	cmdSplitted := strings.Split(buffer.String(), " -")
	for i := range cmdSplitted[1:] {
		cmdSplitted[i+1] = "-" + strings.Trim(cmdSplitted[i+1], " ")
	}

	// Prefix the command with the directory where the worker binary has been downloaded
	log.Debug(ctx, "Command exec: %v", cmdSplitted)
	var cmd *exec.Cmd
	if spawnArgs.RegisterOnly {
		cmdSplitted[0] = "register"
		cmd = h.LocalWorkerRunner.NewCmd(context.Background(), cmdSplitted[0], cmdSplitted...)
	} else {
		cmd = h.LocalWorkerRunner.NewCmd(context.Background(), cmdSplitted[0], cmdSplitted[1:]...)
	}

	cmd.Dir = udataParam.BaseDir

	// Clearenv
	env := os.Environ()
	for _, e := range env {
		if !strings.HasPrefix(e, "CDS") && !strings.HasPrefix(e, "HATCHERY") {
			cmd.Env = append(cmd.Env, e)
		}
	}
	for k, v := range udataParam.InjectEnvVars {
		cmd.Env = append(cmd.Env, k+"="+v)
	}

	// Wait in a goroutine so that when process exits, Wait() update cmd.ProcessState
	go func() {
		log.Debug(ctx, "hatchery> local> starting worker: %s", spawnArgs.WorkerName)
		if err := h.startCmd(spawnArgs.WorkerName, cmd, localWorkerLogger{spawnArgs.WorkerName}); err != nil {
			log.Error(ctx, "hatchery> local> %v", err)
		}
	}()

	return nil
}

type Logger interface {
	Logf(fmt string, values ...interface{})
	Errorf(fmt string, values ...interface{})
	Fatalf(fmt string, values ...interface{})
}

func (h *HatcheryLocal) startCmd(name string, cmd *exec.Cmd, logger Logger) error {
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("Failure due to internal error: unable to capture stdout: %v", err)
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("Failure due to internal error: unable to capture stderr: %v", err)
	}

	stdoutreader := bufio.NewReader(stdout)
	stderrreader := bufio.NewReader(stderr)

	outchan := make(chan bool)
	go func() {
		for {
			line, err := stdoutreader.ReadString('\n')
			if line != "" {
				logger.Logf(line)
			}
			if err != nil {
				stdout.Close()
				close(outchan)
				return
			}
		}
	}()

	errchan := make(chan bool)
	go func() {
		for {
			line, err := stderrreader.ReadString('\n')
			if line != "" {
				logger.Logf(line)
			}
			if err != nil {
				stderr.Close()
				close(errchan)
				return
			}
		}
	}()

	if err := cmd.Start(); err != nil {
		return fmt.Errorf("unable to start command: %v", err)
	}

	h.Lock()
	h.workers[name] = workerCmd{cmd: cmd, created: time.Now()}
	h.Unlock()

	<-outchan
	<-errchan
	if err := cmd.Wait(); err != nil {
		return fmt.Errorf("command failure: %v", err)
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 cmac

import (
	"crypto/aes"
	"crypto/cipher"
	"fmt"
	"github.com/jacobsa/aes/common"
	"hash"
)

type cmacHash struct {
	// An AES cipher configured with the original key.
	ciph cipher.Block

	// Generated sub-keys.
	k1 []byte
	k2 []byte

	// Data that has been seen since the last block was disposed of (i.e. since
	// we finished an iteration of the for loop in RFC 4493's AES-CMAC algorithm
	// and were sure we were going into a new one).
	data []byte

	// The current value of X, as defined in the AES-CMAC algorithm in RFC 4493.
	// Initially this is a 128-bit zero, and it is updated with the current block
	// when we're sure it's not the last one.
	x []byte
}

func (h *cmacHash) Write(p []byte) (n int, err error) {
	// Consume the data.
	n = len(p)
	h.data = append(h.data, p...)

	// Consume any blocks that we're sure aren't the last.
	blocksToProcess := len(h.data) / 16
	if blocksToProcess > 0 && len(h.data)%16 == 0 {
		blocksToProcess--
	}

	for i := 0; i < blocksToProcess; i++ {
		block := h.data[16*i : 16*(i+1)]
		y := common.Xor(h.x, block)
		h.ciph.Encrypt(h.x, y)
	}

	h.data = h.data[16*blocksToProcess:]

	return
}

func (h *cmacHash) Sum(b []byte) []byte {
	dataLen := len(h.data)

	// We should have at most one block left.
	if dataLen > 16 {
		panic(fmt.Sprintf("Unexpected data: %x", h.data))
	}

	// Calculate M_last.
	var mLast []byte
	if dataLen == 16 {
		mLast = common.Xor(h.data, h.k1)
	} else {
		mLast = common.Xor(common.PadBlock(h.data), h.k2)
	}

	y := common.Xor(mLast, h.x)
	result := make([]byte, 16)
	h.ciph.Encrypt(result, y)

	b = append(b, result...)
	return b
}

func (h *cmacHash) Reset() {
	h.data = []byte{}
	h.x = make([]byte, 16)
}

func (h *cmacHash) Size() int {
	return 16
}

func (h *cmacHash) BlockSize() int {
	return 16
}

// New returns an AES-CMAC hash using the supplied key. The key must be 16, 24,
// or 32 bytes long.
func New(key []byte) (hash.Hash, error) {
	if len(key) != 16 {
		return nil, fmt.Errorf("AES-CMAC requires a 16-byte key.")
	}

	// Create a cipher.
	ciph, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("aes.NewCipher: %v", err)
	}

	// Set up the hash object.
	h := &cmacHash{ciph: ciph}
	h.k1, h.k2 = generateSubkeys(ciph)
	h.Reset()

	return h, nil
}
<commit_msg>Updated implementation.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 cmac

import (
	"crypto/aes"
	"crypto/cipher"
	"fmt"
	"github.com/jacobsa/aes/common"
	"hash"
)

type cmacHash struct {
	// An AES cipher configured with the original key.
	ciph cipher.Block

	// Generated sub-keys.
	k1 []byte
	k2 []byte

	// Data that has been seen since the last block was disposed of (i.e. since
	// we finished an iteration of the for loop in RFC 4493's AES-CMAC algorithm
	// and were sure we were going into a new one).
	data []byte

	// The current value of X, as defined in the AES-CMAC algorithm in RFC 4493.
	// Initially this is a 128-bit zero, and it is updated with the current block
	// when we're sure it's not the last one.
	x []byte
}

func (h *cmacHash) Write(p []byte) (n int, err error) {
	// Consume the data.
	n = len(p)
	h.data = append(h.data, p...)

	// Consume any blocks that we're sure aren't the last.
	blocksToProcess := len(h.data) / 16
	if blocksToProcess > 0 && len(h.data)%16 == 0 {
		blocksToProcess--
	}

	for i := 0; i < blocksToProcess; i++ {
		block := h.data[16*i : 16*(i+1)]
		y := common.Xor(h.x, block)
		h.ciph.Encrypt(h.x, y)
	}

	h.data = h.data[16*blocksToProcess:]

	return
}

func (h *cmacHash) Sum(b []byte) []byte {
	dataLen := len(h.data)

	// We should have at most one block left.
	if dataLen > 16 {
		panic(fmt.Sprintf("Unexpected data: %x", h.data))
	}

	// Calculate M_last.
	var mLast []byte
	if dataLen == 16 {
		mLast = common.Xor(h.data, h.k1)
	} else {
		mLast = common.Xor(common.PadBlock(h.data), h.k2)
	}

	y := common.Xor(mLast, h.x)
	result := make([]byte, 16)
	h.ciph.Encrypt(result, y)

	b = append(b, result...)
	return b
}

func (h *cmacHash) Reset() {
	h.data = []byte{}
	h.x = make([]byte, 16)
}

func (h *cmacHash) Size() int {
	return 16
}

func (h *cmacHash) BlockSize() int {
	return h.ciph.BlockSize()
}

// New returns an AES-CMAC hash using the supplied key. The key must be 16, 24,
// or 32 bytes long.
func New(key []byte) (hash.Hash, error) {
	switch len(key) {
	case 16, 24, 32:
	default:
		return nil, fmt.Errorf("AES-CMAC requires a 16-, 24-, or 32-byte key.")
	}

	// Create a cipher.
	ciph, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("aes.NewCipher: %v", err)
	}

	// Set up the hash object.
	h := &cmacHash{ciph: ciph}
	h.k1, h.k2 = generateSubkeys(ciph)
	h.Reset()

	return h, nil
}
<|endoftext|>


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

package main

import (
	"bufio"
	"encoding/csv"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"sync"
	"time"

	"github.com/richardlehane/siegfried"
	"github.com/richardlehane/siegfried/config"
	"github.com/richardlehane/siegfried/pkg/core"

	//_ "net/http/pprof"
)

const (
	CONCURRENT = 4
	PROCS      = -1
)

var (
	update  = flag.Bool("update", false, "update or install the default signature file")
	version = flag.Bool("version", false, "display version information")
	debug   = flag.Bool("debug", false, "scan in debug mode")
	nr      = flag.Bool("nr", false, "prevent automatic directory recursion")
	csvo    = flag.Bool("csv", false, "CSV output format")
	sig     = flag.String("sig", config.SignatureBase(), "set the signature file")
	home    = flag.String("home", config.Home(), "override the default home directory")
	serve   = flag.String("serve", "false", "not yet implemented - coming with v1")
)

var (
	csvWriter  *csv.Writer
	yamlWriter *bufio.Writer
)

func getHttp(url string) ([]byte, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	_, timeout, transport := config.UpdateOptions()
	req.Header.Add("User-Agent", "siegfried/siegbot (+https://github.com/richardlehane/siegfried)")
	req.Header.Add("Cache-Control", "no-cache")
	timer := time.AfterFunc(timeout, func() {
		transport.CancelRequest(req)
	})
	defer timer.Stop()
	client := http.Client{
		Transport: transport,
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return ioutil.ReadAll(resp.Body)
}

type Update struct {
	SfVersion  [3]int
	SigCreated string
	GobSize    int
	LatestURL  string
}

func updateSigs() (string, error) {
	url, _, _ := config.UpdateOptions()
	if url == "" {
		return "Update is not available for this distribution of Siegfried", nil
	}
	response, err := getHttp(url)
	if err != nil {
		return "", err
	}
	var u Update
	if err := json.Unmarshal(response, &u); err != nil {
		return "", err
	}
	version := config.Version()
	if version[0] < u.SfVersion[0] || (u.SfVersion[0] == version[0] && version[1] < u.SfVersion[1]) {
		return "Your version of Siegfried is out of date; please install latest from http://www.itforarchivists.com/siegfried before continuing.", nil
	}
	s, err := siegfried.Load(config.Signature())
	if err == nil {
		if !s.Update(u.SigCreated) {
			return "You are already up to date!", nil
		}
	} else {
		// this hairy bit of golang exception handling is thanks to Ross! :)
		if _, err = os.Stat(config.Home()); err != nil {
			if os.IsNotExist(err) {
				err = os.MkdirAll(config.Home(), os.ModePerm)
				if err != nil {
					return "", fmt.Errorf("Siegfried: cannot create home directory %s, %v", config.Home(), err)
				}
			} else {
				return "", fmt.Errorf("Siegfried: error opening directory %s, %v", config.Home(), err)
			}
		}
	}
	fmt.Println("... downloading latest signature file ...")
	response, err = getHttp(u.LatestURL)
	if err != nil {
		return "", err
	}
	if len(response) != u.GobSize {
		return "", fmt.Errorf("Siegfried: error retrieving pronom.gob; expecting %d bytes, got %d bytes", u.GobSize, len(response))
	}
	err = ioutil.WriteFile(config.Signature(), response, os.ModePerm)
	if err != nil {
		return "", fmt.Errorf("Siegfried: error writing to directory, %v", err)
	}
	fmt.Printf("... writing %s ...\n", config.Signature())
	return "Your signature file has been updated", nil
}

func load() (*siegfried.Siegfried, error) {
	s, err := siegfried.Load(config.Signature())
	if err != nil {
		return nil, err
	}
	return s, nil
}

func identify(s *siegfried.Siegfried, p string) ([]string, error) {
	ids := make([]string, 0)
	file, err := os.Open(p)
	if err != nil {
		return nil, fmt.Errorf("failed to open %v, got: %v", p, err)
	}
	c, err := s.Identify(p, file)
	if c == nil {
		return nil, fmt.Errorf("failed to identify %v, got: %v", p, err)
	}
	for i := range c {
		ids = append(ids, i.String())
	}
	err = file.Close()
	if err != nil {
		return nil, err
	}
	return ids, nil
}

func multiIdentify(s *siegfried.Siegfried, r string) ([][]string, error) {
	set := make([][]string, 0)
	wf := func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			if *nr && path != r {
				return filepath.SkipDir
			}
			return nil
		}
		ids, err := identify(s, path)
		if err != nil {
			return err
		}
		set = append(set, ids)
		return nil
	}
	err := filepath.Walk(r, wf)
	return set, err
}

type res struct {
	path string
	sz   int64
	c    []core.Identification
	err  error
}

func printer(resc chan chan res, wg *sync.WaitGroup) {
	var csvRecord []string
	if *csvo {
		csvRecord = make([]string, 10)
	}
	for rr := range resc {
		r := <-rr
		if !config.Debug() && !*csvo {
			yamlWriter.WriteString(fileString(r.path, r.sz, r.err))
		}
		for _, v := range r.c {
			switch {
			case config.Debug():
			case *csvo:
				var errStr string
				if r.err != nil {
					errStr = r.err.Error()
				}
				csvRecord[0], csvRecord[1], csvRecord[2] = r.path, strconv.Itoa(int(r.sz)), errStr
				copy(csvRecord[3:], v.Csv())
				csvWriter.Write(csvRecord)
			default:
				yamlWriter.WriteString(v.Yaml())
			}
		}
		wg.Done()
	}
}

var lastPath string

func quitter() {
	timer := time.NewTimer(time.Minute * 25)
	<-timer.C
	panic(lastPath)
}

func multiIdentifyP(s *siegfried.Siegfried, r string) {
	go quitter()
	wg := &sync.WaitGroup{}
	runtime.GOMAXPROCS(PROCS)
	resc := make(chan chan res, CONCURRENT)
	go printer(resc, wg)
	wf := func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			if *nr && path != r {
				return filepath.SkipDir
			}
			return nil
		}
		lastPath = path
		wg.Add(1)
		rchan := make(chan res, 1)
		resc <- rchan
		go func() {
			file, err := os.Open(path)
			if err != nil {
				rchan <- res{"", 0, nil, fmt.Errorf("failed to open %v, got: %v", path, err)}
				return
			}
			c, err := s.Identify(path, file)
			if c == nil {
				file.Close()
				rchan <- res{"", 0, nil, fmt.Errorf("failed to identify %v, got: %v", path, err)}
				return
			}
			ids := make([]core.Identification, 0, 1)
			for id := range c {
				ids = append(ids, id)
			}
			rchan <- res{path, info.Size(), ids, err}
			file.Close()
		}()
		return nil
	}
	filepath.Walk(r, wf)
	wg.Wait()
	close(resc)
}

func PrintFile(name string, sz int64, err error) {
	fmt.Print(fileString(name, sz, err))
}

func fileString(name string, sz int64, err error) string {
	var errStr string
	if err != nil {
		errStr = fmt.Sprintf("\"%s\"", err.Error())
	}
	return fmt.Sprintf("---\nfilename : \"%s\"\nfilesize : %d\nerrors   : %s\nmatches  :\n", name, sz, errStr)
}

func main() {
	/*
		go func() {
			log.Println(http.ListenAndServe("localhost:6060", nil))
		}()
	*/
	flag.Parse()

	if *csvo {
		csvWriter = csv.NewWriter(os.Stdout)
		csvWriter.Write([]string{"filename", "filesize", "errors", "identifier", "id", "format name", "format version", "mimetype", "basis", "warning"})
	} else {
		yamlWriter = bufio.NewWriter(os.Stdout)
	}

	if *home != config.Home() {
		config.SetHome(*home)
	}

	if *sig != config.SignatureBase() {
		config.SetSignature(*sig)
	}

	if *version {
		version := config.Version()
		fmt.Printf("Siegfried version: %d.%d.%d\n", version[0], version[1], version[2])
		return
	}

	if *debug {
		config.SetDebug()
	}

	if *update {
		msg, err := updateSigs()
		if err != nil {
			log.Fatalf("Siegfried: error updating signature file, %v", err)
		}
		fmt.Println(msg)
		return
	}

	if *serve != "false" {
		fmt.Println("sf server not yet implemented; expect by v1")
	}

	if flag.NArg() != 1 {
		log.Fatal("Error: expecting a single file or directory argument")
	}

	var err error
	file, err := os.Open(flag.Arg(0))
	if err != nil {
		log.Fatalf("Error: error opening %v, got: %v", flag.Arg(0), err)
	}
	info, err := file.Stat()
	if err != nil {
		log.Fatalf("Error: error getting info for %v, got: %v", flag.Arg(0), err)
	}
	s, err := load()
	if err != nil {
		log.Fatalf("Error: error loading signature file, got: %v", err)

	}
	if info.IsDir() {
		file.Close()
		if config.Debug() {
			log.Fatalln("Error: when scanning in debug mode, give a file rather than a directory argument")
		}
		if !*csvo {
			yamlWriter.WriteString(s.Yaml())
		}
		multiIdentifyP(s, flag.Arg(0))
		if *csvo {
			csvWriter.Flush()
		} else {
			yamlWriter.Flush()
		}
		os.Exit(0)
	}
	c, err := s.Identify(flag.Arg(0), file)
	if c == nil {
		file.Close()
		log.Fatal(err)
	}
	if !config.Debug() && !*csvo {
		fmt.Print(s.Yaml())
		PrintFile(flag.Arg(0), info.Size(), err)
	}
	var csvRecord []string
	if *csvo {
		csvRecord = make([]string, 10)
	}
	for i := range c {
		switch {
		case config.Debug():
		case *csvo:
			var errStr string
			if err != nil {
				errStr = err.Error()
			}
			csvRecord[0], csvRecord[1], csvRecord[2] = flag.Arg(0), strconv.Itoa(int(info.Size())), errStr
			copy(csvRecord[3:], i.Csv())
			csvWriter.Write(csvRecord)
		default:
			fmt.Print(i.Yaml())
		}
	}
	file.Close()
	if *csvo {
		csvWriter.Flush()
	}

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

package main

import (
	"bufio"
	"encoding/csv"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"sync"
	"time"

	"github.com/richardlehane/siegfried"
	"github.com/richardlehane/siegfried/config"
	"github.com/richardlehane/siegfried/pkg/core"

	//_ "net/http/pprof"
)

const (
	CONCURRENT = 4
	PROCS      = -1
)

var (
	update  = flag.Bool("update", false, "update or install the default signature file")
	version = flag.Bool("version", false, "display version information")
	debug   = flag.Bool("debug", false, "scan in debug mode")
	nr      = flag.Bool("nr", false, "prevent automatic directory recursion")
	csvo    = flag.Bool("csv", false, "CSV output format")
	sig     = flag.String("sig", config.SignatureBase(), "set the signature file")
	home    = flag.String("home", config.Home(), "override the default home directory")
	serve   = flag.String("serve", "false", "not yet implemented - coming with v1")
)

var (
	csvWriter  *csv.Writer
	yamlWriter *bufio.Writer
)

func getHttp(url string) ([]byte, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	_, timeout, transport := config.UpdateOptions()
	req.Header.Add("User-Agent", "siegfried/siegbot (+https://github.com/richardlehane/siegfried)")
	req.Header.Add("Cache-Control", "no-cache")
	timer := time.AfterFunc(timeout, func() {
		transport.CancelRequest(req)
	})
	defer timer.Stop()
	client := http.Client{
		Transport: transport,
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return ioutil.ReadAll(resp.Body)
}

type Update struct {
	SfVersion  [3]int
	SigCreated string
	GobSize    int
	LatestURL  string
}

func updateSigs() (string, error) {
	url, _, _ := config.UpdateOptions()
	if url == "" {
		return "Update is not available for this distribution of Siegfried", nil
	}
	response, err := getHttp(url)
	if err != nil {
		return "", err
	}
	var u Update
	if err := json.Unmarshal(response, &u); err != nil {
		return "", err
	}
	version := config.Version()
	if version[0] < u.SfVersion[0] || (u.SfVersion[0] == version[0] && version[1] < u.SfVersion[1]) {
		return "Your version of Siegfried is out of date; please install latest from http://www.itforarchivists.com/siegfried before continuing.", nil
	}
	s, err := siegfried.Load(config.Signature())
	if err == nil {
		if !s.Update(u.SigCreated) {
			return "You are already up to date!", nil
		}
	} else {
		// this hairy bit of golang exception handling is thanks to Ross! :)
		if _, err = os.Stat(config.Home()); err != nil {
			if os.IsNotExist(err) {
				err = os.MkdirAll(config.Home(), os.ModePerm)
				if err != nil {
					return "", fmt.Errorf("Siegfried: cannot create home directory %s, %v", config.Home(), err)
				}
			} else {
				return "", fmt.Errorf("Siegfried: error opening directory %s, %v", config.Home(), err)
			}
		}
	}
	fmt.Println("... downloading latest signature file ...")
	response, err = getHttp(u.LatestURL)
	if err != nil {
		return "", err
	}
	if len(response) != u.GobSize {
		return "", fmt.Errorf("Siegfried: error retrieving pronom.gob; expecting %d bytes, got %d bytes", u.GobSize, len(response))
	}
	err = ioutil.WriteFile(config.Signature(), response, os.ModePerm)
	if err != nil {
		return "", fmt.Errorf("Siegfried: error writing to directory, %v", err)
	}
	fmt.Printf("... writing %s ...\n", config.Signature())
	return "Your signature file has been updated", nil
}

func load() (*siegfried.Siegfried, error) {
	s, err := siegfried.Load(config.Signature())
	if err != nil {
		return nil, err
	}
	return s, nil
}

func identify(s *siegfried.Siegfried, p string) ([]string, error) {
	ids := make([]string, 0)
	file, err := os.Open(p)
	if err != nil {
		return nil, fmt.Errorf("failed to open %v, got: %v", p, err)
	}
	c, err := s.Identify(p, file)
	if c == nil {
		return nil, fmt.Errorf("failed to identify %v, got: %v", p, err)
	}
	for i := range c {
		ids = append(ids, i.String())
	}
	err = file.Close()
	if err != nil {
		return nil, err
	}
	return ids, nil
}

func multiIdentify(s *siegfried.Siegfried, r string) ([][]string, error) {
	set := make([][]string, 0)
	wf := func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			if *nr && path != r {
				return filepath.SkipDir
			}
			return nil
		}
		ids, err := identify(s, path)
		if err != nil {
			return err
		}
		set = append(set, ids)
		return nil
	}
	err := filepath.Walk(r, wf)
	return set, err
}

type res struct {
	path string
	sz   int64
	c    []core.Identification
	err  error
}

func printer(resc chan chan res, wg *sync.WaitGroup) {
	var csvRecord []string
	if *csvo {
		csvRecord = make([]string, 10)
	}
	for rr := range resc {
		r := <-rr
		if !config.Debug() && !*csvo {
			yamlWriter.WriteString(fileString(r.path, r.sz, r.err))
		}
		for _, v := range r.c {
			switch {
			case config.Debug():
			case *csvo:
				var errStr string
				if r.err != nil {
					errStr = r.err.Error()
				}
				csvRecord[0], csvRecord[1], csvRecord[2] = r.path, strconv.Itoa(int(r.sz)), errStr
				copy(csvRecord[3:], v.Csv())
				csvWriter.Write(csvRecord)
			default:
				yamlWriter.WriteString(v.Yaml())
			}
		}
		wg.Done()
	}
}

/*
var lastPath string

func quitter() {
	timer := time.NewTimer(time.Minute * 25)
	<-timer.C
	panic(lastPath)
}
*/
func multiIdentifyP(s *siegfried.Siegfried, r string) {
	//go quitter()
	wg := &sync.WaitGroup{}
	runtime.GOMAXPROCS(PROCS)
	resc := make(chan chan res, CONCURRENT)
	go printer(resc, wg)
	wf := func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			if *nr && path != r {
				return filepath.SkipDir
			}
			return nil
		}
		lastPath = path
		wg.Add(1)
		rchan := make(chan res, 1)
		resc <- rchan
		go func() {
			file, err := os.Open(path)
			if err != nil {
				rchan <- res{"", 0, nil, fmt.Errorf("failed to open %v, got: %v", path, err)}
				return
			}
			c, err := s.Identify(path, file)
			if c == nil {
				file.Close()
				rchan <- res{"", 0, nil, fmt.Errorf("failed to identify %v, got: %v", path, err)}
				return
			}
			ids := make([]core.Identification, 0, 1)
			for id := range c {
				ids = append(ids, id)
			}
			rchan <- res{path, info.Size(), ids, err}
			file.Close()
		}()
		return nil
	}
	filepath.Walk(r, wf)
	wg.Wait()
	close(resc)
}

func PrintFile(name string, sz int64, err error) {
	fmt.Print(fileString(name, sz, err))
}

func fileString(name string, sz int64, err error) string {
	var errStr string
	if err != nil {
		errStr = fmt.Sprintf("\"%s\"", err.Error())
	}
	return fmt.Sprintf("---\nfilename : \"%s\"\nfilesize : %d\nerrors   : %s\nmatches  :\n", name, sz, errStr)
}

func main() {
	/*
		go func() {
			log.Println(http.ListenAndServe("localhost:6060", nil))
		}()
	*/
	flag.Parse()

	if *csvo {
		csvWriter = csv.NewWriter(os.Stdout)
		csvWriter.Write([]string{"filename", "filesize", "errors", "identifier", "id", "format name", "format version", "mimetype", "basis", "warning"})
	} else {
		yamlWriter = bufio.NewWriter(os.Stdout)
	}

	if *home != config.Home() {
		config.SetHome(*home)
	}

	if *sig != config.SignatureBase() {
		config.SetSignature(*sig)
	}

	if *version {
		version := config.Version()
		fmt.Printf("Siegfried version: %d.%d.%d\n", version[0], version[1], version[2])
		return
	}

	if *debug {
		config.SetDebug()
	}

	if *update {
		msg, err := updateSigs()
		if err != nil {
			log.Fatalf("Siegfried: error updating signature file, %v", err)
		}
		fmt.Println(msg)
		return
	}

	if *serve != "false" {
		fmt.Println("sf server not yet implemented; expect by v1")
	}

	if flag.NArg() != 1 {
		log.Fatal("Error: expecting a single file or directory argument")
	}

	var err error
	file, err := os.Open(flag.Arg(0))
	if err != nil {
		log.Fatalf("Error: error opening %v, got: %v", flag.Arg(0), err)
	}
	info, err := file.Stat()
	if err != nil {
		log.Fatalf("Error: error getting info for %v, got: %v", flag.Arg(0), err)
	}
	s, err := load()
	if err != nil {
		log.Fatalf("Error: error loading signature file, got: %v", err)

	}
	if info.IsDir() {
		file.Close()
		if config.Debug() {
			log.Fatalln("Error: when scanning in debug mode, give a file rather than a directory argument")
		}
		if !*csvo {
			yamlWriter.WriteString(s.Yaml())
		}
		multiIdentifyP(s, flag.Arg(0))
		if *csvo {
			csvWriter.Flush()
		} else {
			yamlWriter.Flush()
		}
		os.Exit(0)
	}
	c, err := s.Identify(flag.Arg(0), file)
	if c == nil {
		file.Close()
		log.Fatal(err)
	}
	if !config.Debug() && !*csvo {
		fmt.Print(s.Yaml())
		PrintFile(flag.Arg(0), info.Size(), err)
	}
	var csvRecord []string
	if *csvo {
		csvRecord = make([]string, 10)
	}
	for i := range c {
		switch {
		case config.Debug():
		case *csvo:
			var errStr string
			if err != nil {
				errStr = err.Error()
			}
			csvRecord[0], csvRecord[1], csvRecord[2] = flag.Arg(0), strconv.Itoa(int(info.Size())), errStr
			copy(csvRecord[3:], i.Csv())
			csvWriter.Write(csvRecord)
		default:
			fmt.Print(i.Yaml())
		}
	}
	file.Close()
	if *csvo {
		csvWriter.Flush()
	}

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


<commit_before>package main

import (
	"bufio"
	log "code.google.com/p/log4go"
	"github.com/Terry-Mao/goim/define"
	"github.com/Terry-Mao/goim/libs/io/ioutil"
	"net"
	"sync"
	"time"
)

const (
	maxPackLen    = 1 << 10
	rawHeaderLen  = int16(16)
	packLenSize   = 4
	headerLenSize = 2
	maxPackIntBuf = 4
)

// InitTCP listen all tcp.bind and start accept connections.
func InitTCP() (err error) {
	var (
		listener *net.TCPListener
		addr     *net.TCPAddr
	)
	for _, bind := range Conf.TCPBind {
		if addr, err = net.ResolveTCPAddr("tcp4", bind); err != nil {
			log.Error("net.ResolveTCPAddr(\"tcp4\", \"%s\") error(%v)", bind, err)
			return
		}
		if listener, err = net.ListenTCP("tcp4", addr); err != nil {
			log.Error("net.ListenTCP(\"tcp4\", \"%s\") error(%v)", bind, err)
			return
		}
		log.Debug("start tcp listen: \"%s\"", bind)
		// split N core accept
		for i := 0; i < Conf.MaxProc; i++ {
			go acceptTCP(DefaultServer, listener)
		}
	}
	return
}

// Accept accepts connections on the listener and serves requests
// for each incoming connection.  Accept blocks; the caller typically
// invokes it in a go statement.
func acceptTCP(server *Server, lis *net.TCPListener) {
	var (
		conn *net.TCPConn
		err  error
		r    int
	)
	for {
		if conn, err = lis.AcceptTCP(); err != nil {
			// if listener close then return
			log.Error("listener.Accept(\"%s\") error(%v)", lis.Addr().String(), err)
			return
		}
		if err = conn.SetKeepAlive(Conf.TCPKeepalive); err != nil {
			log.Error("conn.SetKeepAlive() error(%v)", err)
			return
		}
		if err = conn.SetReadBuffer(Conf.TCPSndbuf); err != nil {
			log.Error("conn.SetReadBuffer() error(%v)", err)
			return
		}
		if err = conn.SetWriteBuffer(Conf.TCPRcvbuf); err != nil {
			log.Error("conn.SetWriteBuffer() error(%v)", err)
			return
		}
		go serveTCP(server, conn, r)
		if r++; r == maxInt {
			r = 0
		}
	}
}

func serveTCP(server *Server, conn *net.TCPConn, r int) {
	var (
		// bufpool
		rrp = server.round.Reader(r) // reader
		wrp = server.round.Writer(r) // writer
		// timer
		tr = server.round.Timer(r)
		// buf
		rr = NewBufioReaderSize(rrp, conn, Conf.ReadBufSize)  // reader buf
		wr = NewBufioWriterSize(wrp, conn, Conf.WriteBufSize) // writer buf
		// ip addr
		lAddr = conn.LocalAddr().String()
		rAddr = conn.RemoteAddr().String()
	)
	log.Debug("start tcp serve \"%s\" with \"%s\"", lAddr, rAddr)
	server.serveTCP(conn, rrp, wrp, rr, wr, tr)
}

func (server *Server) serveTCP(conn *net.TCPConn, rrp, wrp *sync.Pool, rr *bufio.Reader, wr *bufio.Writer, tr *Timer) {
	var (
		b   *Bucket
		p   *Proto
		key string
		hb  time.Duration // heartbeat
		err error
		trd *TimerData
		ch  = NewChannel(Conf.CliProto, Conf.SvrProto, define.NoRoom)
	)
	// auth
	if trd, err = tr.Add(Conf.HandshakeTimeout, conn); err != nil {
		log.Error("handshake: timer.Add() error(%v)", err)
	} else {
		if key, hb, err = server.authTCP(rr, wr, ch); err != nil {
			log.Error("handshake: server.auth error(%v)", err)
		}
		//deltimer
		tr.Del(trd)
	}
	// failed
	if err != nil {
		if err = conn.Close(); err != nil {
			log.Error("handshake: conn.Close() error(%v)", err)
		}
		PutBufioReader(rrp, rr)
		return
	}
	// register key->channel
	b = server.Bucket(key)
	b.Put(key, ch)
	// hanshake ok start dispatch goroutine
	go server.dispatchTCP(conn, wrp, wr, ch, hb, tr)
	for {
		// fetch a proto from channel free list
		if p, err = ch.CliProto.Set(); err != nil {
			// if full, simply close connection
			log.Error("%s fetch client proto error(%v)", key, err)
			goto failed
		}
		// parse request protocol
		if err = server.readTCPRequest(rr, p); err != nil {
			log.Error("%s read client request error(%v)", key, err)
			goto failed
		}
		// send to writer
		ch.CliProto.SetAdv()
		ch.Signal()
	}
failed:
	// dialog finish
	// may call twice
	if err = conn.Close(); err != nil {
		log.Error("reader: conn.Close() error(%v)", err)
	}
	PutBufioReader(rrp, rr)
	b.Del(key)
	log.Debug("wake up dispatch goroutine")
	ch.Finish()
	if err = server.operator.Disconnect(key, ch.RoomId); err != nil {
		log.Error("%s operator do disconnect error(%v)", key, err)
	}
	log.Debug("%s serverconn goroutine exit", key)
	return
}

// dispatch accepts connections on the listener and serves requests
// for each incoming connection.  dispatch blocks; the caller typically
// invokes it in a go statement.
func (server *Server) dispatchTCP(conn *net.TCPConn, wrp *sync.Pool, wr *bufio.Writer, ch *Channel, hb time.Duration, tr *Timer) {
	var (
		p   *Proto
		err error
		trd *TimerData
	)
	log.Debug("start dispatch goroutine")
	if trd, err = tr.Add(hb, conn); err != nil {
		log.Error("dispatch: timer.Add() error(%v)", err)
		goto failed
	}
	for {
		if !ch.Ready() {
			goto failed
		}
		// fetch message from clibox(client send)
		for {
			if p, err = ch.CliProto.Get(); err != nil {
				// must be empty error
				break
			}
			if p.Operation == define.OP_HEARTBEAT {
				// Use a previous timer value if difference between it and a new
				// value is less than TIMER_LAZY_DELAY milliseconds: this allows
				// to minimize the minheap operations for fast connections.
				if !trd.Lazy(hb) {
					tr.Del(trd)
					if trd, err = tr.Add(hb, conn); err != nil {
						log.Error("dispatch: timer.Add() error(%v)", err)
						goto failed
					}
				}
				// heartbeat
				p.Body = nil
				p.Operation = define.OP_HEARTBEAT_REPLY
			} else {
				// process message
				if err = server.operator.Operate(p); err != nil {
					log.Error("operator.Operate() error(%v)", err)
					goto failed
				}
			}
			if err = server.writeTCPResponse(wr, p); err != nil {
				log.Error("server.writeTCPResponse() error(%v)", err)
				goto failed
			}
			ch.CliProto.GetAdv()
		}
		// fetch message from svrbox(server send)
		for {
			if p, err = ch.SvrProto.Get(); err != nil {
				// must be empty error
				break
			}
			// just forward the message
			if err = server.writeTCPResponse(wr, p); err != nil {
				log.Error("server.writeTCPResponse() error(%v)", err)
				goto failed
			}
			ch.SvrProto.GetAdv()
		}
	}
failed:
	// wake reader up
	if err = conn.Close(); err != nil {
		log.Warn("conn.Close() error(%v)", err)
	}
	// deltimer
	tr.Del(trd)
	PutBufioWriter(wrp, wr)
	log.Debug("dispatch goroutine exit")
	return
}

// auth for goim handshake with client, use rsa & aes.
func (server *Server) authTCP(rr *bufio.Reader, wr *bufio.Writer, ch *Channel) (subKey string, heartbeat time.Duration, err error) {
	var p *Proto
	// WARN
	// don't adv the cli proto, after auth simply discard it.
	if p, err = ch.CliProto.Set(); err != nil {
		return
	}
	if err = server.readTCPRequest(rr, p); err != nil {
		return
	}
	if p.Operation != define.OP_AUTH {
		log.Warn("auth operation not valid: %d", p.Operation)
		err = ErrOperation
		return
	}
	if subKey, ch.RoomId, heartbeat, err = server.operator.Connect(p); err != nil {
		log.Error("operator.Connect error(%v)", err)
		return
	}
	p.Body = nil
	p.Operation = define.OP_AUTH_REPLY
	if err = server.writeTCPResponse(wr, p); err != nil {
		log.Error("[%s] server.sendTCPResponse() error(%v)", subKey, err)
	}
	return
}

// readRequest
func (server *Server) readTCPRequest(rr *bufio.Reader, proto *Proto) (err error) {
	var (
		packLen   int32
		headerLen int16
		bodyLen   int
	)
	if packLen, err = ioutil.ReadBigEndianInt32(rr); err != nil {
		return
	}
	if Conf.Debug {
		log.Debug("packLen: %d", packLen)
	}
	if packLen > maxPackLen {
		return ErrProtoPackLen
	}
	if headerLen, err = ioutil.ReadBigEndianInt16(rr); err != nil {
		return
	}
	if Conf.Debug {
		log.Debug("headerLen: %d", headerLen)
	}
	if headerLen != rawHeaderLen {
		return ErrProtoHeaderLen
	}
	if proto.Ver, err = ioutil.ReadBigEndianInt16(rr); err != nil {
		return
	}
	if Conf.Debug {
		log.Debug("protoVer: %d", proto.Ver)
	}
	if proto.Operation, err = ioutil.ReadBigEndianInt32(rr); err != nil {
		return
	}
	if Conf.Debug {
		log.Debug("operation: %d", proto.Operation)
	}
	if proto.SeqId, err = ioutil.ReadBigEndianInt32(rr); err != nil {
		return
	}
	if Conf.Debug {
		log.Debug("seqId: %d", proto.SeqId)
	}
	bodyLen = int(packLen - int32(headerLen))
	if Conf.Debug {
		log.Debug("read body len: %d", bodyLen)
	}
	if bodyLen > 0 {
		// TODO reuse buf
		proto.Body = make([]byte, bodyLen)
		if err = ioutil.ReadAll(rr, proto.Body); err != nil {
			log.Error("body: ReadAll() error(%v)", err)
			return
		}
	} else {
		proto.Body = nil
	}
	if Conf.Debug {
		log.Debug("read proto: %v", proto)
	}
	return
}

// sendResponse send resp to client, sendResponse must be goroutine safe.
func (server *Server) writeTCPResponse(wr *bufio.Writer, proto *Proto) (err error) {
	if debug {
		log.Debug("write proto: %v", proto)
	}
	if err = ioutil.WriteBigEndianInt32(wr, int32(rawHeaderLen)+int32(len(proto.Body))); err != nil {
		return
	}
	if err = ioutil.WriteBigEndianInt16(wr, rawHeaderLen); err != nil {
		return
	}
	if err = ioutil.WriteBigEndianInt16(wr, proto.Ver); err != nil {
		return
	}
	if err = ioutil.WriteBigEndianInt32(wr, proto.Operation); err != nil {
		return
	}
	if err = ioutil.WriteBigEndianInt32(wr, proto.SeqId); err != nil {
		return
	}
	if proto.Body != nil {
		if _, err = wr.Write(proto.Body); err != nil {
			return
		}
	}
	if err = wr.Flush(); err != nil {
		log.Error("tcp wr.Flush() error(%v)", err)
	}
	proto.Reset()
	return
}
<commit_msg>fix bug<commit_after>package main

import (
	"bufio"
	log "code.google.com/p/log4go"
	"github.com/Terry-Mao/goim/define"
	"github.com/Terry-Mao/goim/libs/io/ioutil"
	"net"
	"sync"
	"time"
)

const (
	maxPackLen    = 1 << 10
	rawHeaderLen  = int16(16)
	packLenSize   = 4
	headerLenSize = 2
	maxPackIntBuf = 4
)

// InitTCP listen all tcp.bind and start accept connections.
func InitTCP() (err error) {
	var (
		listener *net.TCPListener
		addr     *net.TCPAddr
	)
	for _, bind := range Conf.TCPBind {
		if addr, err = net.ResolveTCPAddr("tcp4", bind); err != nil {
			log.Error("net.ResolveTCPAddr(\"tcp4\", \"%s\") error(%v)", bind, err)
			return
		}
		if listener, err = net.ListenTCP("tcp4", addr); err != nil {
			log.Error("net.ListenTCP(\"tcp4\", \"%s\") error(%v)", bind, err)
			return
		}
		log.Debug("start tcp listen: \"%s\"", bind)
		// split N core accept
		for i := 0; i < Conf.MaxProc; i++ {
			go acceptTCP(DefaultServer, listener)
		}
	}
	return
}

// Accept accepts connections on the listener and serves requests
// for each incoming connection.  Accept blocks; the caller typically
// invokes it in a go statement.
func acceptTCP(server *Server, lis *net.TCPListener) {
	var (
		conn *net.TCPConn
		err  error
		r    int
	)
	for {
		if conn, err = lis.AcceptTCP(); err != nil {
			// if listener close then return
			log.Error("listener.Accept(\"%s\") error(%v)", lis.Addr().String(), err)
			return
		}
		if err = conn.SetKeepAlive(Conf.TCPKeepalive); err != nil {
			log.Error("conn.SetKeepAlive() error(%v)", err)
			return
		}
		if err = conn.SetReadBuffer(Conf.TCPSndbuf); err != nil {
			log.Error("conn.SetReadBuffer() error(%v)", err)
			return
		}
		if err = conn.SetWriteBuffer(Conf.TCPRcvbuf); err != nil {
			log.Error("conn.SetWriteBuffer() error(%v)", err)
			return
		}
		go serveTCP(server, conn, r)
		if r++; r == maxInt {
			r = 0
		}
	}
}

func serveTCP(server *Server, conn *net.TCPConn, r int) {
	var (
		// bufpool
		rrp = server.round.Reader(r) // reader
		wrp = server.round.Writer(r) // writer
		// timer
		tr = server.round.Timer(r)
		// buf
		rr = NewBufioReaderSize(rrp, conn, Conf.ReadBufSize)  // reader buf
		wr = NewBufioWriterSize(wrp, conn, Conf.WriteBufSize) // writer buf
		// ip addr
		lAddr = conn.LocalAddr().String()
		rAddr = conn.RemoteAddr().String()
	)
	log.Debug("start tcp serve \"%s\" with \"%s\"", lAddr, rAddr)
	server.serveTCP(conn, rrp, wrp, rr, wr, tr)
}

func (server *Server) serveTCP(conn *net.TCPConn, rrp, wrp *sync.Pool, rr *bufio.Reader, wr *bufio.Writer, tr *Timer) {
	var (
		b   *Bucket
		p   *Proto
		key string
		hb  time.Duration // heartbeat
		err error
		trd *TimerData
		ch  = NewChannel(Conf.CliProto, Conf.SvrProto, define.NoRoom)
	)
	// auth
	if trd, err = tr.Add(Conf.HandshakeTimeout, conn); err != nil {
		log.Error("handshake: timer.Add() error(%v)", err)
	} else {
		if key, hb, err = server.authTCP(rr, wr, ch); err != nil {
			log.Error("handshake: server.auth error(%v)", err)
		}
		//deltimer
		tr.Del(trd)
	}
	// failed
	if err != nil {
		if err = conn.Close(); err != nil {
			log.Error("handshake: conn.Close() error(%v)", err)
		}
		PutBufioReader(rrp, rr)
		return
	}
	// register key->channel
	b = server.Bucket(key)
	b.Put(key, ch)
	// hanshake ok start dispatch goroutine
	go server.dispatchTCP(conn, wrp, wr, ch, hb, tr)
	for {
		// fetch a proto from channel free list
		if p, err = ch.CliProto.Set(); err != nil {
			// if full, simply close connection
			log.Error("%s fetch client proto error(%v)", key, err)
			goto failed
		}
		// parse request protocol
		if err = server.readTCPRequest(rr, p); err != nil {
			log.Error("%s read client request error(%v)", key, err)
			goto failed
		}
		// send to writer
		ch.CliProto.SetAdv()
		ch.Signal()
	}
failed:
	// dialog finish
	// may call twice
	if err = conn.Close(); err != nil {
		log.Error("reader: conn.Close() error(%v)", err)
	}
	PutBufioReader(rrp, rr)
	b.Del(key)
	log.Debug("wake up dispatch goroutine")
	ch.Finish()
	if err = server.operator.Disconnect(key, ch.RoomId); err != nil {
		log.Error("%s operator do disconnect error(%v)", key, err)
	}
	log.Debug("%s serverconn goroutine exit", key)
	return
}

// dispatch accepts connections on the listener and serves requests
// for each incoming connection.  dispatch blocks; the caller typically
// invokes it in a go statement.
func (server *Server) dispatchTCP(conn *net.TCPConn, wrp *sync.Pool, wr *bufio.Writer, ch *Channel, hb time.Duration, tr *Timer) {
	var (
		p   *Proto
		err error
		trd *TimerData
	)
	log.Debug("start dispatch goroutine")
	if trd, err = tr.Add(hb, conn); err != nil {
		log.Error("dispatch: timer.Add() error(%v)", err)
		goto failed
	}
	for {
		if !ch.Ready() {
			goto failed
		}
		// fetch message from clibox(client send)
		for {
			if p, err = ch.CliProto.Get(); err != nil {
				// must be empty error
				break
			}
			if p.Operation == define.OP_HEARTBEAT {
				// Use a previous timer value if difference between it and a new
				// value is less than TIMER_LAZY_DELAY milliseconds: this allows
				// to minimize the minheap operations for fast connections.
				if !trd.Lazy(hb) {
					tr.Del(trd)
					if trd, err = tr.Add(hb, conn); err != nil {
						log.Error("dispatch: timer.Add() error(%v)", err)
						goto failed
					}
				}
				// heartbeat
				p.Body = nil
				p.Operation = define.OP_HEARTBEAT_REPLY
			} else {
				// process message
				if err = server.operator.Operate(p); err != nil {
					log.Error("operator.Operate() error(%v)", err)
					goto failed
				}
			}
			if err = server.writeTCPResponse(wr, p); err != nil {
				log.Error("server.writeTCPResponse() error(%v)", err)
				goto failed
			}
			ch.CliProto.GetAdv()
		}
		// fetch message from svrbox(server send)
		for {
			if p, err = ch.SvrProto.Get(); err != nil {
				// must be empty error
				break
			}
			// just forward the message
			if err = server.writeTCPResponse(wr, p); err != nil {
				log.Error("server.writeTCPResponse() error(%v)", err)
				goto failed
			}
			ch.SvrProto.GetAdv()
		}
	}
failed:
	// wake reader up
	if err = conn.Close(); err != nil {
		log.Warn("conn.Close() error(%v)", err)
	}
	// deltimer
	tr.Del(trd)
	PutBufioWriter(wrp, wr)
	log.Debug("dispatch goroutine exit")
	return
}

// auth for goim handshake with client, use rsa & aes.
func (server *Server) authTCP(rr *bufio.Reader, wr *bufio.Writer, ch *Channel) (subKey string, heartbeat time.Duration, err error) {
	var p *Proto
	// WARN
	// don't adv the cli proto, after auth simply discard it.
	if p, err = ch.CliProto.Set(); err != nil {
		return
	}
	if err = server.readTCPRequest(rr, p); err != nil {
		return
	}
	if p.Operation != define.OP_AUTH {
		log.Warn("auth operation not valid: %d", p.Operation)
		err = ErrOperation
		return
	}
	if subKey, ch.RoomId, heartbeat, err = server.operator.Connect(p); err != nil {
		log.Error("operator.Connect error(%v)", err)
		return
	}
	p.Body = nil
	p.Operation = define.OP_AUTH_REPLY
	if err = server.writeTCPResponse(wr, p); err != nil {
		log.Error("[%s] server.sendTCPResponse() error(%v)", subKey, err)
	}
	return
}

// readRequest
func (server *Server) readTCPRequest(rr *bufio.Reader, proto *Proto) (err error) {
	var (
		packLen   int32
		headerLen int16
		bodyLen   int
	)
	if packLen, err = ioutil.ReadBigEndianInt32(rr); err != nil {
		return
	}
	if Conf.Debug {
		log.Debug("packLen: %d", packLen)
	}
	if packLen > maxPackLen {
		return ErrProtoPackLen
	}
	if headerLen, err = ioutil.ReadBigEndianInt16(rr); err != nil {
		return
	}
	if Conf.Debug {
		log.Debug("headerLen: %d", headerLen)
	}
	if headerLen != rawHeaderLen {
		return ErrProtoHeaderLen
	}
	if proto.Ver, err = ioutil.ReadBigEndianInt16(rr); err != nil {
		return
	}
	if Conf.Debug {
		log.Debug("protoVer: %d", proto.Ver)
	}
	if proto.Operation, err = ioutil.ReadBigEndianInt32(rr); err != nil {
		return
	}
	if Conf.Debug {
		log.Debug("operation: %d", proto.Operation)
	}
	if proto.SeqId, err = ioutil.ReadBigEndianInt32(rr); err != nil {
		return
	}
	if Conf.Debug {
		log.Debug("seqId: %d", proto.SeqId)
	}
	bodyLen = int(packLen - int32(headerLen))
	if Conf.Debug {
		log.Debug("read body len: %d", bodyLen)
	}
	if bodyLen > 0 {
		// TODO reuse buf
		proto.Body = make([]byte, bodyLen)
		if err = ioutil.ReadAll(rr, proto.Body); err != nil {
			log.Error("body: ReadAll() error(%v)", err)
			return
		}
	} else {
		proto.Body = nil
	}
	if Conf.Debug {
		log.Debug("read proto: %v", proto)
	}
	return
}

// sendResponse send resp to client, sendResponse must be goroutine safe.
func (server *Server) writeTCPResponse(wr *bufio.Writer, proto *Proto) (err error) {
	if Conf.Debug {
		log.Debug("write proto: %v", proto)
	}
	if err = ioutil.WriteBigEndianInt32(wr, int32(rawHeaderLen)+int32(len(proto.Body))); err != nil {
		return
	}
	if err = ioutil.WriteBigEndianInt16(wr, rawHeaderLen); err != nil {
		return
	}
	if err = ioutil.WriteBigEndianInt16(wr, proto.Ver); err != nil {
		return
	}
	if err = ioutil.WriteBigEndianInt32(wr, proto.Operation); err != nil {
		return
	}
	if err = ioutil.WriteBigEndianInt32(wr, proto.SeqId); err != nil {
		return
	}
	if proto.Body != nil {
		if _, err = wr.Write(proto.Body); err != nil {
			return
		}
	}
	if err = wr.Flush(); err != nil {
		log.Error("tcp wr.Flush() error(%v)", err)
	}
	proto.Reset()
	return
}
<|endoftext|>


<commit_before>// Description: conf 项目配置项，也用于读取当前目录下的app.conf中的项目
// Author: ZHU HAIHUA
// Since: 2016-02-26 19:19
package conf

import (
	"fmt"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"os"
	"strings"
	"reflect"
)

type Config struct {
	SERVER struct {
		RUNMODE string `yaml:runmode`
		PORT    string `yaml:"port"`
	}

	DATABASE struct {
		HOST     string `yaml:"host"`
		NAME     string `yaml:"name"`
		USER     string `yaml:"user"`
		PASSWORD string `yaml:"password"`
	}

	EXT map[string]interface{} `yaml:"ext,flow"`
}

// Ext will return the value of the EXT config, the keys is a string
// separated by DOT(.). If you provide a default value, this method
// will return the it while the key cannot be found. otherwise it
// will raise a panic!
func (c *Config) Ext(keys string, defaultVal... interface{}) (interface{}) {
	r, e := c.ExtSep(keys, ".")
	if e != nil || r == nil {
		if len(defaultVal) > 0 {
			return defaultVal[0]
		} else {
			panic(e)
		}
	} else {
		return r
	}
}

func (c *Config) ExtString(keys string, defaultVal... interface{}) (string) {
	return fmt.Sprintf("%v", c.Ext(keys, defaultVal...))
}

func (c *Config) ExtInt(keys string, defaultVal... interface{}) (int) {
	return c.Ext(keys, defaultVal...).(int)
}

func (c *Config) ExtInt8(keys string, defaultVal... interface{}) (int8) {
	return c.Ext(keys, defaultVal...).(int8)
}

func (c *Config) ExtInt16(keys string, defaultVal... interface{}) (int16) {
	return c.Ext(keys, defaultVal...).(int16)
}

func (c *Config) ExtInt32(keys string, defaultVal... interface{}) (int32) {
	return c.Ext(keys, defaultVal...).(int32)
}

func (c *Config) ExtInt64(keys string, defaultVal... interface{}) (int64) {
	return c.Ext(keys, defaultVal...).(int64)
}

func (c *Config) ExtBool(keys string, defaultVal... interface{}) (bool) {
	return c.Ext(keys, defaultVal...).(bool)
}

func (c *Config) ExtFloat64(keys string, defaultVal... interface{}) (float64) {
	return c.Ext(keys, defaultVal...).(float64)
}

func (c *Config) ExtFloat32(keys string, defaultVal... interface{}) (float32) {
	return c.Ext(keys, defaultVal...).(float32)
}

// Ext will return the value of the EXT config, the keys is separated
// by the given sep string.
func (c *Config) ExtSep(keys, sep string) (interface{}, error) {
	ks := strings.Split(keys, sep);
	var result interface{}
	var isFinal, success bool
	result = c.EXT
	for _, k := range ks {
		result, isFinal, success = find(result, k)
		if !success {
			return "", fmt.Errorf("no such key: %v", k)
		} else if isFinal {
			break
		}
	}

	if success {
		return result, nil
	} else {
		return "", fmt.Errorf("not found")
	}
}

func find(v interface{}, key interface{}) (result interface{}, isFinal, success bool) {
	switch m := v.(type) {
	case map[string]interface{}:
		result, success = m[key.(string)]
		//success = true
		isFinal = (reflect.TypeOf(result) != nil && reflect.TypeOf(result).Kind() != reflect.Map)
	case map[interface{}]interface{}:
		result, success = m[key]
		//success = true
		isFinal = (reflect.TypeOf(result) != nil && reflect.TypeOf(result).Kind() != reflect.Map)
	}
	return
}

var (
	Conf = Config{}
)

func LoadConf(path string) {
	c, err := ioutil.ReadFile(path)
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		return
	}
	yaml.Unmarshal(c, &Conf)
}

func Ext(keys string, defVal... interface{}) interface {}{
	return Conf.Ext(keys, defVal...)
}

// ExtString is a shorcut for (c *Config) ExtString()
func ExtString(keys string, defVal... interface{}) string {
	return Conf.ExtString(keys, defVal...)
}

// ExtBool is a shorcut for (c *Config) ExtBool()
func ExtBool(keys string, defVal... interface{}) bool {
	return Conf.ExtBool(keys, defVal...)
}

// ExtInt is a shorcut for (c *Config) ExtInt()
func ExtInt(keys string, defVal... interface{}) int {
	return Conf.ExtInt(keys, defVal...)
}

// ExtInt8 is a shorcut for (c *Config) ExtInt8()
func ExtInt8(keys string, defVal... interface{}) int8 {
	return Conf.ExtInt8(keys, defVal...)
}

// ExtInt16 is a shorcut for (c *Config) ExtInt16()
func ExtInt16(keys string, defVal... interface{}) int16 {
	return Conf.ExtInt16(keys, defVal...)
}

// ExtInt32 is a shorcut for (c *Config) ExtInt32()
func ExtInt32(keys string, defVal... interface{}) int32 {
	return Conf.ExtInt32(keys, defVal...)
}

// ExtInt64 is a shorcut for (c *Config) ExtInt64()
func ExtInt64(keys string, defVal... interface{}) int64 {
	return Conf.ExtInt64(keys, defVal...)
}

// ExtFloat32 is a shorcut for (c *Config) ExtFloat32()
func ExtFloat32(keys string, defVal... interface{}) float32 {
	return Conf.ExtFloat32(keys, defVal...)
}

// ExtFloat64 is a shorcut for (c *Config) ExtFloat64()
func ExtFloat64(keys string, defVal... interface{}) float64 {
	return Conf.ExtFloat64(keys, defVal...)
}

func init() {
	LoadConf("conf.yml")
}
<commit_msg>remove unused code<commit_after>// Description: conf 项目配置项，也用于读取当前目录下的app.conf中的项目
// Author: ZHU HAIHUA
// Since: 2016-02-26 19:19
package conf

import (
	"fmt"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"os"
	"strings"
	"reflect"
)

type Config struct {
	SERVER struct {
		RUNMODE string `yaml:runmode`
		PORT    string `yaml:"port"`
	}

	DATABASE struct {
		HOST     string `yaml:"host"`
		NAME     string `yaml:"name"`
		USER     string `yaml:"user"`
		PASSWORD string `yaml:"password"`
	}

	EXT map[string]interface{} `yaml:"ext,flow"`
}

// Ext will return the value of the EXT config, the keys is a string
// separated by DOT(.). If you provide a default value, this method
// will return the it while the key cannot be found. otherwise it
// will raise a panic!
func (c *Config) Ext(keys string, defaultVal... interface{}) (interface{}) {
	r, e := c.ExtSep(keys, ".")
	if e != nil || r == nil {
		if len(defaultVal) > 0 {
			return defaultVal[0]
		} else {
			panic(e)
		}
	} else {
		return r
	}
}

func (c *Config) ExtString(keys string, defaultVal... interface{}) (string) {
	return fmt.Sprintf("%v", c.Ext(keys, defaultVal...))
}

func (c *Config) ExtInt(keys string, defaultVal... interface{}) (int) {
	return c.Ext(keys, defaultVal...).(int)
}

func (c *Config) ExtInt8(keys string, defaultVal... interface{}) (int8) {
	return c.Ext(keys, defaultVal...).(int8)
}

func (c *Config) ExtInt16(keys string, defaultVal... interface{}) (int16) {
	return c.Ext(keys, defaultVal...).(int16)
}

func (c *Config) ExtInt32(keys string, defaultVal... interface{}) (int32) {
	return c.Ext(keys, defaultVal...).(int32)
}

func (c *Config) ExtInt64(keys string, defaultVal... interface{}) (int64) {
	return c.Ext(keys, defaultVal...).(int64)
}

func (c *Config) ExtBool(keys string, defaultVal... interface{}) (bool) {
	return c.Ext(keys, defaultVal...).(bool)
}

func (c *Config) ExtFloat64(keys string, defaultVal... interface{}) (float64) {
	return c.Ext(keys, defaultVal...).(float64)
}

func (c *Config) ExtFloat32(keys string, defaultVal... interface{}) (float32) {
	return c.Ext(keys, defaultVal...).(float32)
}

// Ext will return the value of the EXT config, the keys is separated
// by the given sep string.
func (c *Config) ExtSep(keys, sep string) (interface{}, error) {
	ks := strings.Split(keys, sep);
	var result interface{}
	var isFinal, success bool
	result = c.EXT
	for _, k := range ks {
		result, isFinal, success = find(result, k)
		if !success {
			return "", fmt.Errorf("no such key: %v", k)
		} else if isFinal {
			break
		}
	}

	if success {
		return result, nil
	} else {
		return "", fmt.Errorf("not found")
	}
}

func find(v interface{}, key interface{}) (result interface{}, isFinal, success bool) {
	switch m := v.(type) {
	case map[string]interface{}:
		result, success = m[key.(string)]
		isFinal = (reflect.TypeOf(result) != nil && reflect.TypeOf(result).Kind() != reflect.Map)
	case map[interface{}]interface{}:
		result, success = m[key]
		isFinal = (reflect.TypeOf(result) != nil && reflect.TypeOf(result).Kind() != reflect.Map)
	}
	return
}

var (
	Conf = Config{}
)

func LoadConf(path string) {
	c, err := ioutil.ReadFile(path)
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		return
	}
	yaml.Unmarshal(c, &Conf)
}

func Ext(keys string, defVal... interface{}) interface {}{
	return Conf.Ext(keys, defVal...)
}

// ExtString is a shorcut for (c *Config) ExtString()
func ExtString(keys string, defVal... interface{}) string {
	return Conf.ExtString(keys, defVal...)
}

// ExtBool is a shorcut for (c *Config) ExtBool()
func ExtBool(keys string, defVal... interface{}) bool {
	return Conf.ExtBool(keys, defVal...)
}

// ExtInt is a shorcut for (c *Config) ExtInt()
func ExtInt(keys string, defVal... interface{}) int {
	return Conf.ExtInt(keys, defVal...)
}

// ExtInt8 is a shorcut for (c *Config) ExtInt8()
func ExtInt8(keys string, defVal... interface{}) int8 {
	return Conf.ExtInt8(keys, defVal...)
}

// ExtInt16 is a shorcut for (c *Config) ExtInt16()
func ExtInt16(keys string, defVal... interface{}) int16 {
	return Conf.ExtInt16(keys, defVal...)
}

// ExtInt32 is a shorcut for (c *Config) ExtInt32()
func ExtInt32(keys string, defVal... interface{}) int32 {
	return Conf.ExtInt32(keys, defVal...)
}

// ExtInt64 is a shorcut for (c *Config) ExtInt64()
func ExtInt64(keys string, defVal... interface{}) int64 {
	return Conf.ExtInt64(keys, defVal...)
}

// ExtFloat32 is a shorcut for (c *Config) ExtFloat32()
func ExtFloat32(keys string, defVal... interface{}) float32 {
	return Conf.ExtFloat32(keys, defVal...)
}

// ExtFloat64 is a shorcut for (c *Config) ExtFloat64()
func ExtFloat64(keys string, defVal... interface{}) float64 {
	return Conf.ExtFloat64(keys, defVal...)
}

func init() {
	LoadConf("conf.yml")
}
<|endoftext|>


<commit_before>package ndp

import (
	"net"
	"sync"
	"testing"

	"github.com/mdlayher/ndp/internal/ndptest"

	"github.com/google/go-cmp/cmp"
)

func TestConn(t *testing.T) {
	tests := []struct {
		name string
		fn   func(t *testing.T, c1, c2 *Conn, addr net.IP)
	}{
		{
			name: "echo",
			fn:   testConnEcho,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			withConns(t, tt.fn)
		})
	}
}

func testConnEcho(t *testing.T, c1, c2 *Conn, addr net.IP) {
	// Echo this message between two connections.
	rs := &RouterSolicitation{}

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()

		// Read and bounce the message back to the second Conn.
		m, _, _, err := c2.ReadFrom()
		if err != nil {
			panicf("failed to read from c2: %v", err)
		}

		if err := c2.WriteTo(m, nil, addr); err != nil {
			panicf("failed to write from c2: %v", err)
		}
	}()

	if err := c1.WriteTo(rs, nil, addr); err != nil {
		t.Fatalf("failed to write from c1: %v", err)
	}

	m, _, _, err := c1.ReadFrom()
	if err != nil {
		t.Fatalf("failed to read from c1: %v", err)
	}

	wg.Wait()

	if diff := cmp.Diff(rs, m); diff != "" {
		t.Fatalf("unexpected message (-want +got):\n%s", diff)
	}
}

// withConns invokes fn once with a UDPv6 connection and again with an ICMPv6
// connection, enabling testing with both privileged and unprivileged sockets.
func withConns(t *testing.T, fn func(t *testing.T, c1, c2 *Conn, addr net.IP)) {
	var name string
	var newConn func(t *testing.T) (*Conn, *Conn, net.IP, func())

	for i := 0; i < 2; i++ {
		switch i {
		case 0:
			name = "UDP"
			newConn = testUDPConn
		case 1:
			name = "ICMP"
			newConn = testICMPConn
		default:
			t.Fatalf("unhandled withConns iteration: %d", i)
		}

		t.Run(name, func(t *testing.T) {
			c1, c2, addr, done := newConn(t)
			defer done()

			fn(t, c1, c2, addr)
		})
	}
}

func TestSolicitedNodeMulticast(t *testing.T) {
	tests := []struct {
		name string
		ip   net.IP
		snm  net.IP
		ok   bool
	}{
		{
			name: "bad, IPv4",
			ip:   net.IPv4(192, 168, 1, 1),
		},
		{
			name: "ok, link-local",
			ip:   ndptest.MustIPv6("fe80::1234:5678"),
			snm:  ndptest.MustIPv6("ff02::1:ff34:5678"),
			ok:   true,
		},
		{
			name: "ok, global",
			ip:   ndptest.MustIPv6("2001:db8::dead:beef"),
			snm:  ndptest.MustIPv6("ff02::1:ffad:beef"),
			ok:   true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			snm, err := SolicitedNodeMulticast(tt.ip)

			if err != nil && tt.ok {
				t.Fatalf("unexpected error: %v", err)
			}
			if err == nil && !tt.ok {
				t.Fatal("expected an error, but none occurred")
			}
			if err != nil {
				t.Logf("OK error: %v", err)
				return
			}

			if diff := cmp.Diff(tt.snm, snm); diff != "" {
				t.Fatalf("unexpected solicited-node multicast address (-want +got):\n%s", diff)
			}
		})
	}
}
<commit_msg>ndp: tidy imports<commit_after>package ndp

import (
	"net"
	"sync"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/mdlayher/ndp/internal/ndptest"
)

func TestConn(t *testing.T) {
	tests := []struct {
		name string
		fn   func(t *testing.T, c1, c2 *Conn, addr net.IP)
	}{
		{
			name: "echo",
			fn:   testConnEcho,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			withConns(t, tt.fn)
		})
	}
}

func testConnEcho(t *testing.T, c1, c2 *Conn, addr net.IP) {
	// Echo this message between two connections.
	rs := &RouterSolicitation{}

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()

		// Read and bounce the message back to the second Conn.
		m, _, _, err := c2.ReadFrom()
		if err != nil {
			panicf("failed to read from c2: %v", err)
		}

		if err := c2.WriteTo(m, nil, addr); err != nil {
			panicf("failed to write from c2: %v", err)
		}
	}()

	if err := c1.WriteTo(rs, nil, addr); err != nil {
		t.Fatalf("failed to write from c1: %v", err)
	}

	m, _, _, err := c1.ReadFrom()
	if err != nil {
		t.Fatalf("failed to read from c1: %v", err)
	}

	wg.Wait()

	if diff := cmp.Diff(rs, m); diff != "" {
		t.Fatalf("unexpected message (-want +got):\n%s", diff)
	}
}

// withConns invokes fn once with a UDPv6 connection and again with an ICMPv6
// connection, enabling testing with both privileged and unprivileged sockets.
func withConns(t *testing.T, fn func(t *testing.T, c1, c2 *Conn, addr net.IP)) {
	var name string
	var newConn func(t *testing.T) (*Conn, *Conn, net.IP, func())

	for i := 0; i < 2; i++ {
		switch i {
		case 0:
			name = "UDP"
			newConn = testUDPConn
		case 1:
			name = "ICMP"
			newConn = testICMPConn
		default:
			t.Fatalf("unhandled withConns iteration: %d", i)
		}

		t.Run(name, func(t *testing.T) {
			c1, c2, addr, done := newConn(t)
			defer done()

			fn(t, c1, c2, addr)
		})
	}
}

func TestSolicitedNodeMulticast(t *testing.T) {
	tests := []struct {
		name string
		ip   net.IP
		snm  net.IP
		ok   bool
	}{
		{
			name: "bad, IPv4",
			ip:   net.IPv4(192, 168, 1, 1),
		},
		{
			name: "ok, link-local",
			ip:   ndptest.MustIPv6("fe80::1234:5678"),
			snm:  ndptest.MustIPv6("ff02::1:ff34:5678"),
			ok:   true,
		},
		{
			name: "ok, global",
			ip:   ndptest.MustIPv6("2001:db8::dead:beef"),
			snm:  ndptest.MustIPv6("ff02::1:ffad:beef"),
			ok:   true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			snm, err := SolicitedNodeMulticast(tt.ip)

			if err != nil && tt.ok {
				t.Fatalf("unexpected error: %v", err)
			}
			if err == nil && !tt.ok {
				t.Fatal("expected an error, but none occurred")
			}
			if err != nil {
				t.Logf("OK error: %v", err)
				return
			}

			if diff := cmp.Diff(tt.snm, snm); diff != "" {
				t.Fatalf("unexpected solicited-node multicast address (-want +got):\n%s", diff)
			}
		})
	}
}
<|endoftext|>


<commit_before>package adsi

import (
	"io"
	"sync"
	"unsafe"

	"github.com/go-ole/go-ole"
	"gopkg.in/adsi.v0/api"
)

// Container provides access to Active Directory container objects.
type Container struct {
	m     sync.RWMutex
	iface *api.IADsContainer
}

func (c *Container) closed() bool {
	return (c.iface == nil)
}

// Close will release resources consumed by the container. It should be
// called when the container is no longer needed.
func (c *Container) Close() {
	c.m.Lock()
	defer c.m.Unlock()
	if c.closed() {
		return
	}
	run(func() error {
		c.iface.Release()
		return nil
	})
	// FIXME: What happens if the run returns an error?
	c.iface = nil
}

// Children returns an object iterator that provides access to the immediate
// children of the container.
func (c *Container) Children() (iter *ObjectIter, err error) {
	c.m.Lock()
	defer c.m.Unlock()
	if c.closed() {
		return nil, ErrClosed
	}
	err = run(func() error {
		iunknown, err := c.iface.NewEnum()
		if err != nil {
			return err
		}
		defer iunknown.Release()
		idispatch, err := iunknown.QueryInterface(ole.IID_IEnumVariant)
		if err != nil {
			return err
		}
		iface := (*ole.IEnumVARIANT)(unsafe.Pointer(idispatch))
		iter = NewObjectIter(iface)
		return nil
	})
	return
}

// ObjectIter provides an iterator for a set of objects.
type ObjectIter struct {
	m     sync.RWMutex
	iface *ole.IEnumVARIANT
}

// NewObjectIter returns an object iterator that provides access to the objects
// contained in the given enumerator.
func NewObjectIter(enumerator *ole.IEnumVARIANT) *ObjectIter {
	// TODO: Call ADsBuildEnumerator here?
	return &ObjectIter{iface: enumerator}
}

// Next moves the iterator to the next object and returns a pointer to it. If it
// has reached the end of the set it will return io.EOF. It the iterator has
// already been closed it will return ErrClosed.
//
// FIXME: Make sure that io.EOF is being returned as expected. We might have
// to intercept an internal error.
func (iter *ObjectIter) Next() (obj *Object, err error) {
	iter.m.Lock()
	defer iter.m.Unlock()
	if iter.closed() {
		return nil, ErrClosed
	}
	err = run(func() error {
		// See https://msdn.microsoft.com/library/aa705990
		array, length, err := iter.iface.Next(1)
		if err != nil {
			return err
		}
		defer array.Clear()
		if length == 0 {
			return io.EOF
		}

		idispatch := array.ToIDispatch()
		if idispatch == nil {
			return ErrNonDispatchVariant
		}
		defer idispatch.Release()

		iresult, err := idispatch.QueryInterface(api.IID_IADs)
		if err != nil {
			return err
		}
		iface := (*api.IADs)(unsafe.Pointer(iresult))
		obj = &Object{iface: iface}
		return nil
	})
	return
}

func (iter *ObjectIter) closed() bool {
	return (iter.iface == nil)
}

// Close will release resources consumed by the iterator. It should be
// called when the iterator is no longer needed.
func (iter *ObjectIter) Close() {
	iter.m.Lock()
	defer iter.m.Unlock()
	if iter.closed() {
		return
	}
	run(func() error {
		iter.iface.Release()
		return nil
	})
	// FIXME: What happens if the run returns an error?
	iter.iface = nil
}
<commit_msg>Fixed critical double-release bug in ObjectIter<commit_after>package adsi

import (
	"io"
	"sync"
	"unsafe"

	"github.com/go-ole/go-ole"
	"gopkg.in/adsi.v0/api"
)

// Container provides access to Active Directory container objects.
type Container struct {
	m     sync.RWMutex
	iface *api.IADsContainer
}

func (c *Container) closed() bool {
	return (c.iface == nil)
}

// Close will release resources consumed by the container. It should be
// called when the container is no longer needed.
func (c *Container) Close() {
	c.m.Lock()
	defer c.m.Unlock()
	if c.closed() {
		return
	}
	run(func() error {
		c.iface.Release()
		return nil
	})
	// FIXME: What happens if the run returns an error?
	c.iface = nil
}

// Children returns an object iterator that provides access to the immediate
// children of the container.
func (c *Container) Children() (iter *ObjectIter, err error) {
	c.m.Lock()
	defer c.m.Unlock()
	if c.closed() {
		return nil, ErrClosed
	}
	err = run(func() error {
		iunknown, err := c.iface.NewEnum()
		if err != nil {
			return err
		}
		defer iunknown.Release()
		idispatch, err := iunknown.QueryInterface(ole.IID_IEnumVariant)
		if err != nil {
			return err
		}
		iface := (*ole.IEnumVARIANT)(unsafe.Pointer(idispatch))
		iter = NewObjectIter(iface)
		return nil
	})
	return
}

// ObjectIter provides an iterator for a set of objects.
type ObjectIter struct {
	m     sync.RWMutex
	iface *ole.IEnumVARIANT
}

// NewObjectIter returns an object iterator that provides access to the objects
// contained in the given enumerator.
func NewObjectIter(enumerator *ole.IEnumVARIANT) *ObjectIter {
	// TODO: Call ADsBuildEnumerator here?
	return &ObjectIter{iface: enumerator}
}

// Next moves the iterator to the next object and returns a pointer to it. If it
// has reached the end of the set it will return io.EOF. It the iterator has
// already been closed it will return ErrClosed.
//
// FIXME: Make sure that io.EOF is being returned as expected. We might have
// to intercept an internal error.
func (iter *ObjectIter) Next() (obj *Object, err error) {
	iter.m.Lock()
	defer iter.m.Unlock()
	if iter.closed() {
		return nil, ErrClosed
	}
	err = run(func() error {
		// See https://msdn.microsoft.com/library/aa705990
		array, length, err := iter.iface.Next(1)
		if err != nil {
			return err
		}
		defer array.Clear()
		if length == 0 {
			return io.EOF
		}

		idispatch := array.ToIDispatch()
		if idispatch == nil {
			return ErrNonDispatchVariant
		}
		// Note: Do *not* call idispatch.Release() here, as it will be called
		//       automatically by array.Clear()

		iresult, err := idispatch.QueryInterface(api.IID_IADs)
		if err != nil {
			return err
		}
		iface := (*api.IADs)(unsafe.Pointer(iresult))
		obj = &Object{iface: iface}
		return nil
	})
	return
}

func (iter *ObjectIter) closed() bool {
	return (iter.iface == nil)
}

// Close will release resources consumed by the iterator. It should be
// called when the iterator is no longer needed.
func (iter *ObjectIter) Close() {
	iter.m.Lock()
	defer iter.m.Unlock()
	if iter.closed() {
		return
	}
	run(func() error {
		iter.iface.Release()
		return nil
	})
	// FIXME: What happens if the run returns an error?
	iter.iface = nil
}
<|endoftext|>


<commit_before>package handlers

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestDefaultCORSHandlerReturnsOk(t *testing.T) {
	methods := []string{"GET", "HEAD", "POST"}

	for _, method := range methods {
		r := newRequest(method, "http://www.example.com/")
		rr := httptest.NewRecorder()

		testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

		CORS()(testHandler).ServeHTTP(rr, r)

		if status := rr.Code; status != http.StatusOK {
			t.Fatalf("bad status: got %v want %v for method %s", status, http.StatusFound, method)
		}
	}
}

func TestCORSHandlerIgnoreOptionsFallsThrough(t *testing.T) {
	r := newRequest("OPTIONS", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusTeapot)
	})

	CORS(IgnoreOptions())(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusTeapot {
		t.Fatalf("bad status: got %v want %v", status, http.StatusTeapot)
	}
}

func TestCORSHandlerSetsExposedHeaders(t *testing.T) {
	methods := []string{"GET", "HEAD", "POST"}

	for _, method := range methods {
		// Test default configuration.
		r := newRequest(method, "http://www.example.com/")
		r.Header.Set("Origin", r.URL.String())

		rr := httptest.NewRecorder()

		testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

		CORS(ExposedHeaders([]string{"X-CORS-TEST"}))(testHandler).ServeHTTP(rr, r)

		if status := rr.Code; status != http.StatusOK {
			t.Fatalf("bad status: got %v want %v for method %s", status, http.StatusOK, method)
		}

		header := rr.HeaderMap.Get(corsExposeHeadersHeader)
		if header != "X-Cors-Test" {
			t.Fatalf("bad header: expected X-Cors-Test header, got empty header for method %s.", method)
		}
	}
}

func TestCORSHandlerUnsetRequethMethodForPreflightBadRequest(t *testing.T) {
	r := newRequest("OPTIONS", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(AllowedMethods([]string{"DELETE"}))(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusBadRequest {
		t.Fatalf("bad status: got %v want %v", status, http.StatusBadRequest)
	}
}

func TestCORSHandlerAllowedMethodForPreflight(t *testing.T) {
	r := newRequest("OPTIONS", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())
	r.Header.Set(corsRequestMethodHeader, "DELETE")

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(AllowedMethods([]string{"DELETE"}))(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
	}

	header := rr.HeaderMap.Get(corsAllowMethodsHeader)
	if header != "DELETE" {
		t.Fatalf("bad header: expected DELETE method header, got empty header.")
	}
}

func TestCORSHandlerAllowedHeaderForPreflight(t *testing.T) {
	r := newRequest("OPTIONS", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())
	r.Header.Set(corsRequestMethodHeader, "POST")
	r.Header.Set(corsRequestHeadersHeader, "Content-Type")

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(AllowedHeaders([]string{"Content-Type"}))(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
	}

	header := rr.HeaderMap.Get(corsAllowHeadersHeader)
	if header != "Content-Type" {
		t.Fatalf("bad header: expected Content-Type header, got empty header.")
	}
}

func TestCORSHandlerMaxAgeForPreflight(t *testing.T) {
	r := newRequest("OPTIONS", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())
	r.Header.Set(corsRequestMethodHeader, "POST")

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(MaxAge(3500))(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
	}

	header := rr.HeaderMap.Get(corsMaxAgeHeader)
	if header != "600" {
		t.Fatalf("bad header: expected %s to be %s, got %s.", corsMaxAgeHeader, "600", header)
	}
}

func TestCORSHandlerAllowedCredentials(t *testing.T) {
	r := newRequest("GET", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(AllowCredentials())(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
	}

	header := rr.HeaderMap.Get(corsAllowCredentialsHeader)
	if header != "true" {
		t.Fatalf("bad header: expected %s to be %s, got %s.", corsAllowCredentialsHeader, "true", header)
	}
}

func TestCORSHandlerMultipleAllowOriginsSetsVaryHeader(t *testing.T) {
	r := newRequest("GET", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(AllowedOrigins([]string{r.URL.String(), "http://google.com"}))(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
	}

	header := rr.HeaderMap.Get(corsVaryHeader)
	if header != corsOriginHeader {
		t.Fatalf("bad header: expected %s to be %s, got %s.", corsVaryHeader, corsOriginHeader, header)
	}
}
<commit_msg>Additional Tests<commit_after>package handlers

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestDefaultCORSHandlerReturnsOk(t *testing.T) {
	r := newRequest("GET", "http://www.example.com/")
	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS()(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatalf("bad status: got %v want %v", status, http.StatusFound)
	}
}

func TestCORSHandlerIgnoreOptionsFallsThrough(t *testing.T) {
	r := newRequest("OPTIONS", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusTeapot)
	})

	CORS(IgnoreOptions())(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusTeapot {
		t.Fatalf("bad status: got %v want %v", status, http.StatusTeapot)
	}
}

func TestCORSHandlerSetsExposedHeaders(t *testing.T) {
	// Test default configuration.
	r := newRequest("GET", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(ExposedHeaders([]string{"X-CORS-TEST"}))(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
	}

	header := rr.HeaderMap.Get(corsExposeHeadersHeader)
	if header != "X-Cors-Test" {
		t.Fatal("bad header: expected X-Cors-Test header, got empty header for method.")
	}
}

func TestCORSHandlerUnsetRequethMethodForPreflightBadRequest(t *testing.T) {
	r := newRequest("OPTIONS", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(AllowedMethods([]string{"DELETE"}))(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusBadRequest {
		t.Fatalf("bad status: got %v want %v", status, http.StatusBadRequest)
	}
}

func TestCORSHandlerAllowedMethodForPreflight(t *testing.T) {
	r := newRequest("OPTIONS", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())
	r.Header.Set(corsRequestMethodHeader, "DELETE")

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(AllowedMethods([]string{"DELETE"}))(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
	}

	header := rr.HeaderMap.Get(corsAllowMethodsHeader)
	if header != "DELETE" {
		t.Fatalf("bad header: expected DELETE method header, got empty header.")
	}
}

func TestCORSHandlerAllowMethodsNotSetForSimpleRequestPreflight(t *testing.T) {
	for _, method := range defaultCorsMethods {
		r := newRequest("OPTIONS", "http://www.example.com/")
		r.Header.Set("Origin", r.URL.String())
		r.Header.Set(corsRequestMethodHeader, method)

		rr := httptest.NewRecorder()

		testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

		CORS()(testHandler).ServeHTTP(rr, r)

		if status := rr.Code; status != http.StatusOK {
			t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
		}

		header := rr.HeaderMap.Get(corsAllowMethodsHeader)
		if header != "" {
			t.Fatalf("bad header: expected empty method header, got %s.", header)
		}
	}
}

func TestCORSHandlerAllowedHeaderNotSetForSimpleRequestPreflight(t *testing.T) {
	for _, simpleHeader := range defaultCorsHeaders {
		r := newRequest("OPTIONS", "http://www.example.com/")
		r.Header.Set("Origin", r.URL.String())
		r.Header.Set(corsRequestMethodHeader, "GET")
		r.Header.Set(corsRequestHeadersHeader, simpleHeader)

		rr := httptest.NewRecorder()

		testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

		CORS()(testHandler).ServeHTTP(rr, r)

		if status := rr.Code; status != http.StatusOK {
			t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
		}

		header := rr.HeaderMap.Get(corsAllowHeadersHeader)
		if header != "" {
			t.Fatalf("bad header: expected empty header, got %s.", header)
		}
	}
}

func TestCORSHandlerAllowedHeaderForPreflight(t *testing.T) {
	r := newRequest("OPTIONS", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())
	r.Header.Set(corsRequestMethodHeader, "POST")
	r.Header.Set(corsRequestHeadersHeader, "Content-Type")

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(AllowedHeaders([]string{"Content-Type"}))(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
	}

	header := rr.HeaderMap.Get(corsAllowHeadersHeader)
	if header != "Content-Type" {
		t.Fatalf("bad header: expected Content-Type header, got empty header.")
	}
}

func TestCORSHandlerMaxAgeForPreflight(t *testing.T) {
	r := newRequest("OPTIONS", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())
	r.Header.Set(corsRequestMethodHeader, "POST")

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(MaxAge(3500))(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
	}

	header := rr.HeaderMap.Get(corsMaxAgeHeader)
	if header != "600" {
		t.Fatalf("bad header: expected %s to be %s, got %s.", corsMaxAgeHeader, "600", header)
	}
}

func TestCORSHandlerAllowedCredentials(t *testing.T) {
	r := newRequest("GET", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(AllowCredentials())(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
	}

	header := rr.HeaderMap.Get(corsAllowCredentialsHeader)
	if header != "true" {
		t.Fatalf("bad header: expected %s to be %s, got %s.", corsAllowCredentialsHeader, "true", header)
	}
}

func TestCORSHandlerMultipleAllowOriginsSetsVaryHeader(t *testing.T) {
	r := newRequest("GET", "http://www.example.com/")
	r.Header.Set("Origin", r.URL.String())

	rr := httptest.NewRecorder()

	testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	CORS(AllowedOrigins([]string{r.URL.String(), "http://google.com"}))(testHandler).ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatalf("bad status: got %v want %v", status, http.StatusOK)
	}

	header := rr.HeaderMap.Get(corsVaryHeader)
	if header != corsOriginHeader {
		t.Fatalf("bad header: expected %s to be %s, got %s.", corsVaryHeader, corsOriginHeader, header)
	}
}
<|endoftext|>


<commit_before>package creds

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strings"
)

const (
	mfaCodeRegexString = `^\d{6}$`
)

var mfaCodeRegex = regexp.MustCompile(mfaCodeRegexString)

// MfaPrompt defines an object which recieves an Mfa ARN and returns an Mfa code
type MfaPrompt interface {
	Prompt(string) (string, error)
}

// WritableMfaPrompt defines an MFA Prompt which can store a new secret
type WritableMfaPrompt interface {
	Store(string, string) error
	RetryText(string) string
}

// revive:disable-next-line:flag-parameter
func (c Creds) handleMfa(useMfa bool, mfaCode string, mfaPrompt MfaPrompt) (string, string, error) {
	logger.InfoMsg("handling mfa options")

	if !useMfa && mfaCode == "" {
		logger.InfoMsg("mfa is disabled")
		return "", "", nil
	}
	mfaSerial, err := c.MfaArn()
	if err != nil {
		return "", "", err
	}
	if mfaCode != "" {
		logger.InfoMsg("mfa code already provided")
		err := validateMfa(mfaCode)
		if err != nil {
			return "", "", err
		}

		return mfaCode, mfaSerial, nil
	}
	if mfaPrompt == nil {
		logger.InfoMsg("using default mfa prompt")
		// revive:disable-next-line:modifies-parameter
		mfaPrompt = &DefaultMfaPrompt{}
	}
	logger.InfoMsg("prompting for mfa code")
	// revive:disable-next-line:modifies-parameter
	mfaCode, err = mfaPrompt.Prompt(mfaSerial)
	if err != nil {
		return "", "", err
	}
	return mfaCode, mfaSerial, nil
}

// DefaultMfaPrompt defines the standard CLI-based MFA prompt
type DefaultMfaPrompt struct {
	PromptTemplate string
}

var defaultPromptTemplate = "MFA Code: "

// Prompt asks the user for their MFA token
func (p *DefaultMfaPrompt) Prompt(arn string) (string, error) {
	mfaReader := bufio.NewReader(os.Stdin)
	pt := p.PromptTemplate
	if pt == "" {
		pt = defaultPromptTemplate
	}
	fmt.Fprintf(os.Stderr, pt, arn)
	mfa, err := mfaReader.ReadString('\n')
	if err != nil {
		return "", err
	}
	trimmed := strings.TrimSpace(mfa)
	err = validateMfa(trimmed)
	if err != nil {
		return "", err
	}
	return trimmed, nil
}

func validateMfa(code string) error {
	logger.InfoMsg("validating mfa")
	if mfaCodeRegex.MatchString(code) {
		return nil
	}
	return fmt.Errorf("provided mfa code does not match the necessary format")
}

// MultiMfaPrompt allows a slice of sequential backends to check for Mfa
type MultiMfaPrompt struct {
	Backends []MfaPrompt
}

// Prompt iterates through the backends to find an Mfa code
func (m *MultiMfaPrompt) Prompt(arn string) (string, error) {
	logger.InfoMsgf("looking up %s in multi mfa prompt", arn)
	for index, item := range m.Backends {
		res, err := item.Prompt(arn)
		if err == nil {
			return res, nil
		}
		logger.InfoMsgf("failed backend lookup %d with %s", index, err)
	}
	return "", fmt.Errorf("all backends failed to return mfa code")
}

// Store attempts to store the Mfa in a backend
func (m *MultiMfaPrompt) Store(arn, seed string) error {
	logger.InfoMsgf("attempting to store %s in multi mfa object", arn)
	writer := m.getWriter()
	if writer == nil {
		return fmt.Errorf("no writable mfa backends found")
	}
	return writer.Store(arn, seed)
}

// RetryText returns helper text for retrying a failed mfa storage
func (m *MultiMfaPrompt) RetryText(arn string) string {
	writer := m.getWriter()
	if writer == nil {
		return ""
	}
	return writer.RetryText(arn)
}

func (m *MultiMfaPrompt) getWriter() WritableMfaPrompt {
	for index, item := range m.Backends {
		writer, ok := item.(WritableMfaPrompt)
		if ok {
			logger.InfoMsgf("found mfa writer with index %d", index)
			return writer
		}
	}
	return nil
}
<commit_msg>fix to use func<commit_after>package creds

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strings"
)

const (
	mfaCodeRegexString = `^\d{6}$`
)

var mfaCodeRegex = regexp.MustCompile(mfaCodeRegexString)

// MfaPrompt defines an object which recieves an Mfa ARN and returns an Mfa code
type MfaPrompt interface {
	Prompt(string) (string, error)
}

// WritableMfaPrompt defines an MFA Prompt which can store a new secret
type WritableMfaPrompt interface {
	Store(string, string) error
	RetryText(string) string
}

// revive:disable-next-line:flag-parameter
func (c Creds) handleMfa(useMfa bool, mfaCode string, mfaPrompt MfaPrompt) (string, string, error) {
	logger.InfoMsg("handling mfa options")

	if !useMfa && mfaCode == "" {
		logger.InfoMsg("mfa is disabled")
		return "", "", nil
	}
	mfaSerial, err := c.MfaArn()
	if err != nil {
		return "", "", err
	}
	if mfaCode != "" {
		logger.InfoMsg("mfa code already provided")
		err := validateMfa(mfaCode)
		if err != nil {
			return "", "", err
		}

		return mfaCode, mfaSerial, nil
	}
	if mfaPrompt == nil {
		logger.InfoMsg("using default mfa prompt")
		// revive:disable-next-line:modifies-parameter
		mfaPrompt = &DefaultMfaPrompt{}
	}
	logger.InfoMsg("prompting for mfa code")
	// revive:disable-next-line:modifies-parameter
	mfaCode, err = mfaPrompt.Prompt(mfaSerial)
	if err != nil {
		return "", "", err
	}
	return mfaCode, mfaSerial, nil
}

// DefaultMfaPrompt defines the standard CLI-based MFA prompt
type DefaultMfaPrompt struct {
	PromptTextFunc func(string) string
}

func defaultPromptTextFunc(_ string) string {
	return "MFA Code: "
}

// Prompt asks the user for their MFA token
func (p *DefaultMfaPrompt) Prompt(arn string) (string, error) {
	mfaReader := bufio.NewReader(os.Stdin)
	pf := PromptTextFunc
	if pf == nil {
		pf = defaultPromptTextFunc
	}
	fmt.Fprint(os.Stderr, pf(arn))
	mfa, err := mfaReader.ReadString('\n')
	if err != nil {
		return "", err
	}
	trimmed := strings.TrimSpace(mfa)
	err = validateMfa(trimmed)
	if err != nil {
		return "", err
	}
	return trimmed, nil
}

func validateMfa(code string) error {
	logger.InfoMsg("validating mfa")
	if mfaCodeRegex.MatchString(code) {
		return nil
	}
	return fmt.Errorf("provided mfa code does not match the necessary format")
}

// MultiMfaPrompt allows a slice of sequential backends to check for Mfa
type MultiMfaPrompt struct {
	Backends []MfaPrompt
}

// Prompt iterates through the backends to find an Mfa code
func (m *MultiMfaPrompt) Prompt(arn string) (string, error) {
	logger.InfoMsgf("looking up %s in multi mfa prompt", arn)
	for index, item := range m.Backends {
		res, err := item.Prompt(arn)
		if err == nil {
			return res, nil
		}
		logger.InfoMsgf("failed backend lookup %d with %s", index, err)
	}
	return "", fmt.Errorf("all backends failed to return mfa code")
}

// Store attempts to store the Mfa in a backend
func (m *MultiMfaPrompt) Store(arn, seed string) error {
	logger.InfoMsgf("attempting to store %s in multi mfa object", arn)
	writer := m.getWriter()
	if writer == nil {
		return fmt.Errorf("no writable mfa backends found")
	}
	return writer.Store(arn, seed)
}

// RetryText returns helper text for retrying a failed mfa storage
func (m *MultiMfaPrompt) RetryText(arn string) string {
	writer := m.getWriter()
	if writer == nil {
		return ""
	}
	return writer.RetryText(arn)
}

func (m *MultiMfaPrompt) getWriter() WritableMfaPrompt {
	for index, item := range m.Backends {
		writer, ok := item.(WritableMfaPrompt)
		if ok {
			logger.InfoMsgf("found mfa writer with index %d", index)
			return writer
		}
	}
	return nil
}
<|endoftext|>