<commit_before>package pgp

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// Test interface
var (
	_ Signer   = &GpgSigner{}
	_ Verifier = &GpgVerifier{}
)

// GpgSigner is implementation of Signer interface using gpg as external program
type GpgSigner struct {
	keyRef                     string
	keyring, secretKeyring     string
	passphrase, passphraseFile string
	batch                      bool
}

// SetBatch control --no-tty flag to gpg
func (g *GpgSigner) SetBatch(batch bool) {
	g.batch = batch
}

// SetKey sets key ID to use when signing files
func (g *GpgSigner) SetKey(keyRef string) {
	g.keyRef = keyRef
}

// SetKeyRing allows to set custom keyring and secretkeyring
func (g *GpgSigner) SetKeyRing(keyring, secretKeyring string) {
	g.keyring, g.secretKeyring = keyring, secretKeyring
}

// SetPassphrase sets passhprase params
func (g *GpgSigner) SetPassphrase(passphrase, passphraseFile string) {
	g.passphrase, g.passphraseFile = passphrase, passphraseFile
}

func (g *GpgSigner) gpgArgs() []string {
	args := []string{}
	if g.keyring != "" {
		args = append(args, "--no-auto-check-trustdb", "--no-default-keyring", "--keyring", g.keyring)
	}
	if g.secretKeyring != "" {
		args = append(args, "--secret-keyring", g.secretKeyring)
	}

	if g.keyRef != "" {
		args = append(args, "-u", g.keyRef)
	}

	if g.passphrase != "" || g.passphraseFile != "" {
		args = append(args, "--no-use-agent")
	}

	if g.passphrase != "" {
		args = append(args, "--passphrase", g.passphrase)
	}

	if g.passphraseFile != "" {
		args = append(args, "--passphrase-file", g.passphraseFile)
	}

	if g.batch {
		args = append(args, "--no-tty")
	}

	return args
}

// Init verifies availability of gpg & presence of keys
func (g *GpgSigner) Init() error {
	output, err := exec.Command("gpg", "--list-keys", "--dry-run", "--no-auto-check-trustdb").CombinedOutput()
	if err != nil {
		return fmt.Errorf("unable to execute gpg: %s (is gpg installed?): %s", err, string(output))
	}

	if g.keyring == "" && g.secretKeyring == "" && len(output) == 0 {
		return fmt.Errorf("looks like there are no keys in gpg, please create one (official manual: http://www.gnupg.org/gph/en/manual.html)")
	}

	return err
}

// DetachedSign signs file with detached signature in ASCII format
func (g *GpgSigner) DetachedSign(source string, destination string) error {
	fmt.Printf("Signing file '%s' with gpg, please enter your passphrase when prompted:\n", filepath.Base(source))

	args := []string{"-o", destination, "--digest-algo", "SHA256", "--armor", "--yes"}
	args = append(args, g.gpgArgs()...)
	args = append(args, "--detach-sign", source)
	cmd := exec.Command("gpg", args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// ClearSign clear-signs the file
func (g *GpgSigner) ClearSign(source string, destination string) error {
	fmt.Printf("Clearsigning file '%s' with gpg, please enter your passphrase when prompted:\n", filepath.Base(source))
	args := []string{"-o", destination, "--digest-algo", "SHA256", "--yes"}
	args = append(args, g.gpgArgs()...)
	args = append(args, "--clearsign", source)
	cmd := exec.Command("gpg", args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// GpgVerifier is implementation of Verifier interface using gpgv as external program
type GpgVerifier struct {
	keyRings []string
}

// InitKeyring verifies that gpg is installed and some keys are trusted
func (g *GpgVerifier) InitKeyring() error {
	err := exec.Command("gpgv", "--version").Run()
	if err != nil {
		return fmt.Errorf("unable to execute gpgv: %s (is gpg installed?)", err)
	}

	if len(g.keyRings) == 0 {
		// using default keyring
		output, err := exec.Command("gpg", "--no-default-keyring", "--no-auto-check-trustdb", "--keyring", "trustedkeys.gpg", "--list-keys").Output()
		if err == nil && len(output) == 0 {
			fmt.Printf("\nLooks like your keyring with trusted keys is empty. You might consider importing some keys.\n")
			fmt.Printf("If you're running Debian or Ubuntu, it's a good idea to import current archive keys by running:\n\n")
			fmt.Printf("  gpg --no-default-keyring --keyring /usr/share/keyrings/debian-archive-keyring.gpg --export | gpg --no-default-keyring --keyring trustedkeys.gpg --import\n")
			fmt.Printf("\n(for Ubuntu, use /usr/share/keyrings/ubuntu-archive-keyring.gpg)\n\n")
		}
	}

	return nil
}

// AddKeyring adds custom keyring to GPG parameters
func (g *GpgVerifier) AddKeyring(keyring string) {
	g.keyRings = append(g.keyRings, keyring)
}

func (g *GpgVerifier) argsKeyrings() (args []string) {
	if len(g.keyRings) > 0 {
		args = make([]string, 0, 2*len(g.keyRings))
		for _, keyring := range g.keyRings {
			args = append(args, "--keyring", keyring)
		}
	} else {
		args = []string{"--keyring", "trustedkeys.gpg"}
	}
	return
}

func (g *GpgVerifier) runGpgv(args []string, context string, showKeyTip bool) (*KeyInfo, error) {
	args = append([]string{"--status-fd", "3"}, args...)
	cmd := exec.Command("gpgv", args...)

	tempf, err := ioutil.TempFile("", "aptly-gpg-status")
	if err != nil {
		return nil, err
	}
	defer tempf.Close()

	err = os.Remove(tempf.Name())
	if err != nil {
		return nil, err
	}

	cmd.ExtraFiles = []*os.File{tempf}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, err
	}
	defer stderr.Close()

	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	buffer := &bytes.Buffer{}

	_, err = io.Copy(io.MultiWriter(os.Stderr, buffer), stderr)
	if err != nil {
		return nil, err
	}

	cmderr := cmd.Wait()

	tempf.Seek(0, 0)

	statusr := bufio.NewScanner(tempf)

	result := &KeyInfo{}

	for statusr.Scan() {
		line := strings.TrimSpace(statusr.Text())

		if strings.HasPrefix(line, "[GNUPG:] GOODSIG ") {
			result.GoodKeys = append(result.GoodKeys, Key(strings.Fields(line)[2]))
		} else if strings.HasPrefix(line, "[GNUPG:] NO_PUBKEY ") {
			result.MissingKeys = append(result.MissingKeys, Key(strings.Fields(line)[2]))
		}
	}

	if err = statusr.Err(); err != nil {
		return nil, err
	}

	if cmderr != nil {
		if showKeyTip && len(g.keyRings) == 0 && len(result.MissingKeys) > 0 {
			fmt.Printf("\nLooks like some keys are missing in your trusted keyring, you may consider importing them from keyserver:\n\n")

			keys := make([]string, len(result.MissingKeys))

			for i := range result.MissingKeys {
				keys[i] = string(result.MissingKeys[i])
			}

			fmt.Printf("gpg --no-default-keyring --keyring trustedkeys.gpg --keyserver keys.gnupg.net --recv-keys %s\n\n",
				strings.Join(keys, " "))

			fmt.Printf("Sometimes keys are stored in repository root in file named Release.key, to import such key:\n\n")
			fmt.Printf("wget -O - https://some.repo/repository/Release.key | gpg --no-default-keyring --keyring trustedkeys.gpg --import\n\n")
		}
		return result, fmt.Errorf("verification of %s failed: %s", context, cmderr)
	}
	return result, nil
}

// VerifyDetachedSignature verifies combination of signature and cleartext using gpgv
func (g *GpgVerifier) VerifyDetachedSignature(signature, cleartext io.Reader, showKeyTip bool) error {
	args := g.argsKeyrings()

	sigf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return err
	}
	defer os.Remove(sigf.Name())
	defer sigf.Close()

	_, err = io.Copy(sigf, signature)
	if err != nil {
		return err
	}

	clearf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return err
	}
	defer os.Remove(clearf.Name())
	defer clearf.Close()

	_, err = io.Copy(clearf, cleartext)
	if err != nil {
		return err
	}

	args = append(args, sigf.Name(), clearf.Name())
	_, err = g.runGpgv(args, "detached signature", showKeyTip)
	return err
}

// IsClearSigned returns true if file contains signature
func (g *GpgVerifier) IsClearSigned(clearsigned io.Reader) (bool, error) {
	scanner := bufio.NewScanner(clearsigned)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), "BEGIN PGP SIGN") {
			return true, nil
		}
	}

	return false, scanner.Err()
}

// VerifyClearsigned verifies clearsigned file using gpgv
func (g *GpgVerifier) VerifyClearsigned(clearsigned io.Reader, showKeyTip bool) (*KeyInfo, error) {
	args := g.argsKeyrings()

	clearf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return nil, err
	}
	defer os.Remove(clearf.Name())
	defer clearf.Close()

	_, err = io.Copy(clearf, clearsigned)
	if err != nil {
		return nil, err
	}

	args = append(args, clearf.Name())
	return g.runGpgv(args, "clearsigned file", showKeyTip)
}

// ExtractClearsigned extracts cleartext from clearsigned file WITHOUT signature verification
func (g *GpgVerifier) ExtractClearsigned(clearsigned io.Reader) (text *os.File, err error) {
	clearf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return
	}
	defer os.Remove(clearf.Name())
	defer clearf.Close()

	_, err = io.Copy(clearf, clearsigned)
	if err != nil {
		return
	}

	text, err = ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return
	}
	defer os.Remove(text.Name())

	args := []string{"--no-auto-check-trustdb", "--decrypt", "--batch", "--skip-verify", "--output", "-", clearf.Name()}

	cmd := exec.Command("gpg", args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	defer stdout.Close()

	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	_, err = io.Copy(text, stdout)
	if err != nil {
		return nil, err
	}

	err = cmd.Wait()

	if err != nil {
		return nil, fmt.Errorf("extraction of clearsigned file failed: %s", err)
	}

	_, err = text.Seek(0, 0)
	if err != nil {
		return nil, err
	}

	return
}
<commit_msg>Add `--batch` in batch mode (fixes #519)<commit_after>package pgp

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// Test interface
var (
	_ Signer   = &GpgSigner{}
	_ Verifier = &GpgVerifier{}
)

// GpgSigner is implementation of Signer interface using gpg as external program
type GpgSigner struct {
	keyRef                     string
	keyring, secretKeyring     string
	passphrase, passphraseFile string
	batch                      bool
}

// SetBatch control --no-tty flag to gpg
func (g *GpgSigner) SetBatch(batch bool) {
	g.batch = batch
}

// SetKey sets key ID to use when signing files
func (g *GpgSigner) SetKey(keyRef string) {
	g.keyRef = keyRef
}

// SetKeyRing allows to set custom keyring and secretkeyring
func (g *GpgSigner) SetKeyRing(keyring, secretKeyring string) {
	g.keyring, g.secretKeyring = keyring, secretKeyring
}

// SetPassphrase sets passhprase params
func (g *GpgSigner) SetPassphrase(passphrase, passphraseFile string) {
	g.passphrase, g.passphraseFile = passphrase, passphraseFile
}

func (g *GpgSigner) gpgArgs() []string {
	args := []string{}
	if g.keyring != "" {
		args = append(args, "--no-auto-check-trustdb", "--no-default-keyring", "--keyring", g.keyring)
	}
	if g.secretKeyring != "" {
		args = append(args, "--secret-keyring", g.secretKeyring)
	}

	if g.keyRef != "" {
		args = append(args, "-u", g.keyRef)
	}

	if g.passphrase != "" || g.passphraseFile != "" {
		args = append(args, "--no-use-agent")
	}

	if g.passphrase != "" {
		args = append(args, "--passphrase", g.passphrase)
	}

	if g.passphraseFile != "" {
		args = append(args, "--passphrase-file", g.passphraseFile)
	}

	if g.batch {
		args = append(args, "--no-tty", "--batch")
	}

	return args
}

// Init verifies availability of gpg & presence of keys
func (g *GpgSigner) Init() error {
	output, err := exec.Command("gpg", "--list-keys", "--dry-run", "--no-auto-check-trustdb").CombinedOutput()
	if err != nil {
		return fmt.Errorf("unable to execute gpg: %s (is gpg installed?): %s", err, string(output))
	}

	if g.keyring == "" && g.secretKeyring == "" && len(output) == 0 {
		return fmt.Errorf("looks like there are no keys in gpg, please create one (official manual: http://www.gnupg.org/gph/en/manual.html)")
	}

	return err
}

// DetachedSign signs file with detached signature in ASCII format
func (g *GpgSigner) DetachedSign(source string, destination string) error {
	fmt.Printf("Signing file '%s' with gpg, please enter your passphrase when prompted:\n", filepath.Base(source))

	args := []string{"-o", destination, "--digest-algo", "SHA256", "--armor", "--yes"}
	args = append(args, g.gpgArgs()...)
	args = append(args, "--detach-sign", source)
	cmd := exec.Command("gpg", args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// ClearSign clear-signs the file
func (g *GpgSigner) ClearSign(source string, destination string) error {
	fmt.Printf("Clearsigning file '%s' with gpg, please enter your passphrase when prompted:\n", filepath.Base(source))
	args := []string{"-o", destination, "--digest-algo", "SHA256", "--yes"}
	args = append(args, g.gpgArgs()...)
	args = append(args, "--clearsign", source)
	cmd := exec.Command("gpg", args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// GpgVerifier is implementation of Verifier interface using gpgv as external program
type GpgVerifier struct {
	keyRings []string
}

// InitKeyring verifies that gpg is installed and some keys are trusted
func (g *GpgVerifier) InitKeyring() error {
	err := exec.Command("gpgv", "--version").Run()
	if err != nil {
		return fmt.Errorf("unable to execute gpgv: %s (is gpg installed?)", err)
	}

	if len(g.keyRings) == 0 {
		// using default keyring
		output, err := exec.Command("gpg", "--no-default-keyring", "--no-auto-check-trustdb", "--keyring", "trustedkeys.gpg", "--list-keys").Output()
		if err == nil && len(output) == 0 {
			fmt.Printf("\nLooks like your keyring with trusted keys is empty. You might consider importing some keys.\n")
			fmt.Printf("If you're running Debian or Ubuntu, it's a good idea to import current archive keys by running:\n\n")
			fmt.Printf("  gpg --no-default-keyring --keyring /usr/share/keyrings/debian-archive-keyring.gpg --export | gpg --no-default-keyring --keyring trustedkeys.gpg --import\n")
			fmt.Printf("\n(for Ubuntu, use /usr/share/keyrings/ubuntu-archive-keyring.gpg)\n\n")
		}
	}

	return nil
}

// AddKeyring adds custom keyring to GPG parameters
func (g *GpgVerifier) AddKeyring(keyring string) {
	g.keyRings = append(g.keyRings, keyring)
}

func (g *GpgVerifier) argsKeyrings() (args []string) {
	if len(g.keyRings) > 0 {
		args = make([]string, 0, 2*len(g.keyRings))
		for _, keyring := range g.keyRings {
			args = append(args, "--keyring", keyring)
		}
	} else {
		args = []string{"--keyring", "trustedkeys.gpg"}
	}
	return
}

func (g *GpgVerifier) runGpgv(args []string, context string, showKeyTip bool) (*KeyInfo, error) {
	args = append([]string{"--status-fd", "3"}, args...)
	cmd := exec.Command("gpgv", args...)

	tempf, err := ioutil.TempFile("", "aptly-gpg-status")
	if err != nil {
		return nil, err
	}
	defer tempf.Close()

	err = os.Remove(tempf.Name())
	if err != nil {
		return nil, err
	}

	cmd.ExtraFiles = []*os.File{tempf}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, err
	}
	defer stderr.Close()

	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	buffer := &bytes.Buffer{}

	_, err = io.Copy(io.MultiWriter(os.Stderr, buffer), stderr)
	if err != nil {
		return nil, err
	}

	cmderr := cmd.Wait()

	tempf.Seek(0, 0)

	statusr := bufio.NewScanner(tempf)

	result := &KeyInfo{}

	for statusr.Scan() {
		line := strings.TrimSpace(statusr.Text())

		if strings.HasPrefix(line, "[GNUPG:] GOODSIG ") {
			result.GoodKeys = append(result.GoodKeys, Key(strings.Fields(line)[2]))
		} else if strings.HasPrefix(line, "[GNUPG:] NO_PUBKEY ") {
			result.MissingKeys = append(result.MissingKeys, Key(strings.Fields(line)[2]))
		}
	}

	if err = statusr.Err(); err != nil {
		return nil, err
	}

	if cmderr != nil {
		if showKeyTip && len(g.keyRings) == 0 && len(result.MissingKeys) > 0 {
			fmt.Printf("\nLooks like some keys are missing in your trusted keyring, you may consider importing them from keyserver:\n\n")

			keys := make([]string, len(result.MissingKeys))

			for i := range result.MissingKeys {
				keys[i] = string(result.MissingKeys[i])
			}

			fmt.Printf("gpg --no-default-keyring --keyring trustedkeys.gpg --keyserver keys.gnupg.net --recv-keys %s\n\n",
				strings.Join(keys, " "))

			fmt.Printf("Sometimes keys are stored in repository root in file named Release.key, to import such key:\n\n")
			fmt.Printf("wget -O - https://some.repo/repository/Release.key | gpg --no-default-keyring --keyring trustedkeys.gpg --import\n\n")
		}
		return result, fmt.Errorf("verification of %s failed: %s", context, cmderr)
	}
	return result, nil
}

// VerifyDetachedSignature verifies combination of signature and cleartext using gpgv
func (g *GpgVerifier) VerifyDetachedSignature(signature, cleartext io.Reader, showKeyTip bool) error {
	args := g.argsKeyrings()

	sigf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return err
	}
	defer os.Remove(sigf.Name())
	defer sigf.Close()

	_, err = io.Copy(sigf, signature)
	if err != nil {
		return err
	}

	clearf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return err
	}
	defer os.Remove(clearf.Name())
	defer clearf.Close()

	_, err = io.Copy(clearf, cleartext)
	if err != nil {
		return err
	}

	args = append(args, sigf.Name(), clearf.Name())
	_, err = g.runGpgv(args, "detached signature", showKeyTip)
	return err
}

// IsClearSigned returns true if file contains signature
func (g *GpgVerifier) IsClearSigned(clearsigned io.Reader) (bool, error) {
	scanner := bufio.NewScanner(clearsigned)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), "BEGIN PGP SIGN") {
			return true, nil
		}
	}

	return false, scanner.Err()
}

// VerifyClearsigned verifies clearsigned file using gpgv
func (g *GpgVerifier) VerifyClearsigned(clearsigned io.Reader, showKeyTip bool) (*KeyInfo, error) {
	args := g.argsKeyrings()

	clearf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return nil, err
	}
	defer os.Remove(clearf.Name())
	defer clearf.Close()

	_, err = io.Copy(clearf, clearsigned)
	if err != nil {
		return nil, err
	}

	args = append(args, clearf.Name())
	return g.runGpgv(args, "clearsigned file", showKeyTip)
}

// ExtractClearsigned extracts cleartext from clearsigned file WITHOUT signature verification
func (g *GpgVerifier) ExtractClearsigned(clearsigned io.Reader) (text *os.File, err error) {
	clearf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return
	}
	defer os.Remove(clearf.Name())
	defer clearf.Close()

	_, err = io.Copy(clearf, clearsigned)
	if err != nil {
		return
	}

	text, err = ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return
	}
	defer os.Remove(text.Name())

	args := []string{"--no-auto-check-trustdb", "--decrypt", "--batch", "--skip-verify", "--output", "-", clearf.Name()}

	cmd := exec.Command("gpg", args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	defer stdout.Close()

	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	_, err = io.Copy(text, stdout)
	if err != nil {
		return nil, err
	}

	err = cmd.Wait()

	if err != nil {
		return nil, fmt.Errorf("extraction of clearsigned file failed: %s", err)
	}

	_, err = text.Seek(0, 0)
	if err != nil {
		return nil, err
	}

	return
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package color

// RGBToYCbCr converts an RGB triple to a Y'CbCr triple.
func RGBToYCbCr(r, g, b uint8) (uint8, uint8, uint8) {
	// The JFIF specification says:
	//	Y' =  0.2990*R + 0.5870*G + 0.1140*B
	//	Cb = -0.1687*R - 0.3313*G + 0.5000*B + 128
	//	Cr =  0.5000*R - 0.4187*G - 0.0813*B + 128
	// http://www.w3.org/Graphics/JPEG/jfif3.pdf says Y but means Y'.
	r1 := int(r)
	g1 := int(g)
	b1 := int(b)
	yy := (19595*r1 + 38470*g1 + 7471*b1 + 1<<15) >> 16
	cb := (-11056*r1 - 21712*g1 + 32768*b1 + 257<<15) >> 16
	cr := (32768*r1 - 27440*g1 - 5328*b1 + 257<<15) >> 16
	if yy < 0 {
		yy = 0
	} else if yy > 255 {
		yy = 255
	}
	if cb < 0 {
		cb = 0
	} else if cb > 255 {
		cb = 255
	}
	if cr < 0 {
		cr = 0
	} else if cr > 255 {
		cr = 255
	}
	return uint8(yy), uint8(cb), uint8(cr)
}

// YCbCrToRGB converts a Y'CbCr triple to an RGB triple.
func YCbCrToRGB(y, cb, cr uint8) (uint8, uint8, uint8) {
	// The JFIF specification says:
	//	R = Y' + 1.40200*(Cr-128)
	//	G = Y' - 0.34414*(Cb-128) - 0.71414*(Cr-128)
	//	B = Y' + 1.77200*(Cb-128)
	// http://www.w3.org/Graphics/JPEG/jfif3.pdf says Y but means Y'.
	yy1 := int(y)<<16 + 1<<15
	cb1 := int(cb) - 128
	cr1 := int(cr) - 128
	r := (yy1 + 91881*cr1) >> 16
	g := (yy1 - 22554*cb1 - 46802*cr1) >> 16
	b := (yy1 + 116130*cb1) >> 16
	if r < 0 {
		r = 0
	} else if r > 255 {
		r = 255
	}
	if g < 0 {
		g = 0
	} else if g > 255 {
		g = 255
	}
	if b < 0 {
		b = 0
	} else if b > 255 {
		b = 255
	}
	return uint8(r), uint8(g), uint8(b)
}

// YCbCr represents a fully opaque 24-bit Y'CbCr color, having 8 bits each for
// one luma and two chroma components.
//
// JPEG, VP8, the MPEG family and other codecs use this color model. Such
// codecs often use the terms YUV and Y'CbCr interchangeably, but strictly
// speaking, the term YUV applies only to analog video signals, and Y' (luma)
// is Y (luminance) after applying gamma correction.
//
// Conversion between RGB and Y'CbCr is lossy and there are multiple, slightly
// different formulae for converting between the two. This package follows
// the JFIF specification at http://www.w3.org/Graphics/JPEG/jfif3.pdf.
type YCbCr struct {
	Y, Cb, Cr uint8
}

func (c YCbCr) RGBA() (uint32, uint32, uint32, uint32) {
	r, g, b := YCbCrToRGB(c.Y, c.Cb, c.Cr)
	return uint32(r) * 0x101, uint32(g) * 0x101, uint32(b) * 0x101, 0xffff
}

// YCbCrModel is the Model for Y'CbCr colors.
var YCbCrModel Model = ModelFunc(modelYCbCr)

func modelYCbCr(c Color) Color {
	if _, ok := c.(YCbCr); ok {
		return c
	}
	r, g, b, _ := c.RGBA()
	y, u, v := RGBToYCbCr(uint8(r>>8), uint8(g>>8), uint8(b>>8))
	return YCbCr{y, u, v}
}
<commit_msg>image/color: rename modelYCbCr to yCbCrModel.<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package color

// RGBToYCbCr converts an RGB triple to a Y'CbCr triple.
func RGBToYCbCr(r, g, b uint8) (uint8, uint8, uint8) {
	// The JFIF specification says:
	//	Y' =  0.2990*R + 0.5870*G + 0.1140*B
	//	Cb = -0.1687*R - 0.3313*G + 0.5000*B + 128
	//	Cr =  0.5000*R - 0.4187*G - 0.0813*B + 128
	// http://www.w3.org/Graphics/JPEG/jfif3.pdf says Y but means Y'.
	r1 := int(r)
	g1 := int(g)
	b1 := int(b)
	yy := (19595*r1 + 38470*g1 + 7471*b1 + 1<<15) >> 16
	cb := (-11056*r1 - 21712*g1 + 32768*b1 + 257<<15) >> 16
	cr := (32768*r1 - 27440*g1 - 5328*b1 + 257<<15) >> 16
	if yy < 0 {
		yy = 0
	} else if yy > 255 {
		yy = 255
	}
	if cb < 0 {
		cb = 0
	} else if cb > 255 {
		cb = 255
	}
	if cr < 0 {
		cr = 0
	} else if cr > 255 {
		cr = 255
	}
	return uint8(yy), uint8(cb), uint8(cr)
}

// YCbCrToRGB converts a Y'CbCr triple to an RGB triple.
func YCbCrToRGB(y, cb, cr uint8) (uint8, uint8, uint8) {
	// The JFIF specification says:
	//	R = Y' + 1.40200*(Cr-128)
	//	G = Y' - 0.34414*(Cb-128) - 0.71414*(Cr-128)
	//	B = Y' + 1.77200*(Cb-128)
	// http://www.w3.org/Graphics/JPEG/jfif3.pdf says Y but means Y'.
	yy1 := int(y)<<16 + 1<<15
	cb1 := int(cb) - 128
	cr1 := int(cr) - 128
	r := (yy1 + 91881*cr1) >> 16
	g := (yy1 - 22554*cb1 - 46802*cr1) >> 16
	b := (yy1 + 116130*cb1) >> 16
	if r < 0 {
		r = 0
	} else if r > 255 {
		r = 255
	}
	if g < 0 {
		g = 0
	} else if g > 255 {
		g = 255
	}
	if b < 0 {
		b = 0
	} else if b > 255 {
		b = 255
	}
	return uint8(r), uint8(g), uint8(b)
}

// YCbCr represents a fully opaque 24-bit Y'CbCr color, having 8 bits each for
// one luma and two chroma components.
//
// JPEG, VP8, the MPEG family and other codecs use this color model. Such
// codecs often use the terms YUV and Y'CbCr interchangeably, but strictly
// speaking, the term YUV applies only to analog video signals, and Y' (luma)
// is Y (luminance) after applying gamma correction.
//
// Conversion between RGB and Y'CbCr is lossy and there are multiple, slightly
// different formulae for converting between the two. This package follows
// the JFIF specification at http://www.w3.org/Graphics/JPEG/jfif3.pdf.
type YCbCr struct {
	Y, Cb, Cr uint8
}

func (c YCbCr) RGBA() (uint32, uint32, uint32, uint32) {
	r, g, b := YCbCrToRGB(c.Y, c.Cb, c.Cr)
	return uint32(r) * 0x101, uint32(g) * 0x101, uint32(b) * 0x101, 0xffff
}

// YCbCrModel is the Model for Y'CbCr colors.
var YCbCrModel Model = ModelFunc(yCbCrModel)

func yCbCrModel(c Color) Color {
	if _, ok := c.(YCbCr); ok {
		return c
	}
	r, g, b, _ := c.RGBA()
	y, u, v := RGBToYCbCr(uint8(r>>8), uint8(g>>8), uint8(b>>8))
	return YCbCr{y, u, v}
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package strconv_test

import (
	"runtime"
	. "strconv"
	"testing"
)

type itob64Test struct {
	in   int64
	base int
	out  string
}

var itob64tests = []itob64Test{
	{0, 10, "0"},
	{1, 10, "1"},
	{-1, 10, "-1"},
	{12345678, 10, "12345678"},
	{-987654321, 10, "-987654321"},
	{1<<31 - 1, 10, "2147483647"},
	{-1<<31 + 1, 10, "-2147483647"},
	{1 << 31, 10, "2147483648"},
	{-1 << 31, 10, "-2147483648"},
	{1<<31 + 1, 10, "2147483649"},
	{-1<<31 - 1, 10, "-2147483649"},
	{1<<32 - 1, 10, "4294967295"},
	{-1<<32 + 1, 10, "-4294967295"},
	{1 << 32, 10, "4294967296"},
	{-1 << 32, 10, "-4294967296"},
	{1<<32 + 1, 10, "4294967297"},
	{-1<<32 - 1, 10, "-4294967297"},
	{1 << 50, 10, "1125899906842624"},
	{1<<63 - 1, 10, "9223372036854775807"},
	{-1<<63 + 1, 10, "-9223372036854775807"},
	{-1 << 63, 10, "-9223372036854775808"},

	{0, 2, "0"},
	{10, 2, "1010"},
	{-1, 2, "-1"},
	{1 << 15, 2, "1000000000000000"},

	{-8, 8, "-10"},
	{057635436545, 8, "57635436545"},
	{1 << 24, 8, "100000000"},

	{16, 16, "10"},
	{-0x123456789abcdef, 16, "-123456789abcdef"},
	{1<<63 - 1, 16, "7fffffffffffffff"},
	{1<<63 - 1, 2, "111111111111111111111111111111111111111111111111111111111111111"},

	{16, 17, "g"},
	{25, 25, "10"},
	{(((((17*35+24)*35+21)*35+34)*35+12)*35+24)*35 + 32, 35, "holycow"},
	{(((((17*36+24)*36+21)*36+34)*36+12)*36+24)*36 + 32, 36, "holycow"},
}

func TestItoa(t *testing.T) {
	for _, test := range itob64tests {
		s := FormatInt(test.in, test.base)
		if s != test.out {
			t.Errorf("FormatInt(%v, %v) = %v want %v",
				test.in, test.base, s, test.out)
		}
		x := AppendInt([]byte("abc"), test.in, test.base)
		if string(x) != "abc"+test.out {
			t.Errorf("AppendInt(%q, %v, %v) = %q want %v",
				"abc", test.in, test.base, x, test.out)
		}

		if test.in >= 0 {
			s := FormatUint(uint64(test.in), test.base)
			if s != test.out {
				t.Errorf("FormatUint(%v, %v) = %v want %v",
					test.in, test.base, s, test.out)
			}
			x := AppendUint(nil, uint64(test.in), test.base)
			if string(x) != test.out {
				t.Errorf("AppendUint(%q, %v, %v) = %q want %v",
					"abc", uint64(test.in), test.base, x, test.out)
			}
		}

		if test.base == 10 && int64(int(test.in)) == test.in {
			s := Itoa(int(test.in))
			if s != test.out {
				t.Errorf("Itoa(%v) = %v want %v",
					test.in, s, test.out)
			}
		}
	}
}

type uitob64Test struct {
	in   uint64
	base int
	out  string
}

var uitob64tests = []uitob64Test{
	{1<<63 - 1, 10, "9223372036854775807"},
	{1 << 63, 10, "9223372036854775808"},
	{1<<63 + 1, 10, "9223372036854775809"},
	{1<<64 - 2, 10, "18446744073709551614"},
	{1<<64 - 1, 10, "18446744073709551615"},
	{1<<64 - 1, 2, "1111111111111111111111111111111111111111111111111111111111111111"},
}

func TestUitoa(t *testing.T) {
	for _, test := range uitob64tests {
		s := FormatUint(test.in, test.base)
		if s != test.out {
			t.Errorf("FormatUint(%v, %v) = %v want %v",
				test.in, test.base, s, test.out)
		}
		x := AppendUint([]byte("abc"), test.in, test.base)
		if string(x) != "abc"+test.out {
			t.Errorf("AppendUint(%q, %v, %v) = %q want %v",
				"abc", test.in, test.base, x, test.out)
		}

	}
}

func numAllocations(f func()) int {
	memstats := new(runtime.MemStats)
	runtime.ReadMemStats(memstats)
	n0 := memstats.Mallocs
	f()
	runtime.ReadMemStats(memstats)
	return int(memstats.Mallocs - n0)
}

var globalBuf [64]byte

func TestAppendUintDoesntAllocate(t *testing.T) {
	n := numAllocations(func() {
		var buf [64]byte
		AppendInt(buf[:0], 123, 10)
	})
	want := 1 // TODO(bradfitz): this might be 0, once escape analysis is better
	if n != want {
		t.Errorf("with local buffer, did %d allocations, want %d", n, want)
	}
	n = numAllocations(func() {
		AppendInt(globalBuf[:0], 123, 10)
	})
	if n != 0 {
		t.Errorf("with reused buffer, did %d allocations, want 0", n)
	}
}

func BenchmarkFormatInt(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, test := range itob64tests {
			FormatInt(test.in, test.base)
		}
	}
}

func BenchmarkAppendInt(b *testing.B) {
	dst := make([]byte, 0, 30)
	for i := 0; i < b.N; i++ {
		for _, test := range itob64tests {
			AppendInt(dst, test.in, test.base)
		}
	}
}

func BenchmarkFormatUint(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, test := range uitob64tests {
			FormatUint(test.in, test.base)
		}
	}
}

func BenchmarkAppendUint(b *testing.B) {
	dst := make([]byte, 0, 30)
	for i := 0; i < b.N; i++ {
		for _, test := range uitob64tests {
			AppendUint(dst, test.in, test.base)
		}
	}
}
<commit_msg>strconv: run garbage collection before counting allocations in test<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package strconv_test

import (
	"runtime"
	. "strconv"
	"testing"
)

type itob64Test struct {
	in   int64
	base int
	out  string
}

var itob64tests = []itob64Test{
	{0, 10, "0"},
	{1, 10, "1"},
	{-1, 10, "-1"},
	{12345678, 10, "12345678"},
	{-987654321, 10, "-987654321"},
	{1<<31 - 1, 10, "2147483647"},
	{-1<<31 + 1, 10, "-2147483647"},
	{1 << 31, 10, "2147483648"},
	{-1 << 31, 10, "-2147483648"},
	{1<<31 + 1, 10, "2147483649"},
	{-1<<31 - 1, 10, "-2147483649"},
	{1<<32 - 1, 10, "4294967295"},
	{-1<<32 + 1, 10, "-4294967295"},
	{1 << 32, 10, "4294967296"},
	{-1 << 32, 10, "-4294967296"},
	{1<<32 + 1, 10, "4294967297"},
	{-1<<32 - 1, 10, "-4294967297"},
	{1 << 50, 10, "1125899906842624"},
	{1<<63 - 1, 10, "9223372036854775807"},
	{-1<<63 + 1, 10, "-9223372036854775807"},
	{-1 << 63, 10, "-9223372036854775808"},

	{0, 2, "0"},
	{10, 2, "1010"},
	{-1, 2, "-1"},
	{1 << 15, 2, "1000000000000000"},

	{-8, 8, "-10"},
	{057635436545, 8, "57635436545"},
	{1 << 24, 8, "100000000"},

	{16, 16, "10"},
	{-0x123456789abcdef, 16, "-123456789abcdef"},
	{1<<63 - 1, 16, "7fffffffffffffff"},
	{1<<63 - 1, 2, "111111111111111111111111111111111111111111111111111111111111111"},

	{16, 17, "g"},
	{25, 25, "10"},
	{(((((17*35+24)*35+21)*35+34)*35+12)*35+24)*35 + 32, 35, "holycow"},
	{(((((17*36+24)*36+21)*36+34)*36+12)*36+24)*36 + 32, 36, "holycow"},
}

func TestItoa(t *testing.T) {
	for _, test := range itob64tests {
		s := FormatInt(test.in, test.base)
		if s != test.out {
			t.Errorf("FormatInt(%v, %v) = %v want %v",
				test.in, test.base, s, test.out)
		}
		x := AppendInt([]byte("abc"), test.in, test.base)
		if string(x) != "abc"+test.out {
			t.Errorf("AppendInt(%q, %v, %v) = %q want %v",
				"abc", test.in, test.base, x, test.out)
		}

		if test.in >= 0 {
			s := FormatUint(uint64(test.in), test.base)
			if s != test.out {
				t.Errorf("FormatUint(%v, %v) = %v want %v",
					test.in, test.base, s, test.out)
			}
			x := AppendUint(nil, uint64(test.in), test.base)
			if string(x) != test.out {
				t.Errorf("AppendUint(%q, %v, %v) = %q want %v",
					"abc", uint64(test.in), test.base, x, test.out)
			}
		}

		if test.base == 10 && int64(int(test.in)) == test.in {
			s := Itoa(int(test.in))
			if s != test.out {
				t.Errorf("Itoa(%v) = %v want %v",
					test.in, s, test.out)
			}
		}
	}
}

type uitob64Test struct {
	in   uint64
	base int
	out  string
}

var uitob64tests = []uitob64Test{
	{1<<63 - 1, 10, "9223372036854775807"},
	{1 << 63, 10, "9223372036854775808"},
	{1<<63 + 1, 10, "9223372036854775809"},
	{1<<64 - 2, 10, "18446744073709551614"},
	{1<<64 - 1, 10, "18446744073709551615"},
	{1<<64 - 1, 2, "1111111111111111111111111111111111111111111111111111111111111111"},
}

func TestUitoa(t *testing.T) {
	for _, test := range uitob64tests {
		s := FormatUint(test.in, test.base)
		if s != test.out {
			t.Errorf("FormatUint(%v, %v) = %v want %v",
				test.in, test.base, s, test.out)
		}
		x := AppendUint([]byte("abc"), test.in, test.base)
		if string(x) != "abc"+test.out {
			t.Errorf("AppendUint(%q, %v, %v) = %q want %v",
				"abc", test.in, test.base, x, test.out)
		}

	}
}

func numAllocations(f func()) int {
	runtime.GC()
	memstats := new(runtime.MemStats)
	runtime.ReadMemStats(memstats)
	n0 := memstats.Mallocs
	f()
	runtime.ReadMemStats(memstats)
	return int(memstats.Mallocs - n0)
}

var globalBuf [64]byte

func TestAppendUintDoesntAllocate(t *testing.T) {
	n := numAllocations(func() {
		var buf [64]byte
		AppendInt(buf[:0], 123, 10)
	})
	want := 1 // TODO(bradfitz): this might be 0, once escape analysis is better
	if n != want {
		t.Errorf("with local buffer, did %d allocations, want %d", n, want)
	}
	n = numAllocations(func() {
		AppendInt(globalBuf[:0], 123, 10)
	})
	if n != 0 {
		t.Errorf("with reused buffer, did %d allocations, want 0", n)
	}
}

func BenchmarkFormatInt(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, test := range itob64tests {
			FormatInt(test.in, test.base)
		}
	}
}

func BenchmarkAppendInt(b *testing.B) {
	dst := make([]byte, 0, 30)
	for i := 0; i < b.N; i++ {
		for _, test := range itob64tests {
			AppendInt(dst, test.in, test.base)
		}
	}
}

func BenchmarkFormatUint(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, test := range uitob64tests {
			FormatUint(test.in, test.base)
		}
	}
}

func BenchmarkAppendUint(b *testing.B) {
	dst := make([]byte, 0, 30)
	for i := 0; i < b.N; i++ {
		for _, test := range uitob64tests {
			AppendUint(dst, test.in, test.base)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testing

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"time"
)

var matchBenchmarks = flag.String("test.bench", "", "regular expression to select benchmarks to run")

// An internal type but exported because it is cross-package; part of the implementation
// of gotest.
type InternalBenchmark struct {
	Name string
	F    func(b *B)
}

// B is a type passed to Benchmark functions to manage benchmark
// timing and to specify the number of iterations to run.
type B struct {
	N         int
	benchmark InternalBenchmark
	ns        int64
	bytes     int64
	start     int64
}

// StartTimer starts timing a test.  This function is called automatically
// before a benchmark starts, but it can also used to resume timing after
// a call to StopTimer.
func (b *B) StartTimer() { b.start = time.Nanoseconds() }

// StopTimer stops timing a test.  This can be used to pause the timer
// while performing complex initialization that you don't
// want to measure.
func (b *B) StopTimer() {
	if b.start > 0 {
		b.ns += time.Nanoseconds() - b.start
	}
	b.start = 0
}

// ResetTimer stops the timer and sets the elapsed benchmark time to zero.
func (b *B) ResetTimer() {
	b.start = 0
	b.ns = 0
}

// SetBytes records the number of bytes processed in a single operation.
// If this is called, the benchmark will report ns/op and MB/s.
func (b *B) SetBytes(n int64) { b.bytes = n }

func (b *B) nsPerOp() int64 {
	if b.N <= 0 {
		return 0
	}
	return b.ns / int64(b.N)
}

// runN runs a single benchmark for the specified number of iterations.
func (b *B) runN(n int) {
	// Try to get a comparable environment for each run
	// by clearing garbage from previous runs.
	runtime.GC()
	b.N = n
	b.ResetTimer()
	b.StartTimer()
	b.benchmark.F(b)
	b.StopTimer()
}

func min(x, y int) int {
	if x > y {
		return y
	}
	return x
}

func max(x, y int) int {
	if x < y {
		return y
	}
	return x
}

// roundDown10 rounds a number down to the nearest power of 10.
func roundDown10(n int) int {
	var tens = 0
	// tens = floor(log_10(n))
	for n > 10 {
		n = n / 10
		tens++
	}
	// result = 10^tens
	result := 1
	for i := 0; i < tens; i++ {
		result *= 10
	}
	return result
}

// roundUp rounds x up to a number of the form [1eX, 2eX, 5eX].
func roundUp(n int) int {
	base := roundDown10(n)
	if n < (2 * base) {
		return 2 * base
	}
	if n < (5 * base) {
		return 5 * base
	}
	return 10 * base
}

// run times the benchmark function.  It gradually increases the number
// of benchmark iterations until the benchmark runs for a second in order
// to get a reasonable measurement.  It prints timing information in this form
//		testing.BenchmarkHello	100000		19 ns/op
func (b *B) run() BenchmarkResult {
	// Run the benchmark for a single iteration in case it's expensive.
	n := 1
	b.runN(n)
	// Run the benchmark for at least a second.
	for b.ns < 1e9 && n < 1e9 {
		last := n
		// Predict iterations/sec.
		if b.nsPerOp() == 0 {
			n = 1e9
		} else {
			n = 1e9 / int(b.nsPerOp())
		}
		// Run more iterations than we think we'll need for a second (1.5x).
		// Don't grow too fast in case we had timing errors previously.
		// Be sure to run at least one more than last time.
		n = max(min(n+n/2, 100*last), last+1)
		// Round up to something easy to read.
		n = roundUp(n)
		b.runN(n)
	}
	return BenchmarkResult{b.N, b.ns, b.bytes}

}

// The results of a benchmark run.
type BenchmarkResult struct {
	N     int   // The number of iterations.
	Ns    int64 // The total time taken.
	Bytes int64 // The total number of bytes processed.
}

func (r BenchmarkResult) NsPerOp() int64 {
	if r.N <= 0 {
		return 0
	}
	return r.Ns / int64(r.N)
}

func (r BenchmarkResult) String() string {
	ns := r.NsPerOp()
	mb := ""
	if ns > 0 && r.Bytes > 0 {
		mb = fmt.Sprintf("\t%7.2f MB/s", (float64(r.Bytes)/1e6)/(float64(ns)/1e9))
	}
	return fmt.Sprintf("%8d\t%10d ns/op%s", r.N, ns, mb)
}

// An internal function but exported because it is cross-package; part of the implementation
// of gotest.
func RunBenchmarks(matchString func(pat, str string) (bool, os.Error), benchmarks []InternalBenchmark) {
	// If no flag was specified, don't run benchmarks.
	if len(*matchBenchmarks) == 0 {
		return
	}
	for _, Benchmark := range benchmarks {
		matched, err := matchString(*matchBenchmarks, Benchmark.Name)
		if err != nil {
			println("invalid regexp for -test.bench:", err.String())
			os.Exit(1)
		}
		if !matched {
			continue
		}
		b := &B{benchmark: Benchmark}
		r := b.run()
		fmt.Printf("%s\t%v\n", Benchmark.Name, r)
	}
}

// Benchmark benchmarks a single function. Useful for creating
// custom benchmarks that do not use gotest.
func Benchmark(f func(b *B)) BenchmarkResult {
	b := &B{benchmark: InternalBenchmark{"", f}}
	return b.run()
}
<commit_msg>testing: fix MB/s computation, documentation<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testing

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"time"
)

var matchBenchmarks = flag.String("test.bench", "", "regular expression to select benchmarks to run")

// An internal type but exported because it is cross-package; part of the implementation
// of gotest.
type InternalBenchmark struct {
	Name string
	F    func(b *B)
}

// B is a type passed to Benchmark functions to manage benchmark
// timing and to specify the number of iterations to run.
type B struct {
	N         int
	benchmark InternalBenchmark
	ns        int64
	bytes     int64
	start     int64
}

// StartTimer starts timing a test.  This function is called automatically
// before a benchmark starts, but it can also used to resume timing after
// a call to StopTimer.
func (b *B) StartTimer() { b.start = time.Nanoseconds() }

// StopTimer stops timing a test.  This can be used to pause the timer
// while performing complex initialization that you don't
// want to measure.
func (b *B) StopTimer() {
	if b.start > 0 {
		b.ns += time.Nanoseconds() - b.start
	}
	b.start = 0
}

// ResetTimer stops the timer and sets the elapsed benchmark time to zero.
func (b *B) ResetTimer() {
	b.start = 0
	b.ns = 0
}

// SetBytes records the number of bytes processed in a single operation.
// If this is called, the benchmark will report ns/op and MB/s.
func (b *B) SetBytes(n int64) { b.bytes = n }

func (b *B) nsPerOp() int64 {
	if b.N <= 0 {
		return 0
	}
	return b.ns / int64(b.N)
}

// runN runs a single benchmark for the specified number of iterations.
func (b *B) runN(n int) {
	// Try to get a comparable environment for each run
	// by clearing garbage from previous runs.
	runtime.GC()
	b.N = n
	b.ResetTimer()
	b.StartTimer()
	b.benchmark.F(b)
	b.StopTimer()
}

func min(x, y int) int {
	if x > y {
		return y
	}
	return x
}

func max(x, y int) int {
	if x < y {
		return y
	}
	return x
}

// roundDown10 rounds a number down to the nearest power of 10.
func roundDown10(n int) int {
	var tens = 0
	// tens = floor(log_10(n))
	for n > 10 {
		n = n / 10
		tens++
	}
	// result = 10^tens
	result := 1
	for i := 0; i < tens; i++ {
		result *= 10
	}
	return result
}

// roundUp rounds x up to a number of the form [1eX, 2eX, 5eX].
func roundUp(n int) int {
	base := roundDown10(n)
	if n < (2 * base) {
		return 2 * base
	}
	if n < (5 * base) {
		return 5 * base
	}
	return 10 * base
}

// run times the benchmark function.  It gradually increases the number
// of benchmark iterations until the benchmark runs for a second in order
// to get a reasonable measurement.  It prints timing information in this form
//		testing.BenchmarkHello	100000		19 ns/op
func (b *B) run() BenchmarkResult {
	// Run the benchmark for a single iteration in case it's expensive.
	n := 1
	b.runN(n)
	// Run the benchmark for at least a second.
	for b.ns < 1e9 && n < 1e9 {
		last := n
		// Predict iterations/sec.
		if b.nsPerOp() == 0 {
			n = 1e9
		} else {
			n = 1e9 / int(b.nsPerOp())
		}
		// Run more iterations than we think we'll need for a second (1.5x).
		// Don't grow too fast in case we had timing errors previously.
		// Be sure to run at least one more than last time.
		n = max(min(n+n/2, 100*last), last+1)
		// Round up to something easy to read.
		n = roundUp(n)
		b.runN(n)
	}
	return BenchmarkResult{b.N, b.ns, b.bytes}

}

// The results of a benchmark run.
type BenchmarkResult struct {
	N     int   // The number of iterations.
	Ns    int64 // The total time taken.
	Bytes int64 // Bytes processed in one iteration.
}

func (r BenchmarkResult) NsPerOp() int64 {
	if r.N <= 0 {
		return 0
	}
	return r.Ns / int64(r.N)
}

func (r BenchmarkResult) mbPerSec() float64 {
	if r.Bytes <= 0 || r.Ns <= 0 || r.N <= 0 {
		return 0
	}
	return float64(r.Bytes) * float64(r.N) / float64(r.Ns) * 1e3
}

func (r BenchmarkResult) String() string {
	mbs := r.mbPerSec()
	mb := ""
	if mbs != 0 {
		mb = fmt.Sprintf("\t%7.2f MB/s", mbs)
	}
	return fmt.Sprintf("%8d\t%10d ns/op%s", r.N, r.NsPerOp(), mb)
}

// An internal function but exported because it is cross-package; part of the implementation
// of gotest.
func RunBenchmarks(matchString func(pat, str string) (bool, os.Error), benchmarks []InternalBenchmark) {
	// If no flag was specified, don't run benchmarks.
	if len(*matchBenchmarks) == 0 {
		return
	}
	for _, Benchmark := range benchmarks {
		matched, err := matchString(*matchBenchmarks, Benchmark.Name)
		if err != nil {
			println("invalid regexp for -test.bench:", err.String())
			os.Exit(1)
		}
		if !matched {
			continue
		}
		b := &B{benchmark: Benchmark}
		r := b.run()
		fmt.Printf("%s\t%v\n", Benchmark.Name, r)
	}
}

// Benchmark benchmarks a single function. Useful for creating
// custom benchmarks that do not use gotest.
func Benchmark(f func(b *B)) BenchmarkResult {
	b := &B{benchmark: InternalBenchmark{"", f}}
	return b.run()
}
<|endoftext|>


<commit_before>package pipe

import (
	"../common"
	"../ikcp"
	"encoding/binary"
	"errors"
	"log"
	"net"
	"time"
)

const WriteBufferSize = 5000 //udp writer will add some data for checksum or encrypt
const ReadBufferSize = 7000  //so reader must be larger

const dataLimit = 4000

const mainV = 0
const subV = 1

type cache struct {
	b []byte
	l int
	c chan int
}

func init() {
}

type Action struct {
	t    string
	args []interface{}
}

func iclock() int32 {
	return int32((time.Now().UnixNano() / 1000000) & 0xffffffff)
}

func udp_output(buf []byte, _len int32, kcp *ikcp.Ikcpcb, user interface{}) int32 {
	c := user.(*UDPMakeSession)
	//log.Println("send udp", _len, c.remote.String())
	c.sock.WriteTo(buf[:_len], c.remote)
	return 0
}

const (
	Reset     byte = 0
	FirstSYN  byte = 6
	FirstACK  byte = 1
	SndSYN    byte = 2
	SndACK    byte = 3
	Data      byte = 4
	Ping      byte = 5
	Close     byte = 7
	CloseBack byte = 8
	ResetAck  byte = 9
)

func makeEncode(buf []byte, status byte, arg int) []byte {
	buf[0] = status
	binary.LittleEndian.PutUint32(buf[1:], uint32(arg))
	return buf
}

func makeDecode(data []byte) (status byte, arg int32) {
	if len(data) < 5 {
		return Reset, 0
	}
	status = data[0]
	arg = int32(binary.LittleEndian.Uint32(data[1:]))
	return
}

type UDPMakeSession struct {
	id                int
	idstr             string
	status            string
	overTime          int64
	quitChan          chan bool
	closeChan         chan bool
	recvChan          chan cache
	handShakeChan     chan string
	handShakeChanQuit chan bool
	sock              *net.UDPConn
	remote            *net.UDPAddr
	kcp               *ikcp.Ikcpcb
	do                chan Action
	do2               chan Action
	checkCanWrite     chan (chan bool)
	listener          *Listener
	closed            bool

	readBuffer    []byte
	processBuffer []byte
	encodeBuffer  []byte
	timeout       int64
	disBind       bool
	noClose  bool
}

type Listener struct {
	connChan   chan *UDPMakeSession
	quitChan   chan bool
	sock       *net.UDPConn
	readBuffer []byte
	sessions   map[string]*UDPMakeSession
}

func (l *Listener) Accept() (net.Conn, error) {
	var c *UDPMakeSession
	var err error
	select {
	case c = <-l.connChan:
	case <-l.quitChan:
	}
	if c == nil {
		err = errors.New("listener quit")
	}
	return net.Conn(c), err
}

func (l *Listener) Dump() {
	for addr, session := range l.sessions {
		log.Println("listener", addr, session.status)
	}
}

func (l *Listener) inner_loop() {
	sock := l.sock
	for {
		n, from, err := sock.ReadFromUDP(l.readBuffer)
		if err == nil {
			//log.Println("recv", n, from)
			addr := from.String()
			session, bHave := l.sessions[addr]
			if bHave {
				if session.status == "ok" {
					if session.remote.String() == from.String() && n >= int(ikcp.IKCP_OVERHEAD) {
						buf := make([]byte, n)
						copy(buf, l.readBuffer[:n])
						session.DoAction2("input", buf, n)
					}
					continue
				} else {
					session.serverDo(string(l.readBuffer[:n]))
				}
			} else {
				status, _ := makeDecode(l.readBuffer[:n])
				if status != FirstSYN {
					go sock.WriteToUDP([]byte("0"), from)
					//log.Println("invalid package,reset", from, status)
					continue
				}
				sessionId := common.GetId("udp")
				session = &UDPMakeSession{status: "init", overTime: time.Now().Unix() + 10, remote: from, sock: sock, recvChan: make(chan cache), quitChan: make(chan bool), readBuffer: make([]byte, ReadBufferSize), processBuffer: make([]byte, ReadBufferSize), timeout: 30, do: make(chan Action), do2: make(chan Action), id: sessionId, handShakeChan: make(chan string), handShakeChanQuit: make(chan bool), listener: l, closeChan: make(chan bool), encodeBuffer: make([]byte, 5), checkCanWrite: make(chan (chan bool))}
				l.sessions[addr] = session
				session.serverInit(l)
				session.serverDo(string(l.readBuffer[:n]))
			}
			//log.Println("debug out.........")
		} else {
			e, ok := err.(net.Error)
			if !ok || !e.Timeout() {
				log.Println("recv error", err.Error(), from)
				l.remove(from.String())
				//time.Sleep(time.Second)
				break
			}
		}
	}
}

func (l *Listener) remove(addr string) {
	log.Println("listener remove", addr)
	session, bHave := l.sessions[addr]
	if bHave {
		common.RmId("udp", session.id)
	}
	delete(l.sessions, addr)
}

func (l *Listener) Close() error {
	if l.sock != nil {
		l.sock.Close()
		l.sock = nil
	} else {
		return nil
	}
	close(l.quitChan)
	return nil
}

func (l *Listener) Addr() net.Addr {
	return nil
}

func (l *Listener) loop() {
	l.inner_loop()
}

func Listen(addr string) (*Listener, error) {
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return nil, err
	}
	sock, _err := net.ListenUDP("udp", udpAddr)
	if _err != nil {
		return nil, _err
	}

	listener := &Listener{connChan: make(chan *UDPMakeSession), quitChan: make(chan bool), sock: sock, readBuffer: make([]byte, ReadBufferSize), sessions: make(map[string]*UDPMakeSession)}
	go listener.loop()
	return listener, nil
}

func Dial(addr string) (*UDPMakeSession, error) {
	return DialTimeout(addr, 30)
}

func DialTimeout(addr string, timeout int) (*UDPMakeSession, error) {
	bReset := false
	if timeout < 5 {
		bReset = true
		timeout = 5
	} else if timeout > 255 {
		bReset = true
		timeout = 255
	}
	if bReset {
		log.Println("timeout should in [5, 255], force reset timeout to", timeout)
	}
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return nil, err
	}
	sock, _err := net.ListenUDP("udp", &net.UDPAddr{})
	if _err != nil {
		log.Println("dial addr fail", _err.Error())
		return nil, _err
	}
	session := &UDPMakeSession{readBuffer: make([]byte, ReadBufferSize), do: make(chan Action), do2: make(chan Action), quitChan: make(chan bool), recvChan: make(chan cache), processBuffer: make([]byte, ReadBufferSize), closeChan: make(chan bool), encodeBuffer: make([]byte, 5), checkCanWrite: make(chan (chan bool))}
	session.remote = udpAddr
	session.sock = sock
	session.status = "firstsyn"
	session.timeout = int64(timeout)
	_timeout := int(timeout / 2)
	if _timeout < 5 {
		timeout = 5
	}
	arg := int(int32(timeout) + int32(mainV<<24) + int32(subV<<16))
	info := makeEncode(session.encodeBuffer, FirstSYN, arg)
	code := session.doAndWait(func() {
		sock.WriteToUDP(info, udpAddr)
	}, _timeout, func(status byte, arg int32) int {
		if status == ResetAck {
			_mainV, _subV := int(byte(arg>>24)), int(byte(arg>>16))
			log.Printf("pipe version not eq,%d.%d=>%d.%d", mainV, subV, _mainV, _subV)
			return 1
		}
		if status != FirstACK {
			return -1
		} else {
			session.status = "firstack"
			session.id = int(arg)
			return 0
		}
	})
	if code != 0 {
		return nil, errors.New("handshake fail,1")
	}
	code = session.doAndWait(func() {
		sock.WriteToUDP(makeEncode(session.encodeBuffer, SndSYN, session.id), udpAddr)
	}, _timeout, func(status byte, arg int32) int {
		if status != SndACK {
			return -1
		} else if session.id != int(arg) {
			return 2
		} else {
			session.status = "ok"
			return 0
		}
	})
	if code != 0 {
		return nil, errors.New("handshake fail,2")
	}
	session.kcp = ikcp.Ikcp_create(uint32(session.id), session)
	session.kcp.Output = udp_output
	ikcp.Ikcp_wndsize(session.kcp, 128, 128)
	ikcp.Ikcp_nodelay(session.kcp, 1, 10, 2, 1)
	go session.loop()
	return session, nil
}

func (session *UDPMakeSession) SetNoClose() {
	session.noClose = true
}

func (session *UDPMakeSession) GetSock() *net.UDPConn {
	return session.sock
}
func (session *UDPMakeSession) doAndWait(f func(), sec int, readf func(status byte, arg int32) int) (code int) {
	t := time.NewTicker(50 * time.Millisecond)
	currT := time.Now().Unix()
	f()
out:
	for {
		select {
		case <-session.quitChan:
			break out
		case <-t.C:
			if time.Now().Unix()-currT >= int64(sec) {
				code = -1
				break out
			}
			session.sock.SetReadDeadline(time.Now().Add(2 * time.Second))
			n, from, err := session.sock.ReadFromUDP(session.readBuffer)
			if err != nil {
				e, ok := err.(net.Error)
				if !ok || !e.Timeout() {
					log.Println("recv error", err.Error(), from)
					code = -2
					break out
				}
			} else {
				code = readf(makeDecode(session.readBuffer[:n]))
				if code >= 0 {
					break out
				}
			}
			go f()
		}
	}
	t.Stop()
	if code > 0 {
		log.Println("handshake fail,got code", code)
	}
	return
}

func (session *UDPMakeSession) serverDo(s string) {
	go func() {
		//log.Println("prepare handshake", session.remote)
		select {
		case <-session.handShakeChanQuit:
		case session.handShakeChan <- s:
		}
	}()
}
func (session *UDPMakeSession) serverInit(l *Listener) {
	go func() {
		c := time.NewTicker(50 * time.Millisecond)
		defer func() {
			close(session.handShakeChanQuit)
			c.Stop()
			if session.status != "ok" {
				session.Close()
			}
		}()
		overTime := time.Now().Unix() + session.timeout
	out:
		for {
			select {
			case s := <-session.handShakeChan:
				//log.Println("process handshake", session.remote)
				status, arg := makeDecode([]byte(s))
				switch session.status {
				case "init":
					if status != FirstSYN {
						session.sock.WriteToUDP(makeEncode(session.encodeBuffer, Reset, 0), session.remote)
						return
					}
					_mainV, _subV := int(byte(arg>>24)), int(byte(arg>>16))
					if _mainV != mainV || _subV != subV {
						session.sock.WriteToUDP(makeEncode(session.encodeBuffer, ResetAck, (mainV<<24)+(subV<<16)), session.remote)
						log.Printf("pipe version not eq,kickout,%d.%d=>%d.%d", mainV, subV, _mainV, _subV)
						return
					}
					session.status = "firstack"
					session.timeout = int64(arg & 0xff)
					session.sock.WriteToUDP(makeEncode(session.encodeBuffer, FirstACK, session.id), session.remote)
					overTime = time.Now().Unix() + session.timeout
				case "firstack":
					if status != SndSYN {
						return
					}
					session.status = "ok"
					session.kcp = ikcp.Ikcp_create(uint32(session.id), session)
					session.kcp.Output = udp_output
					ikcp.Ikcp_wndsize(session.kcp, 128, 128)
					ikcp.Ikcp_nodelay(session.kcp, 1, 10, 2, 1)
					go session.loop()
					go func() {
						select {
						case l.connChan <- session:
						case <-l.quitChan:
						}
					}()
					session.sock.WriteToUDP(makeEncode(session.encodeBuffer, SndACK, session.id), session.remote)
					overTime = time.Now().Unix() + session.timeout
				}
			case <-session.quitChan:
				break out
			case <-c.C:
				if time.Now().Unix() > overTime {
					return
				}
				switch session.status {
				case "firstack":
					session.sock.WriteToUDP(makeEncode(session.encodeBuffer, FirstACK, session.id), session.remote)
				case "ok":
					buf := make([]byte, 5)
					session.sock.WriteToUDP(makeEncode(buf, SndACK, session.id), session.remote)
				}
			}
		}
	}()
}

func (session *UDPMakeSession) loop() {
	curr := time.Now().Unix()
	session.overTime = curr + session.timeout
	ping := make(chan bool)
	pingC := 0
	callUpdate := false
	updateC := make(chan bool)
	if session.listener == nil {
		go func() {
			tmp := session.readBuffer
			t := time.Time{}
			session.sock.SetReadDeadline(t)
			for {
				n, from, err := session.sock.ReadFromUDP(tmp)
				if err != nil {
					e, ok := err.(net.Error)
					if !ok || !e.Timeout() {
						break
					}
					if session.disBind && ok && e.Timeout() {
						log.Println("force timeout!!!")
						session.sock.SetReadDeadline(time.Time{})
						break
					}
				}
				if session.remote.String() == from.String() {
					if n >= int(ikcp.IKCP_OVERHEAD) || n <= 5 {
						buf := make([]byte, n)
						copy(buf, session.readBuffer[:n])
						session.DoAction2("input", buf, n)
					}
				}
			}
		}()
	}
	updateF := func(n time.Duration) {
		if !callUpdate {
			callUpdate = true
			time.AfterFunc(n*time.Millisecond, func() {
				select {
				case updateC <- true:
				case <-session.quitChan:
				}
			})
		}
	}

	fastCheck := false
	waitList := [](chan bool){}
	recoverChan := make(chan bool)

	var waitRecvCache *cache
	go func() {
	out:
		for {
			select {
			//session.wait.Done()
			case <-ping:
				updateF(50)
				pingC++
				if pingC >= 4 {
					pingC = 0
					if ikcp.Ikcp_waitsnd(session.kcp) <= dataLimit/2 {
						go session.DoWrite(makeEncode(session.encodeBuffer, Ping, 0))
					}
				}
				if time.Now().Unix() > session.overTime {
					log.Println("overtime close", session.LocalAddr().String(), session.RemoteAddr().String())
					go session.Close()
				} else {
					time.AfterFunc(300*time.Millisecond, func() {
						select {
						case ping <- true:
						case <-session.quitChan:
						}
					})
				}
			case <-recoverChan:
				fastCheck = false
				for _, r := range waitList {
					log.Println("recover writing data")
					select {
					case r <- true:
					case <-session.quitChan:
					}
				}
				waitList = [](chan bool){}
			case c := <-session.checkCanWrite:
				if ikcp.Ikcp_waitsnd(session.kcp) > dataLimit {
					log.Println("wait for data limit")
					waitList = append(waitList, c)
					if !fastCheck {
						fastCheck = true
						var f func()
						f = func() {
							n := ikcp.Ikcp_waitsnd(session.kcp)
							//log.Println("fast check!", n, len(waitList))
							if n <= dataLimit/2 {
								select {
								case <-session.quitChan:
								case recoverChan <- true:
								}
							} else {
								updateF(20)
								time.AfterFunc(40*time.Millisecond, f)
							}
						}
						time.AfterFunc(20*time.Millisecond, f)
					}
				} else {
					select {
					case c <- true:
					case <-session.quitChan:
					}
				}
			case ca := <-session.recvChan:
				tmp := session.processBuffer
				for {
					hr := ikcp.Ikcp_recv(session.kcp, tmp, ReadBufferSize)
					if hr > 0 {
						status := tmp[0]
						if status == Data {
							copy(ca.b, tmp[1:hr])
							ca.c <- int(hr - 1)
							break
						} else {
							session.DoAction("recv", status)
						}
					} else {
						waitRecvCache = &ca
						break
					}
				}
			case action := <-session.do2:
				switch action.t {
				case "input":
					session.overTime = time.Now().Unix() + session.timeout
					args := action.args
					s := args[0].([]byte)
					n := args[1].(int)
					if n < 5 {
						//log.Println("recv reset")
						go session._Close(false)
						break
					} else if n == 5 {
						status, _ := makeDecode(s)
						if status == Reset || status == ResetAck {
							//log.Println("recv reset2", status)
							go session._Close(false)
						}
						break
					}
					ikcp.Ikcp_input(session.kcp, s, n)
					if waitRecvCache != nil {
						ca := *waitRecvCache
						tmp := session.processBuffer
						for {
							hr := ikcp.Ikcp_recv(session.kcp, tmp, ReadBufferSize)
							if hr > 0 {
								status := tmp[0]
								if status == Data {
									waitRecvCache = nil
									copy(ca.b, tmp[1:hr])
									ca.c <- int(hr - 1)
									break
								} else {
									session.DoAction("recv", status)
								}
							} else {
								break
							}
						}
					}
					updateF(10)
				case "write":
					b := action.args[0].([]byte)
					ikcp.Ikcp_send(session.kcp, b, len(b))
					updateF(10)
				}
			case <-session.quitChan:
				break out
			case <-updateC:
				now := uint32(iclock())
				ikcp.Ikcp_update(session.kcp, now)
				callUpdate = false
			}
		}
	}()
	select {
	case ping <- true:
	case <-session.quitChan:
	}
out:
	for {
		select {
		case <-session.quitChan:
			if session.disBind {
				break out
			}
		case action := <-session.do:
			switch action.t {
			case "quit":
				if session.closed {
					break
				}
				session._Close(true)
			case "closebegin":
				//A tell B to close and wait
				time.AfterFunc(time.Millisecond*500, func() {
					//log.Println("close over, step3", session.LocalAddr().String(), session.RemoteAddr().String())
					go session.DoAction("closeover")
				})
				buf := make([]byte, 5)
				go session.DoWrite(makeEncode(buf, Close, 0))
			case "closeover":
				//A call timeover
				close(session.closeChan)
				break out
			case "closeend":
				//B call close
				session._Close(false)
				break out
			case "recv":
				status := (action.args[0]).(byte)
				switch status {
				case CloseBack:
					//A call from B
					//log.Println("recv back close, step2", session.LocalAddr().String(), session.RemoteAddr().String())
					go session.DoAction("closeover")
				case Close:
					if session.closed {
						break
					}
					if session.status != "ok" {
						session._Close(false)
					} else {
						//log.Println("recv remote close, step1", session.LocalAddr().String(), session.RemoteAddr().String())
						buf := make([]byte, 5)
						go session.DoWrite(makeEncode(buf, CloseBack, 0))
						time.AfterFunc(time.Millisecond*500, func() {
							//log.Println("close remote over, step4", session.LocalAddr().String(), session.RemoteAddr().String())
							if session.closed {
								return
							}
							go session.DoAction("closeend")
						})
					}
				case Reset:
					//log.Println("recv reset")
					go session._Close(false)
				case Ping:
				default:
					if session.status != "ok" {
						session.Close()
					}
				}
			}
		}
	}
}

func (session *UDPMakeSession) _Close(bFirstCall bool) {
	if session.closed || session.noClose {
		return
	}
	session.closed = true
	//session.wait.Wait()
	go func() {
		//log.Println("pipe begin close", session.LocalAddr().String(), session.RemoteAddr().String())
		if bFirstCall {
			go session.DoAction("closebegin", session.LocalAddr().String())
			<-session.closeChan
		} else {
			close(session.closeChan)
		}
		//log.Println("pipe end close", session.id)
		close(session.quitChan)
		if session.listener != nil {
			session.listener.remove(session.remote.String())
		} else {
			if session.sock != nil && !session.disBind {
				session.sock.Close()
			}
		}
	}()
}

func (session *UDPMakeSession) DisBind() error {
	if session.disBind {
		return nil
	}
	session.disBind = true
	if session.listener == nil {
		session.sock.SetReadDeadline(time.Now())
	}
	if session.closed {
		return nil
	}
	close(session.quitChan)
	if session.listener != nil {
		//session.sock.Close()
		session.listener.remove(session.remote.String())
	}
	return nil
}
func (session *UDPMakeSession) Close() error {
	if session.closed {
		return nil
	}
	if session.status != "ok" {
		session._Close(false)
		return nil
	}
	go session.DoAction("quit")
	<-session.closeChan
	return nil
}

func (session *UDPMakeSession) DoAction2(action string, args ...interface{}) {
	//session.wait.Add(1)
	//log.Println(action, len(args))
	select {
	case session.do2 <- Action{t: action, args: args}:
	case <-session.quitChan:
		//session.wait.Done()
	}
}
func (session *UDPMakeSession) DoAction(action string, args ...interface{}) {
	//session.wait.Add(1)
	//log.Println(action, len(args))
	select {
	case session.do <- Action{t: action, args: args}:
	case <-session.quitChan:
		//session.wait.Done()
	}
}

func (session *UDPMakeSession) processInput(s string, n int) {
}

func (session *UDPMakeSession) LocalAddr() net.Addr {
	return session.sock.LocalAddr()
}

func (session *UDPMakeSession) RemoteAddr() net.Addr {
	return net.Addr(session.remote)
}

func (session *UDPMakeSession) SetDeadline(t time.Time) error {
	return session.sock.SetDeadline(t)
}

func (session *UDPMakeSession) SetReadDeadline(t time.Time) error {
	return session.sock.SetReadDeadline(t)
}

func (session *UDPMakeSession) SetWriteDeadline(t time.Time) error {
	return session.sock.SetWriteDeadline(t)
}

func (session *UDPMakeSession) DoWrite(s []byte) bool {
	wc := make(chan bool)
	select {
	case session.checkCanWrite <- wc:
		select {
		case <-wc:
		case <-session.quitChan:
			return false
		}
		session.DoAction2("write", s)
		return true
	case <-session.quitChan:
		return false
	}
}

func (session *UDPMakeSession) Write(b []byte) (n int, err error) {
	sendL := len(b)
	if sendL == 0 || session.status != "ok" {
		return 0, nil
	}
	data := make([]byte, sendL+1)
	data[0] = Data
	copy(data[1:], b)
	ok := session.DoWrite(data)
	if !ok {
		return 0, errors.New("closed")
	}
	return sendL, err
}

//udp read does not relay on the len(p), please make a big enough array to cache data
func (session *UDPMakeSession) Read(p []byte) (n int, err error) {
	wc := cache{p, 0, make(chan int)}
	select {
	case session.recvChan <- wc:
		select {
		case n = <-wc.c:
		case <-session.quitChan:
			n = -1
		}
	case <-session.quitChan:
		n = -1
	}
	if n == -1 {
		log.Println("force quit for read error", session.remote.String())
		return 0, errors.New("force quit for read error")
	} else {
		return n, nil
	}
}
<commit_msg>remove pipe<commit_after><|endoftext|>


<commit_before>package main

import (
	"log"

	"github.com/mndrix/golog"
)

func main() {
	m := golog.NewMachine().Consult(`
toki(jan_Kesi).
toki(jan_Pola).
toki(jan_Kesi, jan_Pola).
toki(jan_Kesi, toki_pona).
command(ilo_Kesi, toki(ziho, jan_Kesi)).
`)
	if m.CanProve(`toki(jan_Kesi).`) {
		log.Printf("toki(jan_Kesi). -> jan Kesi li toki.")
	}

	solutions := m.ProveAll(`toki(jan_Kesi, X).`)
	for _, solution := range solutions {
		log.Printf("jan_Kesi li toki e %s", solution.ByName_("X").String())
	}

	solutions = m.ProveAll(`command(X, toki(ziho, jan_Kesi)).`)
	for _, solution := range solutions {
		log.Printf("%s o, toki e jan_Kesi", solution.ByName_("X").String())
	}
}
<commit_msg>la-baujmi: glue a repl together<commit_after>package main

import (
	"log"
	"os"
	"strings"

	"github.com/Xe/x/web/tokiponatokens"
	"github.com/mndrix/golog"
	line "github.com/peterh/liner"
)

func main() {
	/*
			m := golog.NewMachine().Consult(`
		toki(jan_Kesi).
		toki(jan_Pola).
		toki(jan_Kesi, jan_Pola).
		toki(jan_Kesi, toki_pona).
		command(ilo_Kesi, toki(ziho, jan_Kesi)).
		`)
			if m.CanProve(`toki(jan_Kesi).`) {
				log.Printf("toki(jan_Kesi). -> jan Kesi li toki.")
			}

			solutions := m.ProveAll(`toki(jan_Kesi, X).`)
			for _, solution := range solutions {
				log.Printf("jan_Kesi li toki e %s", solution.ByName_("X").String())
			}

			solutions = m.ProveAll(`command(X, toki(ziho, jan_Kesi)).`)
			for _, solution := range solutions {
				log.Printf("%s o, toki e jan_Kesi", solution.ByName_("X").String())
			}
	*/

	var m golog.Machine
	m = golog.NewMachine()

	for {
		l := line.NewLiner()
		if inp, err := l.Prompt("|: "); err == nil {
			if inp == "" {
				return
			}

			l.AppendHistory(inp)

			parts, err := tokiponatokens.Tokenize(tokiPonaAPIURL, inp)
			if err != nil {
				log.Printf("error: %v", err)
				continue
			}

			for _, sentence := range parts {
				sbs, err := SentenceToSelbris(sentence)
				if err != nil {
					log.Printf("can't derive facts: %v", err)
					continue
				}

				for _, sb := range sbs {
					f := sb.Fact()

					if strings.Contains(inp, "?") {
						log.Printf("Query: %s", f)

						solutions := m.ProveAll(f)
						for _, solution := range solutions {
							log.Println("found", strings.Replace(f, "A", solution.ByName_("A").String(), -1))
						}

						continue
					}

					log.Printf("registering fact: %s", f)
					m = m.Consult(f)
				}
			}
		} else if err == line.ErrPromptAborted {
			log.Print("Aborted")
			break
		} else {
			log.Print("Error reading line: ", err)
			break
		}
	}

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


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

package ns

import (
	"crypto/rand"
	"fmt"
	"os"
	"path"
	"runtime"
	"sync"

	"golang.org/x/sys/unix"
)

type NetNS interface {
	// Executes the passed closure in this object's network namespace,
	// attemtping to restore the original namespace before returning.
	// However, since each OS thread can have a different network namespace,
	// and Go's thread scheduling is highly variable, callers cannot
	// guarantee any specific namespace is set unless operations that
	// require that namespace are wrapped with Do().  Also, no code called
	// from Do() should call runtime.UnlockOSThread(), or the risk
	// of executing code in an incorrect namespace will be greater.  See
	// https://github.com/golang/go/wiki/LockOSThread for further details.
	Do(toRun func(NetNS) error) error

	// Sets the current network namespace to this object's network namespace.
	// Note that since Go's thread scheduling is highly variable, callers
	// cannot guarantee the requested namespace will be the current namespace
	// after this function is called; to ensure this wrap operations that
	// require the namespace with Do() instead.
	Set() error

	// Returns the filesystem path representing this object's network namespace
	Path() string

	// Returns a file descriptor representing this object's network namespace
	Fd() uintptr

	// Cleans up this instance of the network namespace; if this instance
	// is the last user the namespace will be destroyed
	Close() error
}

type netNS struct {
	file    *os.File
	mounted bool
	closed  bool
}

func getCurrentThreadNetNSPath() string {
	// /proc/self/ns/net returns the namespace of the main thread, not
	// of whatever thread this goroutine is running on.  Make sure we
	// use the thread's net namespace since the thread is switching around
	return fmt.Sprintf("/proc/%d/task/%d/ns/net", os.Getpid(), unix.Gettid())
}

// Returns an object representing the current OS thread's network namespace
func GetCurrentNS() (NetNS, error) {
	return GetNS(getCurrentThreadNetNSPath())
}

// Returns an object representing the namespace referred to by @path
func GetNS(nspath string) (NetNS, error) {
	fd, err := os.Open(nspath)
	if err != nil {
		return nil, err
	}
	return &netNS{file: fd}, nil
}

// Creates a new persistent network namespace and returns an object
// representing that namespace, without switching to it
func NewNS() (NetNS, error) {
	const nsRunDir = "/var/run/netns"

	b := make([]byte, 16)
	_, err := rand.Reader.Read(b)
	if err != nil {
		return nil, fmt.Errorf("failed to generate random netns name: %v", err)
	}

	err = os.MkdirAll(nsRunDir, 0755)
	if err != nil {
		return nil, err
	}

	// create an empty file at the mount point
	nsName := fmt.Sprintf("cni-%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
	nsPath := path.Join(nsRunDir, nsName)
	mountPointFd, err := os.Create(nsPath)
	if err != nil {
		return nil, err
	}
	mountPointFd.Close()

	// Ensure the mount point is cleaned up on errors; if the namespace
	// was successfully mounted this will have no effect because the file
	// is in-use
	defer os.RemoveAll(nsPath)

	var wg sync.WaitGroup
	wg.Add(1)

	// do namespace work in a dedicated goroutine, so that we can safely
	// Lock/Unlock OSThread without upsetting the lock/unlock state of
	// the caller of this function
	var fd *os.File
	go (func() {
		defer wg.Done()
		runtime.LockOSThread()

		var origNS NetNS
		origNS, err = GetNS(getCurrentThreadNetNSPath())
		if err != nil {
			return
		}
		defer origNS.Close()

		// create a new netns on the current thread
		err = unix.Unshare(unix.CLONE_NEWNET)
		if err != nil {
			return
		}
		defer origNS.Set()

		// bind mount the new netns from the current thread onto the mount point
		err = unix.Mount(getCurrentThreadNetNSPath(), nsPath, "none", unix.MS_BIND, "")
		if err != nil {
			return
		}

		fd, err = os.Open(nsPath)
		if err != nil {
			return
		}
	})()
	wg.Wait()

	if err != nil {
		unix.Unmount(nsPath, unix.MNT_DETACH)
		return nil, fmt.Errorf("failed to create namespace: %v", err)
	}

	return &netNS{file: fd, mounted: true}, nil
}

func (ns *netNS) Path() string {
	return ns.file.Name()
}

func (ns *netNS) Fd() uintptr {
	return ns.file.Fd()
}

func (ns *netNS) errorIfClosed() error {
	if ns.closed {
		return fmt.Errorf("%q has already been closed", ns.file.Name())
	}
	return nil
}

func (ns *netNS) Close() error {
	if err := ns.errorIfClosed(); err != nil {
		return err
	}

	if err := ns.file.Close(); err != nil {
		return fmt.Errorf("Failed to close %q: %v", ns.file.Name(), err)
	}
	ns.closed = true

	if ns.mounted {
		if err := unix.Unmount(ns.file.Name(), unix.MNT_DETACH); err != nil {
			return fmt.Errorf("Failed to unmount namespace %s: %v", ns.file.Name(), err)
		}
		if err := os.RemoveAll(ns.file.Name()); err != nil {
			return fmt.Errorf("Failed to clean up namespace %s: %v", ns.file.Name(), err)
		}
		ns.mounted = false
	}

	return nil
}

func (ns *netNS) Do(toRun func(NetNS) error) error {
	if err := ns.errorIfClosed(); err != nil {
		return err
	}

	containedCall := func(hostNS NetNS) error {
		threadNS, err := GetNS(getCurrentThreadNetNSPath())
		if err != nil {
			return fmt.Errorf("failed to open current netns: %v", err)
		}
		defer threadNS.Close()

		// switch to target namespace
		if err = ns.Set(); err != nil {
			return fmt.Errorf("error switching to ns %v: %v", ns.file.Name(), err)
		}
		defer threadNS.Set() // switch back

		return toRun(hostNS)
	}

	// save a handle to current network namespace
	hostNS, err := GetNS(getCurrentThreadNetNSPath())
	if err != nil {
		return fmt.Errorf("Failed to open current namespace: %v", err)
	}
	defer hostNS.Close()

	var wg sync.WaitGroup
	wg.Add(1)

	var innerError error
	go func() {
		defer wg.Done()
		runtime.LockOSThread()
		innerError = containedCall(hostNS)
	}()
	wg.Wait()

	return innerError
}

func (ns *netNS) Set() error {
	if err := ns.errorIfClosed(); err != nil {
		return err
	}

	if _, _, err := unix.Syscall(unix.SYS_SETNS, ns.Fd(), uintptr(unix.CLONE_NEWNET), 0); err != 0 {
		return fmt.Errorf("Error switching to ns %v: %v", ns.file.Name(), err)
	}

	return nil
}

// WithNetNSPath executes the passed closure under the given network
// namespace, restoring the original namespace afterwards.
func WithNetNSPath(nspath string, toRun func(NetNS) error) error {
	ns, err := GetNS(nspath)
	if err != nil {
		return fmt.Errorf("Failed to open %v: %v", nspath, err)
	}
	defer ns.Close()
	return ns.Do(toRun)
}
<commit_msg>pkg/ns: verify netns when initialized with GetNS<commit_after>// Copyright 2015 CNI authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ns

import (
	"crypto/rand"
	"fmt"
	"os"
	"path"
	"runtime"
	"sync"
	"syscall"

	"golang.org/x/sys/unix"
)

type NetNS interface {
	// Executes the passed closure in this object's network namespace,
	// attemtping to restore the original namespace before returning.
	// However, since each OS thread can have a different network namespace,
	// and Go's thread scheduling is highly variable, callers cannot
	// guarantee any specific namespace is set unless operations that
	// require that namespace are wrapped with Do().  Also, no code called
	// from Do() should call runtime.UnlockOSThread(), or the risk
	// of executing code in an incorrect namespace will be greater.  See
	// https://github.com/golang/go/wiki/LockOSThread for further details.
	Do(toRun func(NetNS) error) error

	// Sets the current network namespace to this object's network namespace.
	// Note that since Go's thread scheduling is highly variable, callers
	// cannot guarantee the requested namespace will be the current namespace
	// after this function is called; to ensure this wrap operations that
	// require the namespace with Do() instead.
	Set() error

	// Returns the filesystem path representing this object's network namespace
	Path() string

	// Returns a file descriptor representing this object's network namespace
	Fd() uintptr

	// Cleans up this instance of the network namespace; if this instance
	// is the last user the namespace will be destroyed
	Close() error
}

type netNS struct {
	file    *os.File
	mounted bool
	closed  bool
}

func getCurrentThreadNetNSPath() string {
	// /proc/self/ns/net returns the namespace of the main thread, not
	// of whatever thread this goroutine is running on.  Make sure we
	// use the thread's net namespace since the thread is switching around
	return fmt.Sprintf("/proc/%d/task/%d/ns/net", os.Getpid(), unix.Gettid())
}

// Returns an object representing the current OS thread's network namespace
func GetCurrentNS() (NetNS, error) {
	return GetNS(getCurrentThreadNetNSPath())
}

// Returns an object representing the namespace referred to by @path
func GetNS(nspath string) (NetNS, error) {
	fd, err := os.Open(nspath)
	if err != nil {
		return nil, fmt.Errorf("Failed to open %v: %v", nspath, err)
	}

	isNSFS, err := IsNSFS(nspath)
	if err != nil {
		fd.Close()
		return nil, err
	}
	if !isNSFS {
		fd.Close()
		return nil, fmt.Errorf("%v is not of type NSFS", nspath)
	}

	return &netNS{file: fd}, nil
}

// Returns whether or not the nspath argument points to a network namespace
func IsNSFS(nspath string) (bool, error) {
	const NSFS_MAGIC = 0x6e736673

	stat := syscall.Statfs_t{}
	if err := syscall.Statfs(nspath, &stat); err != nil {
		return false, fmt.Errorf("failed to Statfs %q: %v", nspath, err)
	}

	return stat.Type == NSFS_MAGIC, nil
}

// Creates a new persistent network namespace and returns an object
// representing that namespace, without switching to it
func NewNS() (NetNS, error) {
	const nsRunDir = "/var/run/netns"

	b := make([]byte, 16)
	_, err := rand.Reader.Read(b)
	if err != nil {
		return nil, fmt.Errorf("failed to generate random netns name: %v", err)
	}

	err = os.MkdirAll(nsRunDir, 0755)
	if err != nil {
		return nil, err
	}

	// create an empty file at the mount point
	nsName := fmt.Sprintf("cni-%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
	nsPath := path.Join(nsRunDir, nsName)
	mountPointFd, err := os.Create(nsPath)
	if err != nil {
		return nil, err
	}
	mountPointFd.Close()

	// Ensure the mount point is cleaned up on errors; if the namespace
	// was successfully mounted this will have no effect because the file
	// is in-use
	defer os.RemoveAll(nsPath)

	var wg sync.WaitGroup
	wg.Add(1)

	// do namespace work in a dedicated goroutine, so that we can safely
	// Lock/Unlock OSThread without upsetting the lock/unlock state of
	// the caller of this function
	var fd *os.File
	go (func() {
		defer wg.Done()
		runtime.LockOSThread()

		var origNS NetNS
		origNS, err = GetNS(getCurrentThreadNetNSPath())
		if err != nil {
			return
		}
		defer origNS.Close()

		// create a new netns on the current thread
		err = unix.Unshare(unix.CLONE_NEWNET)
		if err != nil {
			return
		}
		defer origNS.Set()

		// bind mount the new netns from the current thread onto the mount point
		err = unix.Mount(getCurrentThreadNetNSPath(), nsPath, "none", unix.MS_BIND, "")
		if err != nil {
			return
		}

		fd, err = os.Open(nsPath)
		if err != nil {
			return
		}
	})()
	wg.Wait()

	if err != nil {
		unix.Unmount(nsPath, unix.MNT_DETACH)
		return nil, fmt.Errorf("failed to create namespace: %v", err)
	}

	return &netNS{file: fd, mounted: true}, nil
}

func (ns *netNS) Path() string {
	return ns.file.Name()
}

func (ns *netNS) Fd() uintptr {
	return ns.file.Fd()
}

func (ns *netNS) errorIfClosed() error {
	if ns.closed {
		return fmt.Errorf("%q has already been closed", ns.file.Name())
	}
	return nil
}

func (ns *netNS) Close() error {
	if err := ns.errorIfClosed(); err != nil {
		return err
	}

	if err := ns.file.Close(); err != nil {
		return fmt.Errorf("Failed to close %q: %v", ns.file.Name(), err)
	}
	ns.closed = true

	if ns.mounted {
		if err := unix.Unmount(ns.file.Name(), unix.MNT_DETACH); err != nil {
			return fmt.Errorf("Failed to unmount namespace %s: %v", ns.file.Name(), err)
		}
		if err := os.RemoveAll(ns.file.Name()); err != nil {
			return fmt.Errorf("Failed to clean up namespace %s: %v", ns.file.Name(), err)
		}
		ns.mounted = false
	}

	return nil
}

func (ns *netNS) Do(toRun func(NetNS) error) error {
	if err := ns.errorIfClosed(); err != nil {
		return err
	}

	containedCall := func(hostNS NetNS) error {
		threadNS, err := GetNS(getCurrentThreadNetNSPath())
		if err != nil {
			return fmt.Errorf("failed to open current netns: %v", err)
		}
		defer threadNS.Close()

		// switch to target namespace
		if err = ns.Set(); err != nil {
			return fmt.Errorf("error switching to ns %v: %v", ns.file.Name(), err)
		}
		defer threadNS.Set() // switch back

		return toRun(hostNS)
	}

	// save a handle to current network namespace
	hostNS, err := GetNS(getCurrentThreadNetNSPath())
	if err != nil {
		return fmt.Errorf("Failed to open current namespace: %v", err)
	}
	defer hostNS.Close()

	var wg sync.WaitGroup
	wg.Add(1)

	var innerError error
	go func() {
		defer wg.Done()
		runtime.LockOSThread()
		innerError = containedCall(hostNS)
	}()
	wg.Wait()

	return innerError
}

func (ns *netNS) Set() error {
	if err := ns.errorIfClosed(); err != nil {
		return err
	}

	if _, _, err := unix.Syscall(unix.SYS_SETNS, ns.Fd(), uintptr(unix.CLONE_NEWNET), 0); err != 0 {
		return fmt.Errorf("Error switching to ns %v: %v", ns.file.Name(), err)
	}

	return nil
}

// WithNetNSPath executes the passed closure under the given network
// namespace, restoring the original namespace afterwards.
func WithNetNSPath(nspath string, toRun func(NetNS) error) error {
	ns, err := GetNS(nspath)
	if err != nil {
		return err
	}
	defer ns.Close()
	return ns.Do(toRun)
}
<|endoftext|>


<commit_before>// Copyright 2017 The casbin Authors. 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 xormadapter

import (
	"errors"
	"runtime"

	"github.com/casbin/casbin/model"
	"github.com/casbin/casbin/persist"
	"github.com/go-xorm/xorm"
	"github.com/lib/pq"
)

type CasbinRule struct {
	PType string `xorm:"varchar(100) index"`
	V0    string `xorm:"varchar(100) index"`
	V1    string `xorm:"varchar(100) index"`
	V2    string `xorm:"varchar(100) index"`
	V3    string `xorm:"varchar(100) index"`
	V4    string `xorm:"varchar(100) index"`
	V5    string `xorm:"varchar(100) index"`
}

// Adapter represents the Xorm adapter for policy storage.
type Adapter struct {
	driverName     string
	dataSourceName string
	dbSpecified    bool
	engine         *xorm.Engine
}

// finalizer is the destructor for Adapter.
func finalizer(a *Adapter) {
	a.engine.Close()
}

// NewAdapter is the constructor for Adapter.
// dbSpecified is an optional bool parameter. The default value is false.
// It's up to whether you have specified an existing DB in dataSourceName.
// If dbSpecified == true, you need to make sure the DB in dataSourceName exists.
// If dbSpecified == false, the adapter will automatically create a DB named "casbin".
func NewAdapter(driverName string, dataSourceName string, dbSpecified ...bool) *Adapter {
	a := &Adapter{}
	a.driverName = driverName
	a.dataSourceName = dataSourceName

	if len(dbSpecified) == 0 {
		a.dbSpecified = false
	} else if len(dbSpecified) == 1 {
		a.dbSpecified = dbSpecified[0]
	} else {
		panic(errors.New("invalid parameter: dbSpecified"))
	}

	// Open the DB, create it if not existed.
	a.open()

	// Call the destructor when the object is released.
	runtime.SetFinalizer(a, finalizer)

	return a
}

func (a *Adapter) createDatabase() error {
	var err error
	var engine *xorm.Engine
	if a.driverName == "postgres" {
		engine, err = xorm.NewEngine(a.driverName, a.dataSourceName+" dbname=postgres")
	} else {
		engine, err = xorm.NewEngine(a.driverName, a.dataSourceName)
	}
	if err != nil {
		return err
	}
	defer engine.Close()

	if a.driverName == "postgres" {
		if _, err = engine.Exec("CREATE DATABASE casbin"); err != nil {
			// 42P04 is	duplicate_database
			if err.(*pq.Error).Code == "42P04" {
				return nil
			}
		}
	} else {
		_, err = engine.Exec("CREATE DATABASE IF NOT EXISTS casbin")
	}
	return err
}

func (a *Adapter) open() {
	var err error
	var engine *xorm.Engine

	if a.dbSpecified {
		engine, err = xorm.NewEngine(a.driverName, a.dataSourceName)
		if err != nil {
			panic(err)
		}
	} else {
		if err = a.createDatabase(); err != nil {
			panic(err)
		}

		if a.driverName == "postgres" {
			engine, err = xorm.NewEngine(a.driverName, a.dataSourceName+" dbname=casbin")
		} else {
			engine, err = xorm.NewEngine(a.driverName, a.dataSourceName+"casbin")
		}
		if err != nil {
			panic(err)
		}
	}

	a.engine = engine

	a.createTable()
}

func (a *Adapter) close() {
	a.engine.Close()
	a.engine = nil
}

func (a *Adapter) createTable() {
	err := a.engine.Sync2(new(CasbinRule))
	if err != nil {
		panic(err)
	}
}

func (a *Adapter) dropTable() {
	err := a.engine.DropTables(new(CasbinRule))
	if err != nil {
		panic(err)
	}
}

func loadPolicyLine(line CasbinRule, model model.Model) {
	lineText := line.PType
	if line.V0 != "" {
		lineText += ", " + line.V0
	}
	if line.V1 != "" {
		lineText += ", " + line.V1
	}
	if line.V2 != "" {
		lineText += ", " + line.V2
	}
	if line.V3 != "" {
		lineText += ", " + line.V3
	}
	if line.V4 != "" {
		lineText += ", " + line.V4
	}
	if line.V5 != "" {
		lineText += ", " + line.V5
	}

	persist.LoadPolicyLine(lineText, model)
}

// LoadPolicy loads policy from database.
func (a *Adapter) LoadPolicy(model model.Model) error {
	var lines []CasbinRule
	err := a.engine.Find(&lines)
	if err != nil {
		return err
	}

	for _, line := range lines {
		loadPolicyLine(line, model)
	}

	return nil
}

func savePolicyLine(ptype string, rule []string) CasbinRule {
	line := CasbinRule{}

	line.PType = ptype
	if len(rule) > 0 {
		line.V0 = rule[0]
	}
	if len(rule) > 1 {
		line.V1 = rule[1]
	}
	if len(rule) > 2 {
		line.V2 = rule[2]
	}
	if len(rule) > 3 {
		line.V3 = rule[3]
	}
	if len(rule) > 4 {
		line.V4 = rule[4]
	}
	if len(rule) > 5 {
		line.V5 = rule[5]
	}

	return line
}

// SavePolicy saves policy to database.
func (a *Adapter) SavePolicy(model model.Model) error {
	a.dropTable()
	a.createTable()

	var lines []CasbinRule

	for ptype, ast := range model["p"] {
		for _, rule := range ast.Policy {
			line := savePolicyLine(ptype, rule)
			lines = append(lines, line)
		}
	}

	for ptype, ast := range model["g"] {
		for _, rule := range ast.Policy {
			line := savePolicyLine(ptype, rule)
			lines = append(lines, line)
		}
	}

	_, err := a.engine.Insert(&lines)
	return err
}

// AddPolicy adds a policy rule to the storage.
func (a *Adapter) AddPolicy(sec string, ptype string, rule []string) error {
	line := savePolicyLine(ptype, rule)
	_, err := a.engine.Insert(line)
	return err
}

// RemovePolicy removes a policy rule from the storage.
func (a *Adapter) RemovePolicy(sec string, ptype string, rule []string) error {
	line := savePolicyLine(ptype, rule)
	_, err := a.engine.Delete(line)
	return err
}

// RemoveFilteredPolicy removes policy rules that match the filter from the storage.
func (a *Adapter) RemoveFilteredPolicy(sec string, ptype string, fieldIndex int, fieldValues ...string) error {
	line := CasbinRule{}

	line.PType = ptype
	if fieldIndex <= 0 && 0 < fieldIndex + len(fieldValues) {
		line.V0 = fieldValues[0 - fieldIndex]
	}
	if fieldIndex <= 1 && 1 < fieldIndex + len(fieldValues) {
		line.V1 = fieldValues[1 - fieldIndex]
	}
	if fieldIndex <= 2 && 2 < fieldIndex + len(fieldValues) {
		line.V2 = fieldValues[2 - fieldIndex]
	}
	if fieldIndex <= 3 && 3 < fieldIndex + len(fieldValues) {
		line.V3 = fieldValues[3 - fieldIndex]
	}
	if fieldIndex <= 4 && 4 < fieldIndex + len(fieldValues) {
		line.V4 = fieldValues[4 - fieldIndex]
	}
	if fieldIndex <= 5 && 5 < fieldIndex + len(fieldValues) {
		line.V5 = fieldValues[5 - fieldIndex]
	}

	_, err := a.engine.Delete(line)
	return err
}
<commit_msg>Handle PostgreSQL duplicated db error better, fix: https://github.com/casbin/xorm-adapter/issues/8<commit_after>// Copyright 2017 The casbin Authors. 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 xormadapter

import (
	"errors"
	"runtime"

	"github.com/casbin/casbin/model"
	"github.com/casbin/casbin/persist"
	"github.com/go-xorm/xorm"
	"github.com/lib/pq"
)

type CasbinRule struct {
	PType string `xorm:"varchar(100) index"`
	V0    string `xorm:"varchar(100) index"`
	V1    string `xorm:"varchar(100) index"`
	V2    string `xorm:"varchar(100) index"`
	V3    string `xorm:"varchar(100) index"`
	V4    string `xorm:"varchar(100) index"`
	V5    string `xorm:"varchar(100) index"`
}

// Adapter represents the Xorm adapter for policy storage.
type Adapter struct {
	driverName     string
	dataSourceName string
	dbSpecified    bool
	engine         *xorm.Engine
}

// finalizer is the destructor for Adapter.
func finalizer(a *Adapter) {
	a.engine.Close()
}

// NewAdapter is the constructor for Adapter.
// dbSpecified is an optional bool parameter. The default value is false.
// It's up to whether you have specified an existing DB in dataSourceName.
// If dbSpecified == true, you need to make sure the DB in dataSourceName exists.
// If dbSpecified == false, the adapter will automatically create a DB named "casbin".
func NewAdapter(driverName string, dataSourceName string, dbSpecified ...bool) *Adapter {
	a := &Adapter{}
	a.driverName = driverName
	a.dataSourceName = dataSourceName

	if len(dbSpecified) == 0 {
		a.dbSpecified = false
	} else if len(dbSpecified) == 1 {
		a.dbSpecified = dbSpecified[0]
	} else {
		panic(errors.New("invalid parameter: dbSpecified"))
	}

	// Open the DB, create it if not existed.
	a.open()

	// Call the destructor when the object is released.
	runtime.SetFinalizer(a, finalizer)

	return a
}

func (a *Adapter) createDatabase() error {
	var err error
	var engine *xorm.Engine
	if a.driverName == "postgres" {
		engine, err = xorm.NewEngine(a.driverName, a.dataSourceName+" dbname=postgres")
	} else {
		engine, err = xorm.NewEngine(a.driverName, a.dataSourceName)
	}
	if err != nil {
		return err
	}
	defer engine.Close()

	if a.driverName == "postgres" {
		if _, err = engine.Exec("CREATE DATABASE casbin"); err != nil {
			// 42P04 is	duplicate_database
			if pqerr, ok := err.(*pq.Error); ok && pqerr.Code == "42P04" {
				return nil
			}
		}
	} else {
		_, err = engine.Exec("CREATE DATABASE IF NOT EXISTS casbin")
	}
	return err
}

func (a *Adapter) open() {
	var err error
	var engine *xorm.Engine

	if a.dbSpecified {
		engine, err = xorm.NewEngine(a.driverName, a.dataSourceName)
		if err != nil {
			panic(err)
		}
	} else {
		if err = a.createDatabase(); err != nil {
			panic(err)
		}

		if a.driverName == "postgres" {
			engine, err = xorm.NewEngine(a.driverName, a.dataSourceName+" dbname=casbin")
		} else {
			engine, err = xorm.NewEngine(a.driverName, a.dataSourceName+"casbin")
		}
		if err != nil {
			panic(err)
		}
	}

	a.engine = engine

	a.createTable()
}

func (a *Adapter) close() {
	a.engine.Close()
	a.engine = nil
}

func (a *Adapter) createTable() {
	err := a.engine.Sync2(new(CasbinRule))
	if err != nil {
		panic(err)
	}
}

func (a *Adapter) dropTable() {
	err := a.engine.DropTables(new(CasbinRule))
	if err != nil {
		panic(err)
	}
}

func loadPolicyLine(line CasbinRule, model model.Model) {
	lineText := line.PType
	if line.V0 != "" {
		lineText += ", " + line.V0
	}
	if line.V1 != "" {
		lineText += ", " + line.V1
	}
	if line.V2 != "" {
		lineText += ", " + line.V2
	}
	if line.V3 != "" {
		lineText += ", " + line.V3
	}
	if line.V4 != "" {
		lineText += ", " + line.V4
	}
	if line.V5 != "" {
		lineText += ", " + line.V5
	}

	persist.LoadPolicyLine(lineText, model)
}

// LoadPolicy loads policy from database.
func (a *Adapter) LoadPolicy(model model.Model) error {
	var lines []CasbinRule
	err := a.engine.Find(&lines)
	if err != nil {
		return err
	}

	for _, line := range lines {
		loadPolicyLine(line, model)
	}

	return nil
}

func savePolicyLine(ptype string, rule []string) CasbinRule {
	line := CasbinRule{}

	line.PType = ptype
	if len(rule) > 0 {
		line.V0 = rule[0]
	}
	if len(rule) > 1 {
		line.V1 = rule[1]
	}
	if len(rule) > 2 {
		line.V2 = rule[2]
	}
	if len(rule) > 3 {
		line.V3 = rule[3]
	}
	if len(rule) > 4 {
		line.V4 = rule[4]
	}
	if len(rule) > 5 {
		line.V5 = rule[5]
	}

	return line
}

// SavePolicy saves policy to database.
func (a *Adapter) SavePolicy(model model.Model) error {
	a.dropTable()
	a.createTable()

	var lines []CasbinRule

	for ptype, ast := range model["p"] {
		for _, rule := range ast.Policy {
			line := savePolicyLine(ptype, rule)
			lines = append(lines, line)
		}
	}

	for ptype, ast := range model["g"] {
		for _, rule := range ast.Policy {
			line := savePolicyLine(ptype, rule)
			lines = append(lines, line)
		}
	}

	_, err := a.engine.Insert(&lines)
	return err
}

// AddPolicy adds a policy rule to the storage.
func (a *Adapter) AddPolicy(sec string, ptype string, rule []string) error {
	line := savePolicyLine(ptype, rule)
	_, err := a.engine.Insert(line)
	return err
}

// RemovePolicy removes a policy rule from the storage.
func (a *Adapter) RemovePolicy(sec string, ptype string, rule []string) error {
	line := savePolicyLine(ptype, rule)
	_, err := a.engine.Delete(line)
	return err
}

// RemoveFilteredPolicy removes policy rules that match the filter from the storage.
func (a *Adapter) RemoveFilteredPolicy(sec string, ptype string, fieldIndex int, fieldValues ...string) error {
	line := CasbinRule{}

	line.PType = ptype
	if fieldIndex <= 0 && 0 < fieldIndex + len(fieldValues) {
		line.V0 = fieldValues[0 - fieldIndex]
	}
	if fieldIndex <= 1 && 1 < fieldIndex + len(fieldValues) {
		line.V1 = fieldValues[1 - fieldIndex]
	}
	if fieldIndex <= 2 && 2 < fieldIndex + len(fieldValues) {
		line.V2 = fieldValues[2 - fieldIndex]
	}
	if fieldIndex <= 3 && 3 < fieldIndex + len(fieldValues) {
		line.V3 = fieldValues[3 - fieldIndex]
	}
	if fieldIndex <= 4 && 4 < fieldIndex + len(fieldValues) {
		line.V4 = fieldValues[4 - fieldIndex]
	}
	if fieldIndex <= 5 && 5 < fieldIndex + len(fieldValues) {
		line.V5 = fieldValues[5 - fieldIndex]
	}

	_, err := a.engine.Delete(line)
	return err
}
<|endoftext|>