<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/urfave/cli"
)

type AccessPointState struct {
	Enabled bool               `json:"enabled"`
	Configg *AccessPointConfig `json:"config"`
}

func ApCMD(ctx *cli.Context) error {
	if ctx.IsSet(enableFlag) {
		return EnableApCMD(ctx)
	}
	if ctx.IsSet(disableFlag) {
		return DisableApCMD(ctx)
	}
	if ctx.IsSet(removeFlag) {
		return RemoveApCMD(ctx)
	}
	if ctx.IsSet(configFlag) {
		return ConfigApCMD(ctx)
	}
	return nil
}

func ConfigApCMD(ctx *cli.Context) error {
	base := ctx.String("dir")
	name := ctx.String("name")
	src := ctx.String("config")
	if src == "" {
		return errors.New("fconf: missing argument")
	}
	var b []byte
	var err error
	if src == "stdin" {
		b, err = ReadFromStdin()
		if err != nil {
			return err
		}
	} else {
		b, err = ioutil.ReadFile(src)
		if err != nil {
			return err
		}
	}
	e := &AccessPointConfig{}
	err = json.Unmarshal(b, e)
	if err != nil {
		return err
	}
	err = checkDir(base)
	if err != nil {
		return err
	}
	ap := DefaultAccesPoint()
	ap.Update(e)
	if strings.Contains(name, "%s") {
		name = fmt.Sprintf(name, ap.WifiIface)
	}
	filename := filepath.Join(base, name)
	var buf bytes.Buffer
	_, err = ap.WriteTo(&buf)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(filename, buf.Bytes(), 0644)
	if err != nil {
		return err
	}
	fmt.Printf("successful written access point configuration to %s \n", filename)
	state := &AccessPointState{Configg: ap.State()}
	as, err := accessPointState(ap.WifiIface)
	if err == nil {
		state.Enabled = as.Enabled
	}
	data, err := json.Marshal(state)
	if err != nil {
		return err
	}
	ctx.GlobalSet("interface", ap.WifiIface)
	return keepState(
		fmt.Sprintf(defaultAccessPointConfig, ap.WifiIface), data)
}

func accessPointState(i string) (*AccessPointState, error) {
	dir := os.Getenv("FCONF_CONFIGDIR")
	if dir == "" {
		dir = fconfConfigDir
	}
	b, err := ioutil.ReadFile(filepath.Join(dir,
		fmt.Sprintf(defaultAccessPointConfig, i)))
	if err != nil {
		return nil, err
	}
	a := &AccessPointState{}
	err = json.Unmarshal(b, a)
	if err != nil {
		return nil, err
	}
	if a.Configg == nil {
		return nil, ErrWrongStateFile
	}
	return a, nil
}

func EnableApCMD(ctx *cli.Context) error {
	if ctx.IsSet(configFlag) {
		err := ConfigApCMD(ctx)
		if err != nil {
			return err
		}
	}
	i := getInterface(ctx)
	if i == "" {
		return errors.New("missing interface, you must specify interface")
	}
	state, err := accessPointState(i)
	if err != nil {
		return err
	}
	service := "create_ap"
	err = startService(service)
	if err != nil {
		return err
	}
	err = enableService(service)
	if err != nil {
		return err
	}
	state.Enabled = true
	data, err := json.Marshal(state)
	if err != nil {
		return err
	}
	return keepState(
		fmt.Sprintf(defaultAccessPointConfig, i), data)
}

func DisableApCMD(ctx *cli.Context) error {
	if ctx.IsSet(configFlag) {
		fmt.Println("WARN: config flag will be ignored when diable flag is used")
	}
	i := getInterface(ctx)
	if i == "" {
		return errors.New("missing interface, you must specify interface")
	}
	state, err := accessPointState(i)
	if err != nil {
		return err
	}
	service := "create_ap"
	err = stopService(service)
	if err != nil {
		return err
	}
	err = disableService(service)
	if err != nil {
		return err
	}
	state.Enabled = false
	data, err := json.Marshal(state)
	if err != nil {
		return err
	}
	return keepState(
		fmt.Sprintf(defaultAccessPointConfig, i), data)
}

func RemoveApCMD(ctx *cli.Context) error {
	i := getInterface(ctx)
	if i == "" {
		return errors.New("missing interface, you must specify interface")
	}
	a, err := accessPointState(i)
	if err != nil {
		return err
	}
	if a.Enabled {
		err = DisableApCMD(ctx)
		if err != nil {
			return err
		}
	}

	// remove the state file
	stateFile := filepath.Join(stateDir(),
		fmt.Sprintf(defaultAccessPointConfig, i))
	err = removeFile(stateFile)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
	}
	return nil
}
<commit_msg>Add interface to create_ap service<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/urfave/cli"
)

type AccessPointState struct {
	Enabled bool               `json:"enabled"`
	Configg *AccessPointConfig `json:"config"`
}

func ApCMD(ctx *cli.Context) error {
	if ctx.IsSet(enableFlag) {
		return EnableApCMD(ctx)
	}
	if ctx.IsSet(disableFlag) {
		return DisableApCMD(ctx)
	}
	if ctx.IsSet(removeFlag) {
		return RemoveApCMD(ctx)
	}
	if ctx.IsSet(configFlag) {
		return ConfigApCMD(ctx)
	}
	return nil
}

func ConfigApCMD(ctx *cli.Context) error {
	base := ctx.String("dir")
	name := ctx.String("name")
	src := ctx.String("config")
	if src == "" {
		return errors.New("fconf: missing argument")
	}
	var b []byte
	var err error
	if src == "stdin" {
		b, err = ReadFromStdin()
		if err != nil {
			return err
		}
	} else {
		b, err = ioutil.ReadFile(src)
		if err != nil {
			return err
		}
	}
	e := &AccessPointConfig{}
	err = json.Unmarshal(b, e)
	if err != nil {
		return err
	}
	err = checkDir(base)
	if err != nil {
		return err
	}
	ap := DefaultAccesPoint()
	ap.Update(e)
	if strings.Contains(name, "%s") {
		name = fmt.Sprintf(name, ap.WifiIface)
	}
	filename := filepath.Join(base, name)
	var buf bytes.Buffer
	_, err = ap.WriteTo(&buf)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(filename, buf.Bytes(), 0644)
	if err != nil {
		return err
	}
	fmt.Printf("successful written access point configuration to %s \n", filename)
	state := &AccessPointState{Configg: ap.State()}
	as, err := accessPointState(ap.WifiIface)
	if err == nil {
		state.Enabled = as.Enabled
	}
	data, err := json.Marshal(state)
	if err != nil {
		return err
	}
	ctx.GlobalSet("interface", ap.WifiIface)
	return keepState(
		fmt.Sprintf(defaultAccessPointConfig, ap.WifiIface), data)
}

func accessPointState(i string) (*AccessPointState, error) {
	dir := os.Getenv("FCONF_CONFIGDIR")
	if dir == "" {
		dir = fconfConfigDir
	}
	b, err := ioutil.ReadFile(filepath.Join(dir,
		fmt.Sprintf(defaultAccessPointConfig, i)))
	if err != nil {
		return nil, err
	}
	a := &AccessPointState{}
	err = json.Unmarshal(b, a)
	if err != nil {
		return nil, err
	}
	if a.Configg == nil {
		return nil, ErrWrongStateFile
	}
	return a, nil
}

func EnableApCMD(ctx *cli.Context) error {
	if ctx.IsSet(configFlag) {
		err := ConfigApCMD(ctx)
		if err != nil {
			return err
		}
	}
	i := getInterface(ctx)
	if i == "" {
		return errors.New("missing interface, you must specify interface")
	}
	state, err := accessPointState(i)
	if err != nil {
		return err
	}
	service := "create_ap@" + i
	err = startService(service)
	if err != nil {
		return err
	}
	err = enableService(service)
	if err != nil {
		return err
	}
	state.Enabled = true
	data, err := json.Marshal(state)
	if err != nil {
		return err
	}
	return keepState(
		fmt.Sprintf(defaultAccessPointConfig, i), data)
}

func DisableApCMD(ctx *cli.Context) error {
	if ctx.IsSet(configFlag) {
		fmt.Println("WARN: config flag will be ignored when diable flag is used")
	}
	i := getInterface(ctx)
	if i == "" {
		return errors.New("missing interface, you must specify interface")
	}
	state, err := accessPointState(i)
	if err != nil {
		return err
	}
	service := "create_ap@" + i
	err = stopService(service)
	if err != nil {
		return err
	}
	err = disableService(service)
	if err != nil {
		return err
	}
	state.Enabled = false
	data, err := json.Marshal(state)
	if err != nil {
		return err
	}
	return keepState(
		fmt.Sprintf(defaultAccessPointConfig, i), data)
}

func RemoveApCMD(ctx *cli.Context) error {
	i := getInterface(ctx)
	if i == "" {
		return errors.New("missing interface, you must specify interface")
	}
	a, err := accessPointState(i)
	if err != nil {
		return err
	}
	if a.Enabled {
		err = DisableApCMD(ctx)
		if err != nil {
			return err
		}
	}

	// remove the state file
	stateFile := filepath.Join(stateDir(),
		fmt.Sprintf(defaultAccessPointConfig, i))
	err = removeFile(stateFile)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/docker/pkg/archive"
)

var cmdCp = &Command{
	Exec:        runCp,
	UsageLine:   "cp [OPTIONS] SERVER:PATH|HOSTPATH|- SERVER:PATH|HOSTPATH|-",
	Description: "Copy files/folders from a PATH on the server to a HOSTDIR on the host",
	Help:        "Copy files/folders from a PATH on the server to a HOSTDIR on the host\nrunning the command. Use '-' to write the data as a tar file to STDOUT.",
	Examples: `
    $ scw cp path/to/my/local/file myserver:path
    $ scw cp myserver:path/to/file path/to/my/local/dir
    $ scw cp myserver:path/to/file myserver2:path/to/dir
    $ scw cp myserver:path/to/file - > myserver-pathtofile-backup.tar
    $ scw cp myserver:path/to/file - | tar -tvf -
    $ scw cp path/to/my/local/dir  myserver:path
    $ scw cp myserver:path/to/dir  path/to/my/local/dir
    $ scw cp myserver:path/to/dir  myserver2:path/to/dir
    $ scw cp myserver:path/to/dir  - > myserver-pathtodir-backup.tar
    $ scw cp myserver:path/to/dir  - | tar -tvf -
    $ cat archive.tar | scw cp - myserver:/path
    $ tar -cvf - . | scw cp - myserver:path
`,
}

func init() {
	cmdCp.Flag.BoolVar(&cpHelp, []string{"h", "-help"}, false, "Print usage")
}

// Flags
var cpHelp bool // -h, --help flag

func runCp(cmd *Command, args []string) {
	if cpHelp {
		cmd.PrintUsage()
	}
	if len(args) != 2 {
		cmd.PrintShortUsage()
	}

	var tarOutputStream io.ReadCloser
	var tarErrorStream io.ReadCloser

	// source
	source := args[0]
	if strings.Index(source, ":") > -1 { // source server address
		serverParts := strings.Split(args[0], ":")
		if len(serverParts) != 2 {
			log.Fatalf("usage: scw %s", cmd.UsageLine)
		}

		serverID := cmd.API.GetServerID(serverParts[0])

		server, err := cmd.API.GetServer(serverID)
		if err != nil {
			log.Fatalf("Failed to get server information for %s: %v", serverID, err)
		}

		dir, base := PathToTARPathparts(serverParts[1])

		// remoteCommand is executed on the remote server
		// it streams a tarball raw content
		remoteCommand := []string{"tar"}
		remoteCommand = append(remoteCommand, "-C", dir)
		if os.Getenv("DEBUG") == "1" {
			remoteCommand = append(remoteCommand, "-v")
		}
		remoteCommand = append(remoteCommand, "-cf", "-")
		remoteCommand = append(remoteCommand, base)

		// execCmd contains the ssh connection + the remoteCommand
		execCmd := append(NewSSHExecCmd(server.PublicAddress.IP, false, remoteCommand))
		log.Debugf("Executing: ssh %s", strings.Join(execCmd, " "))
		spawn := exec.Command("ssh", execCmd...)

		tarOutputStream, err = spawn.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		tarErrorStream, err = spawn.StderrPipe()
		if err != nil {
			log.Fatal(err)
		}

		err = spawn.Start()
		if err != nil {
			log.Fatalf("Failed to start ssh command: %v", err)
		}

		defer spawn.Wait()

		io.Copy(os.Stderr, tarErrorStream)
	} else if source == "-" { // stdin
		tarOutputStream = os.Stdin
	} else { // source host path
		log.Debugf("Creating tarball of local path %s", source)
		path, err := filepath.Abs(source)
		if err != nil {
			log.Fatalf("Cannot tar local path: %v", err)
		}
		path, err = filepath.EvalSymlinks(path)
		if err != nil {
			log.Fatalf("Cannot tar local path: %v", err)
		}
		log.Debugf("Real local path is %s", path)

		dir, base := PathToTARPathparts(path)

		tarOutputStream, err = archive.TarWithOptions(dir, &archive.TarOptions{
			Compression:  archive.Uncompressed,
			IncludeFiles: []string{base},
		})
		if err != nil {
			log.Fatalf("Cannot tar local path: %v", err)
		}
	}

	// destination
	destination := args[1]
	if strings.Index(destination, ":") > -1 { // destination server address
		log.Fatalf("'scw cp ... SERVER' is not yet implemented")

	} else if destination == "-" { // stdout
		log.Debugf("Writing tarOutputStream(%v) to os.Stdout(%v)", tarOutputStream, os.Stdout)
		written, err := io.Copy(os.Stdout, tarOutputStream)
		log.Debugf("%d bytes written", written)
		if err != nil {
			log.Fatal(err)
		}

	} else { // destination host path
		err := archive.Untar(tarOutputStream, destination, &archive.TarOptions{NoLchown: true})
		if err != nil {
			log.Fatalf("Failed to untar the remote archive: %v", err)
		}
	}
}
<commit_msg>wip: Handling 'scw cp ... SERVER:PATH' (#56)<commit_after>package main

import (
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/docker/pkg/archive"
)

var cmdCp = &Command{
	Exec:        runCp,
	UsageLine:   "cp [OPTIONS] SERVER:PATH|HOSTPATH|- SERVER:PATH|HOSTPATH|-",
	Description: "Copy files/folders from a PATH on the server to a HOSTDIR on the host",
	Help:        "Copy files/folders from a PATH on the server to a HOSTDIR on the host\nrunning the command. Use '-' to write the data as a tar file to STDOUT.",
	Examples: `
    $ scw cp path/to/my/local/file myserver:path
    $ scw cp myserver:path/to/file path/to/my/local/dir
    $ scw cp myserver:path/to/file myserver2:path/to/dir
    $ scw cp myserver:path/to/file - > myserver-pathtofile-backup.tar
    $ scw cp myserver:path/to/file - | tar -tvf -
    $ scw cp path/to/my/local/dir  myserver:path
    $ scw cp myserver:path/to/dir  path/to/my/local/dir
    $ scw cp myserver:path/to/dir  myserver2:path/to/dir
    $ scw cp myserver:path/to/dir  - > myserver-pathtodir-backup.tar
    $ scw cp myserver:path/to/dir  - | tar -tvf -
    $ cat archive.tar | scw cp - myserver:/path
    $ tar -cvf - . | scw cp - myserver:path
`,
}

func init() {
	cmdCp.Flag.BoolVar(&cpHelp, []string{"h", "-help"}, false, "Print usage")
}

// Flags
var cpHelp bool // -h, --help flag

func runCp(cmd *Command, args []string) {
	if cpHelp {
		cmd.PrintUsage()
	}
	if len(args) != 2 {
		cmd.PrintShortUsage()
	}

	var tarOutputStream io.ReadCloser
	var tarErrorStream io.ReadCloser

	// source
	source := args[0]
	if strings.Index(source, ":") > -1 { // source server address
		serverParts := strings.Split(args[0], ":")
		if len(serverParts) != 2 {
			log.Fatalf("usage: scw %s", cmd.UsageLine)
		}

		serverID := cmd.API.GetServerID(serverParts[0])

		server, err := cmd.API.GetServer(serverID)
		if err != nil {
			log.Fatalf("Failed to get server information for %s: %v", serverID, err)
		}

		dir, base := PathToTARPathparts(serverParts[1])

		// remoteCommand is executed on the remote server
		// it streams a tarball raw content
		remoteCommand := []string{"tar"}
		remoteCommand = append(remoteCommand, "-C", dir)
		if os.Getenv("DEBUG") == "1" {
			remoteCommand = append(remoteCommand, "-v")
		}
		remoteCommand = append(remoteCommand, "-cf", "-")
		remoteCommand = append(remoteCommand, base)

		// execCmd contains the ssh connection + the remoteCommand
		execCmd := append(NewSSHExecCmd(server.PublicAddress.IP, false, remoteCommand))
		log.Debugf("Executing: ssh %s", strings.Join(execCmd, " "))
		spawn := exec.Command("ssh", execCmd...)

		tarOutputStream, err = spawn.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		tarErrorStream, err = spawn.StderrPipe()
		if err != nil {
			log.Fatal(err)
		}

		err = spawn.Start()
		if err != nil {
			log.Fatalf("Failed to start ssh command: %v", err)
		}

		defer spawn.Wait()

		io.Copy(os.Stderr, tarErrorStream)
	} else if source == "-" { // stdin
		tarOutputStream = os.Stdin
	} else { // source host path
		log.Debugf("Creating tarball of local path %s", source)
		path, err := filepath.Abs(source)
		if err != nil {
			log.Fatalf("Cannot tar local path: %v", err)
		}
		path, err = filepath.EvalSymlinks(path)
		if err != nil {
			log.Fatalf("Cannot tar local path: %v", err)
		}
		log.Debugf("Real local path is %s", path)

		dir, base := PathToTARPathparts(path)

		tarOutputStream, err = archive.TarWithOptions(dir, &archive.TarOptions{
			Compression:  archive.Uncompressed,
			IncludeFiles: []string{base},
		})
		if err != nil {
			log.Fatalf("Cannot tar local path: %v", err)
		}
	}

	// destination
	destination := args[1]
	if strings.Index(destination, ":") > -1 { // destination server address
		serverParts := strings.Split(destination, ":")
		if len(serverParts) != 2 {
			log.Fatalf("usage: scw %s", cmd.UsageLine)
		}

		serverID := cmd.API.GetServerID(serverParts[0])

		server, err := cmd.API.GetServer(serverID)
		if err != nil {
			log.Fatalf("Failed to get server information for %s: %v", serverID, err)
		}

		// remoteCommand is executed on the remote server
		// it streams a tarball raw content
		remoteCommand := []string{"tar"}
		remoteCommand = append(remoteCommand, "-C", serverParts[1])
		if os.Getenv("DEBUG") == "1" {
			remoteCommand = append(remoteCommand, "-v")
		}
		remoteCommand = append(remoteCommand, "-tf", "-")

		// execCmd contains the ssh connection + the remoteCommand
		execCmd := append(NewSSHExecCmd(server.PublicAddress.IP, false, remoteCommand))
		log.Debugf("Executing: ssh %s", strings.Join(execCmd, " "))
		spawn := exec.Command("ssh", execCmd...)

		untarInputStream, err := spawn.StdinPipe()
		if err != nil {
			log.Fatal(err)
		}
		untarErrorStream, err := spawn.StderrPipe()
		if err != nil {
			log.Fatal(err)
		}
		untarOutputStream, err := spawn.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}

		err = spawn.Start()
		if err != nil {
			log.Fatalf("Failed to start ssh command: %v", err)
		}

		defer spawn.Wait()

		io.Copy(untarInputStream, tarOutputStream)
		io.Copy(os.Stderr, untarErrorStream)
		io.Copy(os.Stdout, untarOutputStream)
	} else if destination == "-" { // stdout
		log.Debugf("Writing tarOutputStream(%v) to os.Stdout(%v)", tarOutputStream, os.Stdout)
		written, err := io.Copy(os.Stdout, tarOutputStream)
		log.Debugf("%d bytes written", written)
		if err != nil {
			log.Fatal(err)
		}

	} else { // destination host path
		err := archive.Untar(tarOutputStream, destination, &archive.TarOptions{NoLchown: true})
		if err != nil {
			log.Fatalf("Failed to untar the remote archive: %v", err)
		}
	}
}
<|endoftext|>


<commit_before>package pg // import "gopkg.in/pg.v4"

import (
	"io"
	"log"
	"net"
	"time"

	"gopkg.in/pg.v4/orm"
	"gopkg.in/pg.v4/types"
)

const defaultBackoff = 100 * time.Millisecond

// Database connection options.
type Options struct {
	// The network type, either tcp or unix.
	// Default is tcp.
	Network  string
	Addr     string
	User     string
	Password string
	Database string
	// Whether to use secure TCP/IP connections (TLS).
	SSL bool

	// Run-time configuration parameters to be set on connection.
	Params map[string]interface{}

	// The deadline for establishing new connections. If reached,
	// dial will fail with a timeout.
	// Default is 5 seconds.
	DialTimeout time.Duration
	// The timeout for socket reads. If reached, commands will fail
	// with a timeout error instead of blocking.
	// Default is no timeout.
	ReadTimeout time.Duration
	// The timeout for socket writes. If reached, commands will fail
	// with a timeout error instead of blocking.
	// Default is no timeout.
	WriteTimeout time.Duration

	// The maximum number of open socket connections.
	// Default is 10 connections.
	PoolSize int
	// The amount of time client waits for free connection if all
	// connections are busy before returning an error.
	// Default is 5 seconds.
	PoolTimeout time.Duration
	// The amount of time after which client closes idle connections.
	// Default is to not close idle connections.
	IdleTimeout time.Duration
	// The frequency of idle checks.
	// Default is 1 minute.
	IdleCheckFrequency time.Duration
}

func (opt *Options) getNetwork() string {
	if opt == nil || opt.Network == "" {
		return "tcp"
	}
	return opt.Network
}

func (opt *Options) getAddr() string {
	if opt.Addr != "" {
		return opt.Addr
	}
	if opt.getNetwork() == "unix" {
		return "/var/run/postgresql/.s.PGSQL.5432"
	}
	return "localhost:5432"
}

func (opt *Options) getUser() string {
	if opt == nil || opt.User == "" {
		return ""
	}
	return opt.User
}

func (opt *Options) getPassword() string {
	if opt == nil || opt.Password == "" {
		return ""
	}
	return opt.Password
}

func (opt *Options) getDatabase() string {
	if opt == nil || opt.Database == "" {
		return ""
	}
	return opt.Database
}

func (opt *Options) getPoolSize() int {
	if opt == nil || opt.PoolSize == 0 {
		return 10
	}
	return opt.PoolSize
}

func (opt *Options) getPoolTimeout() time.Duration {
	if opt == nil || opt.PoolTimeout == 0 {
		return 5 * time.Second
	}
	return opt.PoolTimeout
}

func (opt *Options) getDialTimeout() time.Duration {
	if opt.DialTimeout == 0 {
		return 5 * time.Second
	}
	return opt.DialTimeout
}

func (opt *Options) getIdleTimeout() time.Duration {
	return opt.IdleTimeout
}

func (opt *Options) getIdleCheckFrequency() time.Duration {
	if opt.IdleCheckFrequency == 0 {
		return time.Minute
	}
	return opt.IdleCheckFrequency
}

func (opt *Options) getSSL() bool {
	if opt == nil {
		return false
	}
	return opt.SSL
}

// Connect connects to a database using provided options.
//
// The returned DB is safe for concurrent use by multiple goroutines
// and maintains its own connection pool.
func Connect(opt *Options) *DB {
	return &DB{
		opt:  opt,
		pool: newConnPool(opt),
	}
}

// DB is a database handle representing a pool of zero or more
// underlying connections. It's safe for concurrent use by multiple
// goroutines.
type DB struct {
	opt  *Options
	pool *connPool
}

// Options returns read-only Options that were used to connect to the DB.
func (db *DB) Options() *Options {
	return db.opt
}

// WithTimeout returns a DB that uses d as the read/write timeout.
func (db *DB) WithTimeout(d time.Duration) *DB {
	newopt := *db.opt
	newopt.ReadTimeout = d
	newopt.WriteTimeout = d
	return &DB{
		opt:  &newopt,
		pool: db.pool,
	}
}

func (db *DB) conn() (*conn, error) {
	cn, _, err := db.pool.Get()
	if err != nil {
		return nil, err
	}

	cn.SetReadTimeout(db.opt.ReadTimeout)
	cn.SetWriteTimeout(db.opt.WriteTimeout)
	return cn, nil
}

func (db *DB) freeConn(cn *conn, err error) error {
	if err == nil {
		return db.pool.Put(cn)
	}
	if cn.rd.Buffered() > 0 {
		return db.pool.Remove(cn, err)
	}
	if pgerr, ok := err.(Error); ok && pgerr.Field('S') != "FATAL" {
		return db.pool.Put(cn)
	}
	if _, ok := err.(dbError); ok {
		return db.pool.Put(cn)
	}
	if neterr, ok := err.(net.Error); ok && neterr.Timeout() {
		if err := db.cancelRequest(cn.processId, cn.secretKey); err != nil {
			log.Printf("pg: cancelRequest failed: %s", err)
		}
	}
	return db.pool.Remove(cn, err)
}

// Close closes the database client, releasing any open resources.
//
// It is rare to Close a DB, as the DB handle is meant to be
// long-lived and shared between many goroutines.
func (db *DB) Close() error {
	return db.pool.Close()
}

// Exec executes a query ignoring returned rows. The params are for
// any placeholder parameters in the query.
func (db *DB) Exec(query interface{}, params ...interface{}) (res types.Result, err error) {
	backoff := defaultBackoff
	for i := 0; i < 3; i++ {
		var cn *conn

		cn, err = db.conn()
		if err != nil {
			return nil, err
		}

		res, err = simpleQuery(cn, query, params...)
		db.freeConn(cn, err)
		if !canRetry(err) {
			break
		}

		time.Sleep(backoff)
		backoff *= 2
	}
	return
}

// ExecOne acts like Exec, but query must affect only one row. It
// returns ErrNoRows error when query returns zero rows or
// ErrMultiRows when query returns multiple rows.
func (db *DB) ExecOne(query interface{}, params ...interface{}) (types.Result, error) {
	res, err := db.Exec(query, params...)
	if err != nil {
		return nil, err
	}
	return assertOneAffected(res, nil)
}

// Query executes a query that returns rows, typically a SELECT. The
// params are for any placeholder parameters in the query.
func (db *DB) Query(model, query interface{}, params ...interface{}) (res types.Result, err error) {
	backoff := defaultBackoff
	for i := 0; i < 3; i++ {
		var cn *conn

		cn, err = db.conn()
		if err != nil {
			break
		}

		res, err = simpleQueryData(cn, model, query, params...)
		db.freeConn(cn, err)
		if !canRetry(err) {
			break
		}

		time.Sleep(backoff)
		backoff *= 2
	}
	return
}

// QueryOne acts like Query, but query must return only one row. It
// returns ErrNoRows error when query returns zero rows or
// ErrMultiRows when query returns multiple rows.
func (db *DB) QueryOne(model, query interface{}, params ...interface{}) (types.Result, error) {
	mod, err := newSingleModel(model)
	if err != nil {
		return nil, err
	}
	res, err := db.Query(mod, query, params...)
	if err != nil {
		return nil, err
	}
	return assertOneAffected(res, mod)
}

// Listen listens for notifications sent by NOTIFY statement.
func (db *DB) Listen(channels ...string) (*Listener, error) {
	l := &Listener{
		db: db,
	}
	if err := l.Listen(channels...); err != nil {
		l.Close()
		return nil, err
	}
	return l, nil
}

// CopyFrom copies data from reader to a table.
func (db *DB) CopyFrom(r io.Reader, q string, params ...interface{}) (types.Result, error) {
	cn, err := db.conn()
	if err != nil {
		return nil, err
	}
	res, err := copyFrom(cn, r, q, params...)
	if err != nil {
		db.freeConn(cn, err)
		return nil, err
	}
	db.pool.Put(cn)
	return res, nil
}

// CopyTo copies data from a table to writer.
func (db *DB) CopyTo(w io.WriteCloser, q string, params ...interface{}) (types.Result, error) {
	cn, err := db.conn()
	if err != nil {
		return nil, err
	}

	if err := writeQueryMsg(cn.buf, q, params...); err != nil {
		db.pool.Put(cn)
		return nil, err
	}

	if err := cn.FlushWrite(); err != nil {
		db.freeConn(cn, err)
		return nil, err
	}

	if err := readCopyOutResponse(cn); err != nil {
		db.freeConn(cn, err)
		return nil, err
	}

	res, err := readCopyData(cn, w)
	if err != nil {
		db.freeConn(cn, err)
		return nil, err
	}

	return res, nil
}

func (db *DB) Model(model interface{}) *orm.Query {
	return orm.NewQuery(db, model)
}

func (db *DB) Create(model interface{}) error {
	return orm.Create(db, model)
}

func (db *DB) Update(model interface{}) error {
	return orm.Update(db, model)
}

func (db *DB) Delete(model interface{}) error {
	return orm.Delete(db, model)
}

func setParams(cn *conn, params map[string]interface{}) error {
	for key, value := range params {
		_, err := simpleQuery(cn, "SET ? = ?", F(key), value)
		if err != nil {
			return err
		}
	}
	return nil
}

func (db *DB) cancelRequest(processId, secretKey int32) error {
	cn, err := dial(db.opt)
	if err != nil {
		return err
	}

	buf := newBuffer()
	writeCancelRequestMsg(buf, processId, secretKey)
	_, err = cn.Write(buf.Flush())
	if err != nil {
		return err
	}

	return cn.Close()
}

func simpleQuery(cn *conn, query interface{}, params ...interface{}) (types.Result, error) {
	if err := writeQueryMsg(cn.buf, query, params...); err != nil {
		return nil, err
	}

	if err := cn.FlushWrite(); err != nil {
		return nil, err
	}

	res, err := readSimpleQuery(cn)
	if err != nil {
		return nil, err
	}

	return res, nil
}

func simpleQueryData(cn *conn, model, query interface{}, params ...interface{}) (types.Result, error) {
	if err := writeQueryMsg(cn.buf, query, params...); err != nil {
		return nil, err
	}

	if err := cn.FlushWrite(); err != nil {
		return nil, err
	}

	res, err := readSimpleQueryData(cn, model)
	if err != nil {
		return nil, err
	}

	return res, nil
}

type singleModel struct {
	model orm.ColumnScanner
	len   int
}

var _ orm.Collection = (*singleModel)(nil)

func newSingleModel(model interface{}) (*singleModel, error) {
	mod, ok := model.(orm.Model)
	if !ok {
		var err error
		mod, err = orm.NewModel(model)
		if err != nil {
			return nil, err
		}
	}
	return &singleModel{
		model: mod,
	}, nil
}

func (m *singleModel) NewModel() orm.ColumnScanner {
	m.len++
	return m.model
}

func (m *singleModel) Len() int {
	return m.len
}

func assertOne(l int) error {
	switch {
	case l == 0:
		return ErrNoRows
	case l > 1:
		return ErrMultiRows
	default:
		return nil
	}
}

func assertOneAffected(res types.Result, model *singleModel) (types.Result, error) {
	if err := assertOne(res.Affected()); err != nil {
		return nil, err
	}
	if model != nil {
		if err := assertOne(model.Len()); err != nil {
			return nil, err
		}
	}
	return res, nil
}

func copyFrom(cn *conn, r io.Reader, query interface{}, params ...interface{}) (types.Result, error) {
	if err := writeQueryMsg(cn.buf, query, params...); err != nil {
		return nil, err
	}

	if err := cn.FlushWrite(); err != nil {
		return nil, err
	}

	if err := readCopyInResponse(cn); err != nil {
		return nil, err
	}

	ready := make(chan struct{})
	var res types.Result
	var err error
	go func() {
		res, err = readReadyForQuery(cn)
		close(ready)
	}()

	for {
		select {
		case <-ready:
			break
		default:
		}

		_, err := writeCopyData(cn.buf, r)
		if err == io.EOF {
			break
		}

		if err := cn.FlushWrite(); err != nil {
			return nil, err
		}
	}

	select {
	case <-ready:
	default:
	}

	writeCopyDone(cn.buf)
	if err := cn.FlushWrite(); err != nil {
		return nil, err
	}

	<-ready
	if err != nil {
		return nil, err
	}
	return res, nil
}
<commit_msg>Remove duplicate check.<commit_after>package pg // import "gopkg.in/pg.v4"

import (
	"io"
	"log"
	"net"
	"time"

	"gopkg.in/pg.v4/orm"
	"gopkg.in/pg.v4/types"
)

const defaultBackoff = 100 * time.Millisecond

// Database connection options.
type Options struct {
	// The network type, either tcp or unix.
	// Default is tcp.
	Network  string
	Addr     string
	User     string
	Password string
	Database string
	// Whether to use secure TCP/IP connections (TLS).
	SSL bool

	// Run-time configuration parameters to be set on connection.
	Params map[string]interface{}

	// The deadline for establishing new connections. If reached,
	// dial will fail with a timeout.
	// Default is 5 seconds.
	DialTimeout time.Duration
	// The timeout for socket reads. If reached, commands will fail
	// with a timeout error instead of blocking.
	// Default is no timeout.
	ReadTimeout time.Duration
	// The timeout for socket writes. If reached, commands will fail
	// with a timeout error instead of blocking.
	// Default is no timeout.
	WriteTimeout time.Duration

	// The maximum number of open socket connections.
	// Default is 10 connections.
	PoolSize int
	// The amount of time client waits for free connection if all
	// connections are busy before returning an error.
	// Default is 5 seconds.
	PoolTimeout time.Duration
	// The amount of time after which client closes idle connections.
	// Default is to not close idle connections.
	IdleTimeout time.Duration
	// The frequency of idle checks.
	// Default is 1 minute.
	IdleCheckFrequency time.Duration
}

func (opt *Options) getNetwork() string {
	if opt == nil || opt.Network == "" {
		return "tcp"
	}
	return opt.Network
}

func (opt *Options) getAddr() string {
	if opt.Addr != "" {
		return opt.Addr
	}
	if opt.getNetwork() == "unix" {
		return "/var/run/postgresql/.s.PGSQL.5432"
	}
	return "localhost:5432"
}

func (opt *Options) getUser() string {
	if opt == nil || opt.User == "" {
		return ""
	}
	return opt.User
}

func (opt *Options) getPassword() string {
	if opt == nil || opt.Password == "" {
		return ""
	}
	return opt.Password
}

func (opt *Options) getDatabase() string {
	if opt == nil || opt.Database == "" {
		return ""
	}
	return opt.Database
}

func (opt *Options) getPoolSize() int {
	if opt == nil || opt.PoolSize == 0 {
		return 10
	}
	return opt.PoolSize
}

func (opt *Options) getPoolTimeout() time.Duration {
	if opt == nil || opt.PoolTimeout == 0 {
		return 5 * time.Second
	}
	return opt.PoolTimeout
}

func (opt *Options) getDialTimeout() time.Duration {
	if opt.DialTimeout == 0 {
		return 5 * time.Second
	}
	return opt.DialTimeout
}

func (opt *Options) getIdleTimeout() time.Duration {
	return opt.IdleTimeout
}

func (opt *Options) getIdleCheckFrequency() time.Duration {
	if opt.IdleCheckFrequency == 0 {
		return time.Minute
	}
	return opt.IdleCheckFrequency
}

func (opt *Options) getSSL() bool {
	if opt == nil {
		return false
	}
	return opt.SSL
}

// Connect connects to a database using provided options.
//
// The returned DB is safe for concurrent use by multiple goroutines
// and maintains its own connection pool.
func Connect(opt *Options) *DB {
	return &DB{
		opt:  opt,
		pool: newConnPool(opt),
	}
}

// DB is a database handle representing a pool of zero or more
// underlying connections. It's safe for concurrent use by multiple
// goroutines.
type DB struct {
	opt  *Options
	pool *connPool
}

// Options returns read-only Options that were used to connect to the DB.
func (db *DB) Options() *Options {
	return db.opt
}

// WithTimeout returns a DB that uses d as the read/write timeout.
func (db *DB) WithTimeout(d time.Duration) *DB {
	newopt := *db.opt
	newopt.ReadTimeout = d
	newopt.WriteTimeout = d
	return &DB{
		opt:  &newopt,
		pool: db.pool,
	}
}

func (db *DB) conn() (*conn, error) {
	cn, _, err := db.pool.Get()
	if err != nil {
		return nil, err
	}

	cn.SetReadTimeout(db.opt.ReadTimeout)
	cn.SetWriteTimeout(db.opt.WriteTimeout)
	return cn, nil
}

func (db *DB) freeConn(cn *conn, err error) error {
	if err == nil {
		return db.pool.Put(cn)
	}
	if pgerr, ok := err.(Error); ok && pgerr.Field('S') != "FATAL" {
		return db.pool.Put(cn)
	}
	if _, ok := err.(dbError); ok {
		return db.pool.Put(cn)
	}
	if neterr, ok := err.(net.Error); ok && neterr.Timeout() {
		if err := db.cancelRequest(cn.processId, cn.secretKey); err != nil {
			log.Printf("pg: cancelRequest failed: %s", err)
		}
	}
	return db.pool.Remove(cn, err)
}

// Close closes the database client, releasing any open resources.
//
// It is rare to Close a DB, as the DB handle is meant to be
// long-lived and shared between many goroutines.
func (db *DB) Close() error {
	return db.pool.Close()
}

// Exec executes a query ignoring returned rows. The params are for
// any placeholder parameters in the query.
func (db *DB) Exec(query interface{}, params ...interface{}) (res types.Result, err error) {
	backoff := defaultBackoff
	for i := 0; i < 3; i++ {
		var cn *conn

		cn, err = db.conn()
		if err != nil {
			return nil, err
		}

		res, err = simpleQuery(cn, query, params...)
		db.freeConn(cn, err)
		if !canRetry(err) {
			break
		}

		time.Sleep(backoff)
		backoff *= 2
	}
	return
}

// ExecOne acts like Exec, but query must affect only one row. It
// returns ErrNoRows error when query returns zero rows or
// ErrMultiRows when query returns multiple rows.
func (db *DB) ExecOne(query interface{}, params ...interface{}) (types.Result, error) {
	res, err := db.Exec(query, params...)
	if err != nil {
		return nil, err
	}
	return assertOneAffected(res, nil)
}

// Query executes a query that returns rows, typically a SELECT. The
// params are for any placeholder parameters in the query.
func (db *DB) Query(model, query interface{}, params ...interface{}) (res types.Result, err error) {
	backoff := defaultBackoff
	for i := 0; i < 3; i++ {
		var cn *conn

		cn, err = db.conn()
		if err != nil {
			break
		}

		res, err = simpleQueryData(cn, model, query, params...)
		db.freeConn(cn, err)
		if !canRetry(err) {
			break
		}

		time.Sleep(backoff)
		backoff *= 2
	}
	return
}

// QueryOne acts like Query, but query must return only one row. It
// returns ErrNoRows error when query returns zero rows or
// ErrMultiRows when query returns multiple rows.
func (db *DB) QueryOne(model, query interface{}, params ...interface{}) (types.Result, error) {
	mod, err := newSingleModel(model)
	if err != nil {
		return nil, err
	}
	res, err := db.Query(mod, query, params...)
	if err != nil {
		return nil, err
	}
	return assertOneAffected(res, mod)
}

// Listen listens for notifications sent by NOTIFY statement.
func (db *DB) Listen(channels ...string) (*Listener, error) {
	l := &Listener{
		db: db,
	}
	if err := l.Listen(channels...); err != nil {
		l.Close()
		return nil, err
	}
	return l, nil
}

// CopyFrom copies data from reader to a table.
func (db *DB) CopyFrom(r io.Reader, q string, params ...interface{}) (types.Result, error) {
	cn, err := db.conn()
	if err != nil {
		return nil, err
	}
	res, err := copyFrom(cn, r, q, params...)
	if err != nil {
		db.freeConn(cn, err)
		return nil, err
	}
	db.pool.Put(cn)
	return res, nil
}

// CopyTo copies data from a table to writer.
func (db *DB) CopyTo(w io.WriteCloser, q string, params ...interface{}) (types.Result, error) {
	cn, err := db.conn()
	if err != nil {
		return nil, err
	}

	if err := writeQueryMsg(cn.buf, q, params...); err != nil {
		db.pool.Put(cn)
		return nil, err
	}

	if err := cn.FlushWrite(); err != nil {
		db.freeConn(cn, err)
		return nil, err
	}

	if err := readCopyOutResponse(cn); err != nil {
		db.freeConn(cn, err)
		return nil, err
	}

	res, err := readCopyData(cn, w)
	if err != nil {
		db.freeConn(cn, err)
		return nil, err
	}

	return res, nil
}

func (db *DB) Model(model interface{}) *orm.Query {
	return orm.NewQuery(db, model)
}

func (db *DB) Create(model interface{}) error {
	return orm.Create(db, model)
}

func (db *DB) Update(model interface{}) error {
	return orm.Update(db, model)
}

func (db *DB) Delete(model interface{}) error {
	return orm.Delete(db, model)
}

func setParams(cn *conn, params map[string]interface{}) error {
	for key, value := range params {
		_, err := simpleQuery(cn, "SET ? = ?", F(key), value)
		if err != nil {
			return err
		}
	}
	return nil
}

func (db *DB) cancelRequest(processId, secretKey int32) error {
	cn, err := dial(db.opt)
	if err != nil {
		return err
	}

	buf := newBuffer()
	writeCancelRequestMsg(buf, processId, secretKey)
	_, err = cn.Write(buf.Flush())
	if err != nil {
		return err
	}

	return cn.Close()
}

func simpleQuery(cn *conn, query interface{}, params ...interface{}) (types.Result, error) {
	if err := writeQueryMsg(cn.buf, query, params...); err != nil {
		return nil, err
	}

	if err := cn.FlushWrite(); err != nil {
		return nil, err
	}

	res, err := readSimpleQuery(cn)
	if err != nil {
		return nil, err
	}

	return res, nil
}

func simpleQueryData(cn *conn, model, query interface{}, params ...interface{}) (types.Result, error) {
	if err := writeQueryMsg(cn.buf, query, params...); err != nil {
		return nil, err
	}

	if err := cn.FlushWrite(); err != nil {
		return nil, err
	}

	res, err := readSimpleQueryData(cn, model)
	if err != nil {
		return nil, err
	}

	return res, nil
}

type singleModel struct {
	model orm.ColumnScanner
	len   int
}

var _ orm.Collection = (*singleModel)(nil)

func newSingleModel(model interface{}) (*singleModel, error) {
	mod, ok := model.(orm.Model)
	if !ok {
		var err error
		mod, err = orm.NewModel(model)
		if err != nil {
			return nil, err
		}
	}
	return &singleModel{
		model: mod,
	}, nil
}

func (m *singleModel) NewModel() orm.ColumnScanner {
	m.len++
	return m.model
}

func (m *singleModel) Len() int {
	return m.len
}

func assertOne(l int) error {
	switch {
	case l == 0:
		return ErrNoRows
	case l > 1:
		return ErrMultiRows
	default:
		return nil
	}
}

func assertOneAffected(res types.Result, model *singleModel) (types.Result, error) {
	if err := assertOne(res.Affected()); err != nil {
		return nil, err
	}
	if model != nil {
		if err := assertOne(model.Len()); err != nil {
			return nil, err
		}
	}
	return res, nil
}

func copyFrom(cn *conn, r io.Reader, query interface{}, params ...interface{}) (types.Result, error) {
	if err := writeQueryMsg(cn.buf, query, params...); err != nil {
		return nil, err
	}

	if err := cn.FlushWrite(); err != nil {
		return nil, err
	}

	if err := readCopyInResponse(cn); err != nil {
		return nil, err
	}

	ready := make(chan struct{})
	var res types.Result
	var err error
	go func() {
		res, err = readReadyForQuery(cn)
		close(ready)
	}()

	for {
		select {
		case <-ready:
			break
		default:
		}

		_, err := writeCopyData(cn.buf, r)
		if err == io.EOF {
			break
		}

		if err := cn.FlushWrite(); err != nil {
			return nil, err
		}
	}

	select {
	case <-ready:
	default:
	}

	writeCopyDone(cn.buf)
	if err := cn.FlushWrite(); err != nil {
		return nil, err
	}

	<-ready
	if err != nil {
		return nil, err
	}
	return res, nil
}
<|endoftext|>


<commit_before>/*
 * (c) 2014, Caoimhe Chaos <caoimhechaos@protonmail.com>,
 *	     Ancient Solutions. All rights reserved.
 *
 * Redistribution and use in source  and binary forms, with or without
 * modification, are permitted  provided that the following conditions
 * are met:
 *
 * * Redistributions of  source code  must retain the  above copyright
 *   notice, this list of conditions and the following disclaimer.
 * * Redistributions in binary form must reproduce the above copyright
 *   notice, this  list of conditions and the  following disclaimer in
 *   the  documentation  and/or  other  materials  provided  with  the
 *   distribution.
 * * Neither  the  name  of  Ancient Solutions  nor  the  name  of its
 *   contributors may  be used to endorse or  promote products derived
 *   from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS"  AND ANY EXPRESS  OR IMPLIED WARRANTIES  OF MERCHANTABILITY
 * AND FITNESS  FOR A PARTICULAR  PURPOSE ARE DISCLAIMED. IN  NO EVENT
 * SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL,  EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED  TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE,  DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT  LIABILITY,  OR  TORT  (INCLUDING NEGLIGENCE  OR  OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 */

package main

import (
	"bytes"
	"crypto/sha1"
	"database/cassandra"
	"encoding/hex"
	"errors"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	"image/png"
	"time"

	"github.com/nfnt/resize"
)

// Database access routines for the image database.
type PikShrDB struct {
	db *cassandra.RetryCassandraClient
}

type Picture struct {
	Id          string
	Contents    []byte
	ContentType string
	Title       string
	Description string
	AltText     string
}

var Err_ImageNotFound = errors.New("Image not found")

// Connect to the picture database.
func NewPikShrDB(server, keyspace string) (*PikShrDB, error) {
	var ire *cassandra.InvalidRequestException
	var client *cassandra.RetryCassandraClient
	var err error

	client, err = cassandra.NewRetryCassandraClient(server)
	if err != nil {
		return nil, err
	}

	ire, err = client.SetKeyspace(keyspace)
	if ire != nil {
		return nil, errors.New(ire.Why)
	}
	if err != nil {
		return nil, err
	}

	return &PikShrDB{
		db: client,
	}, nil
}

// Retrieve image data and metadata for the given ID.
// Image data will be taken from "column".
func (p *PikShrDB) getPictureColumn(id, column string) (*Picture, error) {
	var r []*cassandra.ColumnOrSuperColumn
	var scol *cassandra.ColumnOrSuperColumn
	var ret *Picture = new(Picture)
	var cp *cassandra.ColumnParent = cassandra.NewColumnParent()
	var pred *cassandra.SlicePredicate = cassandra.NewSlicePredicate()
	var ire *cassandra.InvalidRequestException
	var ue *cassandra.UnavailableException
	var te *cassandra.TimedOutException
	var key []byte
	var err error

	cp.ColumnFamily = "picture"
	pred.ColumnNames = [][]byte{
		[]byte("title"), []byte("description"), []byte("content_type"),
		[]byte("alt_text"),
	}

	if len(column) > 0 {
		pred.ColumnNames = append(pred.ColumnNames, []byte(column))
	}

	ret.Id = id
	key, err = hex.DecodeString(id)
	if err != nil {
		return nil, err
	}

	r, ire, ue, te, err = p.db.GetSlice(
		key, cp, pred, cassandra.ConsistencyLevel_ONE)
	if ire != nil {
		return nil, errors.New(ire.Why)
	}
	if ue != nil {
		return nil, errors.New("Unavailable")
	}
	if te != nil {
		return nil, errors.New("Timed out")
	}
	if err != nil {
		return nil, err
	}

	if len(r) == 0 {
		return nil, Err_ImageNotFound
	}

	for _, scol = range r {
		var col = scol.Column
		var cname string = string(col.Name)

		if cname == "title" {
			ret.Title = string(col.Value)
		} else if cname == "description" {
			ret.Description = string(col.Value)
		} else if cname == "content_type" {
			ret.ContentType = string(col.Value)
		} else if cname == "alt_text" {
			ret.AltText = string(col.Value)
		} else if len(column) > 0 && cname == column {
			ret.Contents = make([]byte, len(col.Value))
			copy(ret.Contents, col.Value)
		}
	}

	return ret, nil
}

// Retrieve the picture with the given ID from the database.
func (p *PikShrDB) GetPicture(id string) (*Picture, error) {
	return p.getPictureColumn(id, "picture")
}

// Retrieve the thumbnail with the given ID from the database.
func (p *PikShrDB) GetThumbnail(id string) (*Picture, error) {
	return p.getPictureColumn(id, "thumbnail")
}

// Retrieve only the metadata of the picture.
func (p *PikShrDB) GetMetadata(id string) (*Picture, error) {
	return p.getPictureColumn(id, "")
}

func makeMutation(name string, value []byte, now time.Time) (ret *cassandra.Mutation) {
	var cos = cassandra.NewColumnOrSuperColumn()
	var col = cassandra.NewColumn()

	col.Name = []byte(name)
	col.Value = value
	col.Timestamp = now.UnixNano()

	cos.Column = col

	ret = cassandra.NewMutation()
	ret.ColumnOrSupercolumn = cos
	return
}

// Insert the given picture and metadata into the database.
// Also, generate a thumbnail.
func (p *PikShrDB) InsertPicture(pic *Picture, creator string) (string, error) {
	var mutation *cassandra.Mutation
	var ire *cassandra.InvalidRequestException
	var ue *cassandra.UnavailableException
	var te *cassandra.TimedOutException
	var mmap = make(map[string]map[string][]*cassandra.Mutation)
	var mlist []*cassandra.Mutation
	var now = time.Now()
	var img, thumbnail image.Image
	var buf *bytes.Buffer = new(bytes.Buffer)
	var rd *bytes.Reader = bytes.NewReader(pic.Contents)
	var keystr string
	var key [sha1.Size]byte
	var err error

	// First, let's read and parse the image.
	img, _, err = image.Decode(rd)
	if err != nil {
		return "", err
	}

	// Write the full image out as a png.
	err = png.Encode(buf, img)
	if err != nil {
		return "", err
	}

	// Compute a hash of the image data.
	key = sha1.Sum(buf.Bytes())
	keystr = hex.EncodeToString(key[:])

	mmap[string(key[:])] = make(map[string][]*cassandra.Mutation)
	mlist = make([]*cassandra.Mutation, 0)

	mutation = makeMutation("picture", buf.Bytes(), now)
	mlist = append(mlist, mutation)

	// Create a thumbnail and save it too.
	thumbnail = resize.Thumbnail(200, 200, img, resize.Lanczos3)
	err = png.Encode(buf, thumbnail)
	if err != nil {
		return keystr, err
	}
	img = nil

	// Write the thumbnail out as a png.
	buf = new(bytes.Buffer)
	err = png.Encode(buf, thumbnail)
	if err != nil {
		return keystr, err
	}
	thumbnail = nil

	mutation = makeMutation("thumbnail", buf.Bytes(), now)
	mlist = append(mlist, mutation)

	// Now for some metadata.
	mutation = makeMutation("title", []byte(pic.Title), now)
	mlist = append(mlist, mutation)

	mutation = makeMutation("description", []byte(pic.Description), now)
	mlist = append(mlist, mutation)

	mutation = makeMutation("owner", []byte(creator), now)
	mlist = append(mlist, mutation)

	mutation = makeMutation("alt_text", []byte(pic.AltText), now)
	mlist = append(mlist, mutation)

	mmap[string(key[:])]["picture"] = mlist

	// Now: write it!
	ire, ue, te, err = p.db.AtomicBatchMutate(mmap, cassandra.ConsistencyLevel_QUORUM)
	if ire != nil {
		return keystr, errors.New(ire.Why)
	}
	if ue != nil {
		return keystr, errors.New("Unavailable")
	}
	if te != nil {
		return keystr, errors.New("Timed out")
	}
	return keystr, err
}

// List the num most recent pictures for the given user. If the user name
// is empty, list the most recent pictures for all users.
func (p *PikShrDB) GetRecentPics(user string, num int32) (ret []*Picture, err error) {
	var cp *cassandra.ColumnParent = cassandra.NewColumnParent()
	var pred *cassandra.SlicePredicate = cassandra.NewSlicePredicate()
	var rng *cassandra.KeyRange = cassandra.NewKeyRange()
	var ks []*cassandra.KeySlice
	var slice *cassandra.KeySlice
	var ire *cassandra.InvalidRequestException
	var ue *cassandra.UnavailableException
	var te *cassandra.TimedOutException

	cp.ColumnFamily = "picture"
	pred.ColumnNames = [][]byte{
		[]byte("title"), []byte("description"), []byte("content_type"),
		[]byte("alt_text"),
	}
	rng.StartKey = make([]byte, 0)
	rng.EndKey = make([]byte, 0)
	rng.Count = num

	if len(user) > 0 {
		var indexexpr = cassandra.NewIndexExpression()

		indexexpr.ColumnName = []byte("owner")
		indexexpr.Op = cassandra.IndexOperator_EQ
		indexexpr.Value = []byte(user)

		rng.RowFilter = make([]*cassandra.IndexExpression, 0)
		rng.RowFilter = append(rng.RowFilter, indexexpr)
	}

	ks, ire, ue, te, err = p.db.GetRangeSlices(
		cp, pred, rng, cassandra.ConsistencyLevel_ONE)
	if ire != nil {
		err = errors.New(ire.Why)
		return
	}
	if ue != nil {
		err = errors.New("Unavailable")
		return
	}
	if te != nil {
		err = errors.New("Timed out")
		return
	}
	if err != nil {
		return
	}

	ret = make([]*Picture, 0)

	for _, slice = range ks {
		var pic *Picture = new(Picture)
		var scol *cassandra.ColumnOrSuperColumn

		pic.Id = hex.EncodeToString(slice.Key)

		for _, scol = range slice.Columns {
			var col = scol.Column
			var cname string = string(col.Name)

			if cname == "title" {
				pic.Title = string(col.Value)
			} else if cname == "description" {
				pic.Description = string(col.Value)
			} else if cname == "content_type" {
				pic.ContentType = string(col.Value)
			} else if cname == "alt_text" {
				pic.AltText = string(col.Value)
			}
		}

		ret = append(ret, pic)
	}

	return
}
<commit_msg>Protect the Cassandra connection with a mutex.<commit_after>/*
 * (c) 2014, Caoimhe Chaos <caoimhechaos@protonmail.com>,
 *	     Ancient Solutions. All rights reserved.
 *
 * Redistribution and use in source  and binary forms, with or without
 * modification, are permitted  provided that the following conditions
 * are met:
 *
 * * Redistributions of  source code  must retain the  above copyright
 *   notice, this list of conditions and the following disclaimer.
 * * Redistributions in binary form must reproduce the above copyright
 *   notice, this  list of conditions and the  following disclaimer in
 *   the  documentation  and/or  other  materials  provided  with  the
 *   distribution.
 * * Neither  the  name  of  Ancient Solutions  nor  the  name  of its
 *   contributors may  be used to endorse or  promote products derived
 *   from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS"  AND ANY EXPRESS  OR IMPLIED WARRANTIES  OF MERCHANTABILITY
 * AND FITNESS  FOR A PARTICULAR  PURPOSE ARE DISCLAIMED. IN  NO EVENT
 * SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL,  EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED  TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE,  DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT  LIABILITY,  OR  TORT  (INCLUDING NEGLIGENCE  OR  OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 */

package main

import (
	"bytes"
	"crypto/sha1"
	"database/cassandra"
	"encoding/hex"
	"errors"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	"image/png"
	"sync"
	"time"

	"github.com/nfnt/resize"
)

// Database access routines for the image database.
type PikShrDB struct {
	db      *cassandra.RetryCassandraClient
	db_lock sync.Mutex
}

type Picture struct {
	Id          string
	Contents    []byte
	ContentType string
	Title       string
	Description string
	AltText     string
}

var Err_ImageNotFound = errors.New("Image not found")

// Connect to the picture database.
func NewPikShrDB(server, keyspace string) (*PikShrDB, error) {
	var ire *cassandra.InvalidRequestException
	var client *cassandra.RetryCassandraClient
	var err error

	client, err = cassandra.NewRetryCassandraClient(server)
	if err != nil {
		return nil, err
	}

	ire, err = client.SetKeyspace(keyspace)
	if ire != nil {
		return nil, errors.New(ire.Why)
	}
	if err != nil {
		return nil, err
	}

	return &PikShrDB{
		db: client,
	}, nil
}

// Retrieve image data and metadata for the given ID.
// Image data will be taken from "column".
func (p *PikShrDB) getPictureColumn(id, column string) (*Picture, error) {
	var r []*cassandra.ColumnOrSuperColumn
	var scol *cassandra.ColumnOrSuperColumn
	var ret *Picture = new(Picture)
	var cp *cassandra.ColumnParent = cassandra.NewColumnParent()
	var pred *cassandra.SlicePredicate = cassandra.NewSlicePredicate()
	var ire *cassandra.InvalidRequestException
	var ue *cassandra.UnavailableException
	var te *cassandra.TimedOutException
	var key []byte
	var err error

	cp.ColumnFamily = "picture"
	pred.ColumnNames = [][]byte{
		[]byte("title"), []byte("description"), []byte("content_type"),
		[]byte("alt_text"),
	}

	if len(column) > 0 {
		pred.ColumnNames = append(pred.ColumnNames, []byte(column))
	}

	ret.Id = id
	key, err = hex.DecodeString(id)
	if err != nil {
		return nil, err
	}

	p.db_lock.Lock()
	r, ire, ue, te, err = p.db.GetSlice(
		key, cp, pred, cassandra.ConsistencyLevel_ONE)
	p.db_lock.Unlock()
	if ire != nil {
		return nil, errors.New(ire.Why)
	}
	if ue != nil {
		return nil, errors.New("Unavailable")
	}
	if te != nil {
		return nil, errors.New("Timed out")
	}
	if err != nil {
		return nil, err
	}

	if len(r) == 0 {
		return nil, Err_ImageNotFound
	}

	for _, scol = range r {
		var col = scol.Column
		var cname string = string(col.Name)

		if cname == "title" {
			ret.Title = string(col.Value)
		} else if cname == "description" {
			ret.Description = string(col.Value)
		} else if cname == "content_type" {
			ret.ContentType = string(col.Value)
		} else if cname == "alt_text" {
			ret.AltText = string(col.Value)
		} else if len(column) > 0 && cname == column {
			ret.Contents = make([]byte, len(col.Value))
			copy(ret.Contents, col.Value)
		}
	}

	return ret, nil
}

// Retrieve the picture with the given ID from the database.
func (p *PikShrDB) GetPicture(id string) (*Picture, error) {
	return p.getPictureColumn(id, "picture")
}

// Retrieve the thumbnail with the given ID from the database.
func (p *PikShrDB) GetThumbnail(id string) (*Picture, error) {
	return p.getPictureColumn(id, "thumbnail")
}

// Retrieve only the metadata of the picture.
func (p *PikShrDB) GetMetadata(id string) (*Picture, error) {
	return p.getPictureColumn(id, "")
}

func makeMutation(name string, value []byte, now time.Time) (ret *cassandra.Mutation) {
	var cos = cassandra.NewColumnOrSuperColumn()
	var col = cassandra.NewColumn()

	col.Name = []byte(name)
	col.Value = value
	col.Timestamp = now.UnixNano()

	cos.Column = col

	ret = cassandra.NewMutation()
	ret.ColumnOrSupercolumn = cos
	return
}

// Insert the given picture and metadata into the database.
// Also, generate a thumbnail.
func (p *PikShrDB) InsertPicture(pic *Picture, creator string) (string, error) {
	var mutation *cassandra.Mutation
	var ire *cassandra.InvalidRequestException
	var ue *cassandra.UnavailableException
	var te *cassandra.TimedOutException
	var mmap = make(map[string]map[string][]*cassandra.Mutation)
	var mlist []*cassandra.Mutation
	var now = time.Now()
	var img, thumbnail image.Image
	var buf *bytes.Buffer = new(bytes.Buffer)
	var rd *bytes.Reader = bytes.NewReader(pic.Contents)
	var keystr string
	var key [sha1.Size]byte
	var err error

	// First, let's read and parse the image.
	img, _, err = image.Decode(rd)
	if err != nil {
		return "", err
	}

	// Write the full image out as a png.
	err = png.Encode(buf, img)
	if err != nil {
		return "", err
	}

	// Compute a hash of the image data.
	key = sha1.Sum(buf.Bytes())
	keystr = hex.EncodeToString(key[:])

	mmap[string(key[:])] = make(map[string][]*cassandra.Mutation)
	mlist = make([]*cassandra.Mutation, 0)

	mutation = makeMutation("picture", buf.Bytes(), now)
	mlist = append(mlist, mutation)

	// Create a thumbnail and save it too.
	thumbnail = resize.Thumbnail(200, 200, img, resize.Lanczos3)
	err = png.Encode(buf, thumbnail)
	if err != nil {
		return keystr, err
	}
	img = nil

	// Write the thumbnail out as a png.
	buf = new(bytes.Buffer)
	err = png.Encode(buf, thumbnail)
	if err != nil {
		return keystr, err
	}
	thumbnail = nil

	mutation = makeMutation("thumbnail", buf.Bytes(), now)
	mlist = append(mlist, mutation)

	// Now for some metadata.
	mutation = makeMutation("title", []byte(pic.Title), now)
	mlist = append(mlist, mutation)

	mutation = makeMutation("description", []byte(pic.Description), now)
	mlist = append(mlist, mutation)

	mutation = makeMutation("owner", []byte(creator), now)
	mlist = append(mlist, mutation)

	mutation = makeMutation("alt_text", []byte(pic.AltText), now)
	mlist = append(mlist, mutation)

	mmap[string(key[:])]["picture"] = mlist

	// Now: write it!
	p.db_lock.Lock()
	defer p.db_lock.Unlock()
	ire, ue, te, err = p.db.AtomicBatchMutate(mmap, cassandra.ConsistencyLevel_QUORUM)
	if ire != nil {
		return keystr, errors.New(ire.Why)
	}
	if ue != nil {
		return keystr, errors.New("Unavailable")
	}
	if te != nil {
		return keystr, errors.New("Timed out")
	}
	return keystr, err
}

// List the num most recent pictures for the given user. If the user name
// is empty, list the most recent pictures for all users.
func (p *PikShrDB) GetRecentPics(user string, num int32) (ret []*Picture, err error) {
	var cp *cassandra.ColumnParent = cassandra.NewColumnParent()
	var pred *cassandra.SlicePredicate = cassandra.NewSlicePredicate()
	var rng *cassandra.KeyRange = cassandra.NewKeyRange()
	var ks []*cassandra.KeySlice
	var slice *cassandra.KeySlice
	var ire *cassandra.InvalidRequestException
	var ue *cassandra.UnavailableException
	var te *cassandra.TimedOutException

	cp.ColumnFamily = "picture"
	pred.ColumnNames = [][]byte{
		[]byte("title"), []byte("description"), []byte("content_type"),
		[]byte("alt_text"),
	}
	rng.StartKey = make([]byte, 0)
	rng.EndKey = make([]byte, 0)
	rng.Count = num

	if len(user) > 0 {
		var indexexpr = cassandra.NewIndexExpression()

		indexexpr.ColumnName = []byte("owner")
		indexexpr.Op = cassandra.IndexOperator_EQ
		indexexpr.Value = []byte(user)

		rng.RowFilter = make([]*cassandra.IndexExpression, 0)
		rng.RowFilter = append(rng.RowFilter, indexexpr)
	}

	p.db_lock.Lock()
	ks, ire, ue, te, err = p.db.GetRangeSlices(
		cp, pred, rng, cassandra.ConsistencyLevel_ONE)
	p.db_lock.Unlock()
	if ire != nil {
		err = errors.New(ire.Why)
		return
	}
	if ue != nil {
		err = errors.New("Unavailable")
		return
	}
	if te != nil {
		err = errors.New("Timed out")
		return
	}
	if err != nil {
		return
	}

	ret = make([]*Picture, 0)

	for _, slice = range ks {
		var pic *Picture = new(Picture)
		var scol *cassandra.ColumnOrSuperColumn

		pic.Id = hex.EncodeToString(slice.Key)

		for _, scol = range slice.Columns {
			var col = scol.Column
			var cname string = string(col.Name)

			if cname == "title" {
				pic.Title = string(col.Value)
			} else if cname == "description" {
				pic.Description = string(col.Value)
			} else if cname == "content_type" {
				pic.ContentType = string(col.Value)
			} else if cname == "alt_text" {
				pic.AltText = string(col.Value)
			}
		}

		ret = append(ret, pic)
	}

	return
}
<|endoftext|>


<commit_before>package maillist

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"reflect"
	"strings"
	"time"

	"gopkg.in/go-playground/validator.v8"

	"github.com/go-gorp/gorp"
)

type table struct {
	name, selectStr string
}

type database struct {
	db     *sql.DB
	dbmap  *gorp.DbMap
	tables map[reflect.Type]table
}

// ErrNotFound returned when an entity is not present in the database
type ErrNotFound struct{}

var validate *validator.Validate

func (err *ErrNotFound) Error() string {
	return "Not found"
}

func openDatabase(address string) (d database, err error) {
	if address == "" {
		return d, errors.New("Database address not set")
	}
	d.db, err = sql.Open("mysql", address+"?charset=utf8&parseTime=True")
	if err != nil {
		return
	}

	dialect := gorp.MySQLDialect{Engine: "InnoDB", Encoding: "UTF8"}
	d.dbmap = &gorp.DbMap{Db: d.db, Dialect: dialect}
	d.tables = make(map[reflect.Type]table)

	// d.addTable(models.MailList{}, "mail_list")
	if validate == nil {
		config := validator.Config{TagName: "validate"}
		validate = validator.New(&config)
	}
	return
}

func (d *database) insert(i interface{}) error {

	reflect.Indirect(reflect.ValueOf(i)).FieldByName("CreateTime").SetInt(time.Now().Unix())

	if err := validate.Struct(i); err != nil {
		return err
	}
	err := d.dbmap.Insert(i)
	return err
}

func (d *database) delete(i interface{}, id int64) error {
	t := reflect.TypeOf(i)
	table, ok := d.tables[t]
	if !ok {
		return fmt.Errorf("Type %s not registered in db", t)
	}

	sql := fmt.Sprintf("update %s set status='deleted' where id=?", table.name)
	_, err := d.dbmap.Exec(sql, id)
	return err
}

func (d *database) update(i interface{}) error {
	if err := validate.Struct(i); err != nil {
		return err
	}
	_, err := d.dbmap.Update(i)
	return err
}

func (d *database) addTable(i interface{}, tableName string) (selectStatement string) {
	tablemap := d.dbmap.AddTableWithName(i, tableName)
	t := reflect.TypeOf(i)

	if _, hasID := t.FieldByName("ID"); hasID {
		tablemap.SetKeys(true, "ID")
	}

	var columns []string
	for _, c := range tablemap.Columns {
		if !c.Transient {
			columns = append(columns, tablemap.TableName+"."+c.ColumnName)
		}
	}

	selectStatement = fmt.Sprintf("%s\n", strings.Join(columns, ","))
	d.tables[t] = table{tableName, selectStatement}
	return selectStatement
}

func (d *database) selectString(i interface{}) string {
	t := reflect.TypeOf(i)
	for t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	table, ok := d.tables[t]
	if !ok {
		log.Fatalf("Type %s not registered in db", t)
	}
	return table.selectStr
}
<commit_msg>adding mysql ref<commit_after>package maillist

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"reflect"
	"strings"
	"time"

	"gopkg.in/go-playground/validator.v8"

	"github.com/go-gorp/gorp"
	_ "github.com/go-sql-driver/mysql" // Overloading mysql for gorp
)

type table struct {
	name, selectStr string
}

type database struct {
	db     *sql.DB
	dbmap  *gorp.DbMap
	tables map[reflect.Type]table
}

// ErrNotFound returned when an entity is not present in the database
type ErrNotFound struct{}

var validate *validator.Validate

func (err *ErrNotFound) Error() string {
	return "Not found"
}

func openDatabase(address string) (d database, err error) {
	if address == "" {
		return d, errors.New("Database address not set")
	}
	d.db, err = sql.Open("mysql", address+"?charset=utf8&parseTime=True")
	if err != nil {
		return
	}

	dialect := gorp.MySQLDialect{Engine: "InnoDB", Encoding: "UTF8"}
	d.dbmap = &gorp.DbMap{Db: d.db, Dialect: dialect}
	d.tables = make(map[reflect.Type]table)

	// d.addTable(models.MailList{}, "mail_list")
	if validate == nil {
		config := validator.Config{TagName: "validate"}
		validate = validator.New(&config)
	}
	return
}

func (d *database) insert(i interface{}) error {

	reflect.Indirect(reflect.ValueOf(i)).FieldByName("CreateTime").SetInt(time.Now().Unix())

	if err := validate.Struct(i); err != nil {
		return err
	}
	err := d.dbmap.Insert(i)
	return err
}

func (d *database) delete(i interface{}, id int64) error {
	t := reflect.TypeOf(i)
	table, ok := d.tables[t]
	if !ok {
		return fmt.Errorf("Type %s not registered in db", t)
	}

	sql := fmt.Sprintf("update %s set status='deleted' where id=?", table.name)
	_, err := d.dbmap.Exec(sql, id)
	return err
}

func (d *database) update(i interface{}) error {
	if err := validate.Struct(i); err != nil {
		return err
	}
	_, err := d.dbmap.Update(i)
	return err
}

func (d *database) addTable(i interface{}, tableName string) (selectStatement string) {
	tablemap := d.dbmap.AddTableWithName(i, tableName)
	t := reflect.TypeOf(i)

	if _, hasID := t.FieldByName("ID"); hasID {
		tablemap.SetKeys(true, "ID")
	}

	var columns []string
	for _, c := range tablemap.Columns {
		if !c.Transient {
			columns = append(columns, tablemap.TableName+"."+c.ColumnName)
		}
	}

	selectStatement = fmt.Sprintf("%s\n", strings.Join(columns, ","))
	d.tables[t] = table{tableName, selectStatement}
	return selectStatement
}

func (d *database) selectString(i interface{}) string {
	t := reflect.TypeOf(i)
	for t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	table, ok := d.tables[t]
	if !ok {
		log.Fatalf("Type %s not registered in db", t)
	}
	return table.selectStr
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	_ "github.com/go-sql-driver/mysql"

	"log"
)

func dbConnection() {
	db, err := sql.Open("mysql",
		"root:@tcp(127.0.0.1:3306)/blog")
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	err = db.Ping()
	if err != nil {
		log.Fatal(err)
	}

}
<commit_msg>Add logic for connecting to database<commit_after>package main

import (
	"database/sql"
	_ "github.com/go-sql-driver/mysql"

	"log"
)

func dbConnection() *sql.DB {
	db, err := sql.Open("mysql",
		"root:@tcp(127.0.0.1:3306)/blog")
	if err != nil {
		log.Fatal(err)
	}

	return db
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	_ "github.com/mattn/go-sqlite3"
	"log"
	"os"
	"reflect"
	"time"
)

type DB interface {
	SetChar(int64, time.Time) error
	GetChar(int64) (time.Time, bool, error)
	OldestChar() (int64, time.Time, error)
	RemoveChar(int64) error
	NumChar() int

	LoadSession() (Session, error)
	SaveSession(s Session) error

	Close() error
}

func createDB() (DB, error) {
	switch t := flags.db["type"]; t {
	case "map":
		log.Printf("Using %q for DB.", t)

		if flags.db["s"] == "" {
			flags.db["s"] = "session.json"
		}

		return make(mapDB), nil

	case "sqlite", "sqlite3":
		log.Printf("Using %q for DB.", t)

		if flags.db["db"] == "" {
			flags.db["db"] = "session.db"
		}

		return newsqliteDB(flags.db["db"])
	}

	return nil, fmt.Errorf("Bad db flag value: %v", flags.db)
}

type mapDB map[int64]time.Time

func (db mapDB) SetChar(id int64, login time.Time) error {
	db[id] = login
	return nil
}

func (db mapDB) GetChar(id int64) (time.Time, bool, error) {
	login, ok := db[id]
	return login, ok, nil
}

func (db mapDB) OldestChar() (int64, time.Time, error) {
	// TODO: Implement.
	panic("Not implemented.")
}

func (db mapDB) RemoveChar(id int64) error {
	delete(db, id)
	return nil
}

func (db mapDB) NumChar() int {
	return len(db)
}

func (db mapDB) LoadSession() (s Session, err error) {
	defer func() {
		s.db = db
	}()

	file, err := os.Open(flags.db["s"])
	if err != nil {
		return s, err
	}
	defer file.Close()

	d := json.NewDecoder(file)
	err = d.Decode(&s)
	return s, err
}

func (db mapDB) SaveSession(s Session) error {
	file, err := os.Create(flags.db["s"])
	if err != nil {
		return err
	}
	defer file.Close()

	e := json.NewEncoder(file)
	return e.Encode(&s)
}

func (db mapDB) Close() error {
	return nil
}

// TODO: Find a way to do some of this asynchronously.
type sqliteDB struct {
	*sql.DB

	add    *sql.Stmt
	get    *sql.Stmt
	oldest *sql.Stmt
	rem    *sql.Stmt
	num    *sql.Stmt

	sadd *sql.Stmt
	sget *sql.Stmt
}

func newsqliteDB(path string) (DB, error) {
	db, err := sql.Open("sqlite3", path)
	if err != nil {
		return nil, err
	}

	_, err = db.Exec(`DROP TABLE IF EXISTS chars`)
	if err != nil {
		return nil, err
	}

	_, err = db.Exec(`CREATE TABLE chars (id INTEGER PRIMARY KEY, login TIMESTAMP)`)
	if err != nil {
		return nil, err
	}

	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS session (id TEXT PRIMARY KEY, valstr TEXT, valint INTEGER)`)
	if err != nil {
		return nil, err
	}

	add, err := db.Prepare(`INSERT OR REPLACE INTO chars (id, login) VALUES (?, ?)`)
	if err != nil {
		return nil, err
	}

	get, err := db.Prepare(`SELECT login FROM chars WHERE id=?`)
	if err != nil {
		return nil, err
	}

	oldest, err := db.Prepare(`SELECT id, login FROM chars ORDER BY login LIMIT 1`)
	if err != nil {
		return nil, err
	}

	rem, err := db.Prepare(`DELETE FROM chars WHERE id=?`)
	if err != nil {
		return nil, err
	}

	num, err := db.Prepare(`SELECT count(id) FROM chars`)
	if err != nil {
		return nil, err
	}

	sadd, err := db.Prepare(`INSERT OR REPLACE INTO session (id, valstr, valint) VALUES (?, ?, ?)`)
	if err != nil {
		return nil, err
	}

	sget, err := db.Prepare(`SELECT valstr, valint FROM session WHERE id=?`)
	if err != nil {
		return nil, err
	}

	return &sqliteDB{
		DB: db,

		add:    add,
		get:    get,
		oldest: oldest,
		rem:    rem,
		num:    num,

		sadd: sadd,
		sget: sget,
	}, nil
}

func (db *sqliteDB) SetChar(id int64, login time.Time) error {
	_, err := db.add.Exec(id, login)
	return err
}

func (db *sqliteDB) GetChar(id int64) (time.Time, bool, error) {
	var login time.Time
	err := db.get.QueryRow(id).Scan(&login)
	if err != nil {
		if err == sql.ErrNoRows {
			return login, false, nil
		}

		return login, false, err
	}

	return login, true, nil
}

func (db *sqliteDB) OldestChar() (id int64, t time.Time, err error) {
	err = db.oldest.QueryRow().Scan(&id, &t)
	return
}

func (db *sqliteDB) RemoveChar(id int64) error {
	_, err := db.rem.Exec(id)
	return err
}

func (db *sqliteDB) NumChar() (n int) {
	err := db.num.QueryRow().Scan(&n)
	if err != nil {
		log.Printf("Failed to get number of active characters: %v", err)
	}

	return n
}

func (db *sqliteDB) LoadSession() (s Session, err error) {
	err = walkStruct(&s, func(name string, field reflect.Value) error {
		var valstr string
		var valint int64
		err := db.sget.QueryRow(name).Scan(&valstr, &valint)
		if err != nil {
			if err == sql.ErrNoRows {
				// Just ignore fields that aren't in the database.
				return nil
			}

			return err
		}

		switch field.Kind() {
		case reflect.String:
			field.SetString(valstr)
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			field.SetInt(valint)
		}

		return nil
	})

	s.db = db
	return
}

func (db *sqliteDB) SaveSession(s Session) error {
	return walkStruct(&s, func(name string, field reflect.Value) (err error) {
		switch field.Kind() {
		case reflect.String:
			_, err = db.sadd.Exec(name, field.String(), 0)
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			_, err = db.sadd.Exec(name, "", field.Int())
		}
		if err != nil {
			log.Fatalf("Failed to save %q (%v): %v", name, field.Interface(), err)
		}

		return
	})
}
<commit_msg>Get oldest session from SQLite without using `order by`.<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	_ "github.com/mattn/go-sqlite3"
	"log"
	"os"
	"reflect"
	"time"
)

type DB interface {
	SetChar(int64, time.Time) error
	GetChar(int64) (time.Time, bool, error)
	OldestChar() (int64, time.Time, error)
	RemoveChar(int64) error
	NumChar() int

	LoadSession() (Session, error)
	SaveSession(s Session) error

	Close() error
}

func createDB() (DB, error) {
	switch t := flags.db["type"]; t {
	case "map":
		log.Printf("Using %q for DB.", t)

		if flags.db["s"] == "" {
			flags.db["s"] = "session.json"
		}

		return make(mapDB), nil

	case "sqlite", "sqlite3":
		log.Printf("Using %q for DB.", t)

		if flags.db["db"] == "" {
			flags.db["db"] = "session.db"
		}

		return newsqliteDB(flags.db["db"])
	}

	return nil, fmt.Errorf("Bad db flag value: %v", flags.db)
}

type mapDB map[int64]time.Time

func (db mapDB) SetChar(id int64, login time.Time) error {
	db[id] = login
	return nil
}

func (db mapDB) GetChar(id int64) (time.Time, bool, error) {
	login, ok := db[id]
	return login, ok, nil
}

func (db mapDB) OldestChar() (int64, time.Time, error) {
	// TODO: Implement.
	panic("Not implemented.")
}

func (db mapDB) RemoveChar(id int64) error {
	delete(db, id)
	return nil
}

func (db mapDB) NumChar() int {
	return len(db)
}

func (db mapDB) LoadSession() (s Session, err error) {
	defer func() {
		s.db = db
	}()

	file, err := os.Open(flags.db["s"])
	if err != nil {
		return s, err
	}
	defer file.Close()

	d := json.NewDecoder(file)
	err = d.Decode(&s)
	return s, err
}

func (db mapDB) SaveSession(s Session) error {
	file, err := os.Create(flags.db["s"])
	if err != nil {
		return err
	}
	defer file.Close()

	e := json.NewEncoder(file)
	return e.Encode(&s)
}

func (db mapDB) Close() error {
	return nil
}

// TODO: Find a way to do some of this asynchronously.
type sqliteDB struct {
	*sql.DB

	add    *sql.Stmt
	get    *sql.Stmt
	oldest *sql.Stmt
	rem    *sql.Stmt
	num    *sql.Stmt

	sadd *sql.Stmt
	sget *sql.Stmt
}

func newsqliteDB(path string) (DB, error) {
	db, err := sql.Open("sqlite3", path)
	if err != nil {
		return nil, err
	}

	_, err = db.Exec(`DROP TABLE IF EXISTS chars`)
	if err != nil {
		return nil, err
	}

	_, err = db.Exec(`CREATE TABLE chars (id INTEGER PRIMARY KEY, login TIMESTAMP)`)
	if err != nil {
		return nil, err
	}

	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS session (id TEXT PRIMARY KEY, valstr TEXT, valint INTEGER)`)
	if err != nil {
		return nil, err
	}

	add, err := db.Prepare(`INSERT OR REPLACE INTO chars (id, login) VALUES (?, ?)`)
	if err != nil {
		return nil, err
	}

	get, err := db.Prepare(`SELECT login FROM chars WHERE id=?`)
	if err != nil {
		return nil, err
	}

	oldest, err := db.Prepare(`SELECT id, login, min(login) FROM chars`)
	if err != nil {
		return nil, err
	}

	rem, err := db.Prepare(`DELETE FROM chars WHERE id=?`)
	if err != nil {
		return nil, err
	}

	num, err := db.Prepare(`SELECT count(id) FROM chars`)
	if err != nil {
		return nil, err
	}

	sadd, err := db.Prepare(`INSERT OR REPLACE INTO session (id, valstr, valint) VALUES (?, ?, ?)`)
	if err != nil {
		return nil, err
	}

	sget, err := db.Prepare(`SELECT valstr, valint FROM session WHERE id=?`)
	if err != nil {
		return nil, err
	}

	return &sqliteDB{
		DB: db,

		add:    add,
		get:    get,
		oldest: oldest,
		rem:    rem,
		num:    num,

		sadd: sadd,
		sget: sget,
	}, nil
}

func (db *sqliteDB) SetChar(id int64, login time.Time) error {
	_, err := db.add.Exec(id, login)
	return err
}

func (db *sqliteDB) GetChar(id int64) (time.Time, bool, error) {
	var login time.Time
	err := db.get.QueryRow(id).Scan(&login)
	if err != nil {
		if err == sql.ErrNoRows {
			return login, false, nil
		}

		return login, false, err
	}

	return login, true, nil
}

func (db *sqliteDB) OldestChar() (id int64, t time.Time, err error) {
	var min string
	err = db.oldest.QueryRow().Scan(&id, &t, &min)
	return
}

func (db *sqliteDB) RemoveChar(id int64) error {
	_, err := db.rem.Exec(id)
	return err
}

func (db *sqliteDB) NumChar() (n int) {
	err := db.num.QueryRow().Scan(&n)
	if err != nil {
		log.Printf("Failed to get number of active characters: %v", err)
	}

	return n
}

func (db *sqliteDB) LoadSession() (s Session, err error) {
	err = walkStruct(&s, func(name string, field reflect.Value) error {
		var valstr string
		var valint int64
		err := db.sget.QueryRow(name).Scan(&valstr, &valint)
		if err != nil {
			if err == sql.ErrNoRows {
				// Just ignore fields that aren't in the database.
				return nil
			}

			return err
		}

		switch field.Kind() {
		case reflect.String:
			field.SetString(valstr)
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			field.SetInt(valint)
		}

		return nil
	})

	s.db = db
	return
}

func (db *sqliteDB) SaveSession(s Session) error {
	return walkStruct(&s, func(name string, field reflect.Value) (err error) {
		switch field.Kind() {
		case reflect.String:
			_, err = db.sadd.Exec(name, field.String(), 0)
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			_, err = db.sadd.Exec(name, "", field.Int())
		}
		if err != nil {
			log.Fatalf("Failed to save %q (%v): %v", name, field.Interface(), err)
		}

		return
	})
}
<|endoftext|>


<commit_before>package query

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"

	"github.com/yuuki/diamondb/lib/model"
	"github.com/yuuki/diamondb/lib/tsdb"
)

func TestEvalTarget_Func(t *testing.T) {
	ctrl := tsdb.SetMockDynamoDB(t, &tsdb.MockDynamoDBParam{
		TableName: tsdb.DynamoDBTableOneHour + "-0",
		ItemEpoch: 0,
		Metrics: []*model.Metric{
			model.NewMetric(
				"Sales.widgets.largeBlue",
				[]*model.DataPoint{
					{60, 10.0},
				},
				60,
			),
		},
	})
	defer ctrl.Finish()

	metrics, err := EvalTarget(
		"alias(Sales.widgets.largeBlue,\"Large Blue Widgets\")",
		time.Unix(0, 0),
		time.Unix(120, 0),
	)
	if assert.NoError(t, err) {
		assert.Exactly(t, 1, len(metrics))
		assert.Exactly(t, metrics[0].Name, "Large Blue Widgets")
	}
}

func TestEvalTarget_FuncNest(t *testing.T) {
	ctrl := tsdb.SetMockDynamoDB(t, &tsdb.MockDynamoDBParam{
		TableName: tsdb.DynamoDBTableOneHour + "-0",
		ItemEpoch: 0,
		Metrics: []*model.Metric{
			model.NewMetric(
				"Sales.widgets.largeBlue",
				[]*model.DataPoint{
					{60, 10.0},
				},
				60,
			),
		},
	})
	defer ctrl.Finish()

	metrics, err := EvalTarget(
		"alias(alias(Sales.widgets.largeBlue,\"Large Blue Widgets\"),\"Large Blue Widgets Sales\")",
		time.Unix(0, 0),
		time.Unix(120, 0),
	)
	if assert.NoError(t, err) {
		assert.Exactly(t, 1, len(metrics))
		assert.Exactly(t, metrics[0].Name, "Large Blue Widgets Sales")
	}
}
<commit_msg>Remove assert from lib/query/engine_test.go<commit_after>package query

import (
	"reflect"
	"testing"
	"time"

	"github.com/yuuki/diamondb/lib/model"
	"github.com/yuuki/diamondb/lib/tsdb"
)

func TestEvalTarget_Func(t *testing.T) {
	points := []*model.DataPoint{
		{60, 10.0},
		{120, 11.0},
	}
	ctrl := tsdb.SetMockDynamoDB(t, &tsdb.MockDynamoDBParam{
		TableName: tsdb.DynamoDBTableOneHour + "-0",
		ItemEpoch: 0,
		Metrics: []*model.Metric{
			model.NewMetric(
				"server1.loadavg5",
				points,
				60,
			),
		},
	})
	defer ctrl.Finish()

	metrics, err := EvalTarget(
		"alias(server1.loadavg5,\"server01.loadavg5\")",
		time.Unix(0, 0),
		time.Unix(120, 0),
	)

	if err != nil {
		t.Fatalf("err: %s", err)
	}
	expected := []*model.Metric{
		{Name: "server01.loadavg5", Start: 60, End: 120, Step: 60, DataPoints: points},
	}
	if !reflect.DeepEqual(metrics, expected) {
		t.Fatalf("\nExpected: %+v\nActual:   %+v", expected, metrics)
	}
}

func TestEvalTarget_FuncNest(t *testing.T) {
	points := []*model.DataPoint{
		{60, 10.0},
		{120, 11.0},
	}
	ctrl := tsdb.SetMockDynamoDB(t, &tsdb.MockDynamoDBParam{
		TableName: tsdb.DynamoDBTableOneHour + "-0",
		ItemEpoch: 0,
		Metrics: []*model.Metric{
			model.NewMetric(
				"server1.loadavg5",
				points,
				60,
			),
		},
	})
	defer ctrl.Finish()

	metrics, err := EvalTarget(
		"alias(alias(server1.loadavg5,\"server01.loadavg5\"),\"server001.loadavg5\")",
		time.Unix(0, 0),
		time.Unix(120, 0),
	)

	if err != nil {
		t.Fatalf("err: %s", err)
	}
	expected := []*model.Metric{
		{Name: "server001.loadavg5", Start: 60, End: 120, Step: 60, DataPoints: points},
	}
	if !reflect.DeepEqual(metrics, expected) {
		t.Fatalf("\nExpected: %+v\nActual:   %+v", expected, metrics)
	}
}
<|endoftext|>


<commit_before>// package peer implements an object used to represent peers in the ipfs network.
package peer

import (
	"bytes"
	"errors"
	"fmt"
	"sync"
	"time"

	b58 "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-base58"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
	mh "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multihash"

	ic "github.com/jbenet/go-ipfs/crypto"
	u "github.com/jbenet/go-ipfs/util"
)

var log = u.Logger("peer")

// ID is a byte slice representing the identity of a peer.
type ID mh.Multihash

// String is utililty function for printing out peer ID strings.
func (id ID) String() string {
	return id.Pretty()
}

// Equal is utililty function for comparing two peer ID's
func (id ID) Equal(other ID) bool {
	return bytes.Equal(id, other)
}

// Pretty returns a b58-encoded string of the ID
func (id ID) Pretty() string {
	return b58.Encode(id)
}

// DecodePrettyID returns a b58-encoded string of the ID
func DecodePrettyID(s string) ID {
	return b58.Decode(s)
}

// IDFromPubKey retrieves a Public Key from the peer given by pk
func IDFromPubKey(pk ic.PubKey) (ID, error) {
	b, err := pk.Bytes()
	if err != nil {
		return nil, err
	}
	hash := u.Hash(b)
	return ID(hash), nil
}

// Map maps Key (string) : *peer (slices are not comparable).
type Map map[u.Key]Peer

// Peer represents the identity information of an IPFS Node, including
// ID, and relevant Addresses.
type Peer interface {

	// TODO reduce the peer interface to be read-only. Force mutations to occur
	// on the peer store eg. peerstore.SetLatency(peerId, value).

	// ID returns the peer's ID
	ID() ID

	// Key returns the ID as a Key (string) for maps.
	Key() u.Key

	// Addresses returns the peer's multiaddrs
	Addresses() []ma.Multiaddr

	// AddAddress adds the given Multiaddr address to Peer's addresses.
	// returns whether this was a newly added address.
	AddAddress(a ma.Multiaddr) bool

	// NetAddress returns the first Multiaddr found for a given network.
	NetAddress(n string) ma.Multiaddr

	// Services returns the peer's services
	// Services() []mux.ProtocolID
	// SetServices([]mux.ProtocolID)

	// Priv/PubKey returns the peer's Private Key
	PrivKey() ic.PrivKey
	PubKey() ic.PubKey

	// LoadAndVerifyKeyPair unmarshalls, loads a private/public key pair.
	// Error if (a) unmarshalling fails, or (b) pubkey does not match id.
	LoadAndVerifyKeyPair(marshalled []byte) error
	VerifyAndSetPrivKey(sk ic.PrivKey) error
	VerifyAndSetPubKey(pk ic.PubKey) error

	// Get/SetLatency manipulate the current latency measurement.
	GetLatency() (out time.Duration)
	SetLatency(laten time.Duration)

	// Update with the data of another peer instance
	Update(Peer) error
}

type peer struct {
	id        ID
	addresses []ma.Multiaddr
	// services  []mux.ProtocolID

	privKey ic.PrivKey
	pubKey  ic.PubKey

	// TODO move latency away from peer into the package that uses it. Instead,
	// within that package, map from ID to latency value.
	latency time.Duration

	sync.RWMutex
}

// String prints out the peer.
func (p *peer) String() string {
	return "[Peer " + p.id.String()[:12] + "]"
}

// Key returns the ID as a Key (string) for maps.
func (p *peer) Key() u.Key {
	return u.Key(p.id)
}

// ID returns the peer's ID
func (p *peer) ID() ID {
	return p.id
}

// PrivKey returns the peer's Private Key
func (p *peer) PrivKey() ic.PrivKey {
	return p.privKey
}

// PubKey returns the peer's Private Key
func (p *peer) PubKey() ic.PubKey {
	return p.pubKey
}

// Addresses returns the peer's multiaddrs
func (p *peer) Addresses() []ma.Multiaddr {
	cp := make([]ma.Multiaddr, len(p.addresses))
	p.RLock()
	copy(cp, p.addresses)
	defer p.RUnlock()
	return cp
}

// AddAddress adds the given Multiaddr address to Peer's addresses.
// Returns whether this address was a newly added address
func (p *peer) AddAddress(a ma.Multiaddr) bool {
	p.Lock()
	defer p.Unlock()

	for _, addr := range p.addresses {
		if addr.Equal(a) {
			return false
		}
	}
	p.addresses = append(p.addresses, a)
	return true
}

// NetAddress returns the first Multiaddr found for a given network.
func (p *peer) NetAddress(n string) ma.Multiaddr {
	p.RLock()
	defer p.RUnlock()

	for _, a := range p.addresses {
		for _, p := range a.Protocols() {
			if p.Name == n {
				return a
			}
		}
	}
	return nil
}

// func (p *peer) Services() []mux.ProtocolID {
// 	p.RLock()
// 	defer p.RUnlock()
// 	return p.services
// }
//
// func (p *peer) SetServices(s []mux.ProtocolID) {
// 	p.Lock()
// 	defer p.Unlock()
// 	p.services = s
// }

// GetLatency retrieves the current latency measurement.
func (p *peer) GetLatency() (out time.Duration) {
	p.RLock()
	out = p.latency
	p.RUnlock()
	return
}

// SetLatency sets the latency measurement.
// TODO: Instead of just keeping a single number,
//		 keep a running average over the last hour or so
// Yep, should be EWMA or something. (-jbenet)
func (p *peer) SetLatency(laten time.Duration) {
	p.Lock()
	if p.latency == 0 {
		p.latency = laten
	} else {
		p.latency = ((p.latency * 9) + laten) / 10
	}
	p.Unlock()
}

// LoadAndVerifyKeyPair unmarshalls, loads a private/public key pair.
// Error if (a) unmarshalling fails, or (b) pubkey does not match id.
func (p *peer) LoadAndVerifyKeyPair(marshalled []byte) error {
	sk, err := ic.UnmarshalPrivateKey(marshalled)
	if err != nil {
		return fmt.Errorf("Failed to unmarshal private key: %v", err)
	}

	return p.VerifyAndSetPrivKey(sk)
}

// VerifyAndSetPrivKey sets private key, given its pubkey matches the peer.ID
func (p *peer) VerifyAndSetPrivKey(sk ic.PrivKey) error {

	// construct and assign pubkey. ensure it matches this peer
	if err := p.VerifyAndSetPubKey(sk.GetPublic()); err != nil {
		return err
	}

	p.Lock()
	defer p.Unlock()

	// if we didn't have the priavte key, assign it
	if p.privKey == nil {
		p.privKey = sk
		return nil
	}

	// if we already had the keys, check they're equal.
	if p.privKey.Equals(sk) {
		return nil // as expected. keep the old objects.
	}

	// keys not equal. invariant violated. this warrants a panic.
	// these keys should be _the same_ because peer.ID = H(pk)
	// this mismatch should never happen.
	log.Errorf("%s had PrivKey: %v -- got %v", p, p.privKey, sk)
	panic("invariant violated: unexpected key mismatch")
}

// VerifyAndSetPubKey sets public key, given it matches the peer.ID
func (p *peer) VerifyAndSetPubKey(pk ic.PubKey) error {
	pkid, err := IDFromPubKey(pk)
	if err != nil {
		return fmt.Errorf("Failed to hash public key: %v", err)
	}

	p.Lock()
	defer p.Unlock()

	if !p.id.Equal(pkid) {
		return fmt.Errorf("Public key does not match peer.ID.")
	}

	// if we didn't have the keys, assign them.
	if p.pubKey == nil {
		p.pubKey = pk
		return nil
	}

	// if we already had the pubkey, check they're equal.
	if p.pubKey.Equals(pk) {
		return nil // as expected. keep the old objects.
	}

	// keys not equal. invariant violated. this warrants a panic.
	// these keys should be _the same_ because peer.ID = H(pk)
	// this mismatch should never happen.
	log.Errorf("%s had PubKey: %v -- got %v", p, p.pubKey, pk)
	panic("invariant violated: unexpected key mismatch")
}

// Updates this peer with information from another peer instance
func (p *peer) Update(other Peer) error {
	if !p.ID().Equal(other.ID()) {
		return errors.New("peer ids do not match")
	}

	for _, a := range other.Addresses() {
		p.AddAddress(a)
	}

	p.SetLatency(other.GetLatency())

	sk := other.PrivKey()
	pk := other.PubKey()
	p.Lock()
	if p.privKey == nil {
		p.privKey = sk
	}

	if p.pubKey == nil {
		p.pubKey = pk
	}
	defer p.Unlock()
	return nil
}

// WithKeyPair returns a Peer object with given keys.
func WithKeyPair(sk ic.PrivKey, pk ic.PubKey) (Peer, error) {
	if sk == nil && pk == nil {
		return nil, fmt.Errorf("PeerWithKeyPair nil keys")
	}

	pk2 := sk.GetPublic()
	if pk == nil {
		pk = pk2
	} else if !pk.Equals(pk2) {
		return nil, fmt.Errorf("key mismatch. pubkey is not privkey's pubkey")
	}

	pkid, err := IDFromPubKey(pk)
	if err != nil {
		return nil, fmt.Errorf("Failed to hash public key: %v", err)
	}

	return &peer{id: pkid, pubKey: pk, privKey: sk}, nil
}

// WithID constructs a peer with given ID.
func WithID(id ID) Peer {
	return &peer{id: id}
}

// WithIDString constructs a peer with given ID (string).
func WithIDString(id string) Peer {
	return WithID(ID(id))
}
<commit_msg>fix(peer) perform bounds check on peer Id when printing String<commit_after>// package peer implements an object used to represent peers in the ipfs network.
package peer

import (
	"bytes"
	"errors"
	"fmt"
	"sync"
	"time"

	b58 "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-base58"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
	mh "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multihash"

	ic "github.com/jbenet/go-ipfs/crypto"
	u "github.com/jbenet/go-ipfs/util"
)

var log = u.Logger("peer")

// ID is a byte slice representing the identity of a peer.
type ID mh.Multihash

// String is utililty function for printing out peer ID strings.
func (id ID) String() string {
	return id.Pretty()
}

// Equal is utililty function for comparing two peer ID's
func (id ID) Equal(other ID) bool {
	return bytes.Equal(id, other)
}

// Pretty returns a b58-encoded string of the ID
func (id ID) Pretty() string {
	return b58.Encode(id)
}

// DecodePrettyID returns a b58-encoded string of the ID
func DecodePrettyID(s string) ID {
	return b58.Decode(s)
}

// IDFromPubKey retrieves a Public Key from the peer given by pk
func IDFromPubKey(pk ic.PubKey) (ID, error) {
	b, err := pk.Bytes()
	if err != nil {
		return nil, err
	}
	hash := u.Hash(b)
	return ID(hash), nil
}

// Map maps Key (string) : *peer (slices are not comparable).
type Map map[u.Key]Peer

// Peer represents the identity information of an IPFS Node, including
// ID, and relevant Addresses.
type Peer interface {

	// TODO reduce the peer interface to be read-only. Force mutations to occur
	// on the peer store eg. peerstore.SetLatency(peerId, value).

	// ID returns the peer's ID
	ID() ID

	// Key returns the ID as a Key (string) for maps.
	Key() u.Key

	// Addresses returns the peer's multiaddrs
	Addresses() []ma.Multiaddr

	// AddAddress adds the given Multiaddr address to Peer's addresses.
	// returns whether this was a newly added address.
	AddAddress(a ma.Multiaddr) bool

	// NetAddress returns the first Multiaddr found for a given network.
	NetAddress(n string) ma.Multiaddr

	// Services returns the peer's services
	// Services() []mux.ProtocolID
	// SetServices([]mux.ProtocolID)

	// Priv/PubKey returns the peer's Private Key
	PrivKey() ic.PrivKey
	PubKey() ic.PubKey

	// LoadAndVerifyKeyPair unmarshalls, loads a private/public key pair.
	// Error if (a) unmarshalling fails, or (b) pubkey does not match id.
	LoadAndVerifyKeyPair(marshalled []byte) error
	VerifyAndSetPrivKey(sk ic.PrivKey) error
	VerifyAndSetPubKey(pk ic.PubKey) error

	// Get/SetLatency manipulate the current latency measurement.
	GetLatency() (out time.Duration)
	SetLatency(laten time.Duration)

	// Update with the data of another peer instance
	Update(Peer) error
}

type peer struct {
	id        ID
	addresses []ma.Multiaddr
	// services  []mux.ProtocolID

	privKey ic.PrivKey
	pubKey  ic.PubKey

	// TODO move latency away from peer into the package that uses it. Instead,
	// within that package, map from ID to latency value.
	latency time.Duration

	sync.RWMutex
}

// String prints out the peer.
func (p *peer) String() string {
	pid := p.id.String()
	maxRunes := 12
	if len(pid) < maxRunes {
		maxRunes = len(pid)
	}
	return "[Peer " + pid[:maxRunes] + "]"
}

// Key returns the ID as a Key (string) for maps.
func (p *peer) Key() u.Key {
	return u.Key(p.id)
}

// ID returns the peer's ID
func (p *peer) ID() ID {
	return p.id
}

// PrivKey returns the peer's Private Key
func (p *peer) PrivKey() ic.PrivKey {
	return p.privKey
}

// PubKey returns the peer's Private Key
func (p *peer) PubKey() ic.PubKey {
	return p.pubKey
}

// Addresses returns the peer's multiaddrs
func (p *peer) Addresses() []ma.Multiaddr {
	cp := make([]ma.Multiaddr, len(p.addresses))
	p.RLock()
	copy(cp, p.addresses)
	defer p.RUnlock()
	return cp
}

// AddAddress adds the given Multiaddr address to Peer's addresses.
// Returns whether this address was a newly added address
func (p *peer) AddAddress(a ma.Multiaddr) bool {
	p.Lock()
	defer p.Unlock()

	for _, addr := range p.addresses {
		if addr.Equal(a) {
			return false
		}
	}
	p.addresses = append(p.addresses, a)
	return true
}

// NetAddress returns the first Multiaddr found for a given network.
func (p *peer) NetAddress(n string) ma.Multiaddr {
	p.RLock()
	defer p.RUnlock()

	for _, a := range p.addresses {
		for _, p := range a.Protocols() {
			if p.Name == n {
				return a
			}
		}
	}
	return nil
}

// func (p *peer) Services() []mux.ProtocolID {
// 	p.RLock()
// 	defer p.RUnlock()
// 	return p.services
// }
//
// func (p *peer) SetServices(s []mux.ProtocolID) {
// 	p.Lock()
// 	defer p.Unlock()
// 	p.services = s
// }

// GetLatency retrieves the current latency measurement.
func (p *peer) GetLatency() (out time.Duration) {
	p.RLock()
	out = p.latency
	p.RUnlock()
	return
}

// SetLatency sets the latency measurement.
// TODO: Instead of just keeping a single number,
//		 keep a running average over the last hour or so
// Yep, should be EWMA or something. (-jbenet)
func (p *peer) SetLatency(laten time.Duration) {
	p.Lock()
	if p.latency == 0 {
		p.latency = laten
	} else {
		p.latency = ((p.latency * 9) + laten) / 10
	}
	p.Unlock()
}

// LoadAndVerifyKeyPair unmarshalls, loads a private/public key pair.
// Error if (a) unmarshalling fails, or (b) pubkey does not match id.
func (p *peer) LoadAndVerifyKeyPair(marshalled []byte) error {
	sk, err := ic.UnmarshalPrivateKey(marshalled)
	if err != nil {
		return fmt.Errorf("Failed to unmarshal private key: %v", err)
	}

	return p.VerifyAndSetPrivKey(sk)
}

// VerifyAndSetPrivKey sets private key, given its pubkey matches the peer.ID
func (p *peer) VerifyAndSetPrivKey(sk ic.PrivKey) error {

	// construct and assign pubkey. ensure it matches this peer
	if err := p.VerifyAndSetPubKey(sk.GetPublic()); err != nil {
		return err
	}

	p.Lock()
	defer p.Unlock()

	// if we didn't have the priavte key, assign it
	if p.privKey == nil {
		p.privKey = sk
		return nil
	}

	// if we already had the keys, check they're equal.
	if p.privKey.Equals(sk) {
		return nil // as expected. keep the old objects.
	}

	// keys not equal. invariant violated. this warrants a panic.
	// these keys should be _the same_ because peer.ID = H(pk)
	// this mismatch should never happen.
	log.Errorf("%s had PrivKey: %v -- got %v", p, p.privKey, sk)
	panic("invariant violated: unexpected key mismatch")
}

// VerifyAndSetPubKey sets public key, given it matches the peer.ID
func (p *peer) VerifyAndSetPubKey(pk ic.PubKey) error {
	pkid, err := IDFromPubKey(pk)
	if err != nil {
		return fmt.Errorf("Failed to hash public key: %v", err)
	}

	p.Lock()
	defer p.Unlock()

	if !p.id.Equal(pkid) {
		return fmt.Errorf("Public key does not match peer.ID.")
	}

	// if we didn't have the keys, assign them.
	if p.pubKey == nil {
		p.pubKey = pk
		return nil
	}

	// if we already had the pubkey, check they're equal.
	if p.pubKey.Equals(pk) {
		return nil // as expected. keep the old objects.
	}

	// keys not equal. invariant violated. this warrants a panic.
	// these keys should be _the same_ because peer.ID = H(pk)
	// this mismatch should never happen.
	log.Errorf("%s had PubKey: %v -- got %v", p, p.pubKey, pk)
	panic("invariant violated: unexpected key mismatch")
}

// Updates this peer with information from another peer instance
func (p *peer) Update(other Peer) error {
	if !p.ID().Equal(other.ID()) {
		return errors.New("peer ids do not match")
	}

	for _, a := range other.Addresses() {
		p.AddAddress(a)
	}

	p.SetLatency(other.GetLatency())

	sk := other.PrivKey()
	pk := other.PubKey()
	p.Lock()
	if p.privKey == nil {
		p.privKey = sk
	}

	if p.pubKey == nil {
		p.pubKey = pk
	}
	defer p.Unlock()
	return nil
}

// WithKeyPair returns a Peer object with given keys.
func WithKeyPair(sk ic.PrivKey, pk ic.PubKey) (Peer, error) {
	if sk == nil && pk == nil {
		return nil, fmt.Errorf("PeerWithKeyPair nil keys")
	}

	pk2 := sk.GetPublic()
	if pk == nil {
		pk = pk2
	} else if !pk.Equals(pk2) {
		return nil, fmt.Errorf("key mismatch. pubkey is not privkey's pubkey")
	}

	pkid, err := IDFromPubKey(pk)
	if err != nil {
		return nil, fmt.Errorf("Failed to hash public key: %v", err)
	}

	return &peer{id: pkid, pubKey: pk, privKey: sk}, nil
}

// WithID constructs a peer with given ID.
func WithID(id ID) Peer {
	return &peer{id: id}
}

// WithIDString constructs a peer with given ID (string).
func WithIDString(id string) Peer {
	return WithID(ID(id))
}
<|endoftext|>