<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/nutrun/lentil"
	"log"
	"path/filepath"
	"strconv"
)

type Client struct {
	q       *lentil.Beanstalkd
	verbose bool
}

func NewClient(verbose bool) (*Client, error) {
	this := new(Client)
	q, err := lentil.Dial(Config.QueueAddr)
	if err != nil {
		return nil, err
	}
	this.q = q
	this.verbose = verbose
	return this, nil
}

func (this *Client) put(cmd, mailto, workdir, out, tube string, priority, delay int) error {
	msg := make(map[string]string)
	msg["cmd"] = cmd
	msg["mailto"] = mailto
	msg["tube"] = tube
	msg["priority"] = fmt.Sprintf("%d", priority)
	if tube == "" {
		return errors.New("Missing required param -tube")
	}
	msg["delay"] = fmt.Sprintf("%d", delay) // Not used except for debugging
	workdir, e := filepath.Abs(workdir)
	if e != nil {
		return e
	}
	msg["workdir"] = workdir
	msg["out"] = out
	message, e := json.Marshal(msg)
	if this.verbose {
		log.Printf("QUEUEING UP: %s\n", message)
	}
	if e != nil {
		return e
	}
	if tube != "default" {
		e = this.q.Use(tube)
		if e != nil {
			return e
		}
	}
	_, e = this.q.Put(priority, delay, 60*60, message) // An hour TTR?
	return e
}

func (this *Client) putMany(input []byte) error {
	jobs := make([]map[string]string, 0)
	e := json.Unmarshal(input, &jobs)
	if e != nil {
		return e
	}
	for _, job := range jobs {
		priority := 0
		if priorityStr, exists := job["major"]; exists {
			priority, e = strconv.Atoi(priorityStr)
			if e != nil {
				return e
			}
		}
		delay := 0
		if delayStr, exists := job["delay"]; exists {
			delay, e = strconv.Atoi(delayStr)
			if e != nil {
				return e
			}
		}
		out, exists := job["out"]
		if !exists {
			out = "/dev/null"
		}
		workdir := "/tmp"
		dir, exists := job["workdir"]
		if exists {
			workdir = dir
		}
		e = this.put(job["cmd"], job["mailto"], workdir, out, job["tube"], priority, delay)
		if e != nil {
			return e
		}
	}
	return nil
}

func (this *Client) stats() error {
	q := NewJobQueue(this.q, false, make([]string, 0))
	stats, err := json.Marshal(q)
	if err != nil {
		return err
	}
	buffer := bytes.NewBufferString("")
	err = json.Indent(buffer, stats, "", "\t")
	if err != nil {
		return err
	}
	fmt.Printf("%s\n", buffer.String())
	return nil
}

func (this *Client) drain(tube string) error {
	_, err := this.q.Watch(tube)
	if err != nil {
		return err
	}
	_, err = this.q.Ignore("default")
	if err != nil {
		return err
	}
	for job, err := this.q.ReserveWithTimeout(0); err == nil; job, err = this.q.ReserveWithTimeout(0) {
		this.q.Delete(job.Id)
	}
	return err
}
<commit_msg>priority fix<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/nutrun/lentil"
	"log"
	"path/filepath"
	"strconv"
)

type Client struct {
	q       *lentil.Beanstalkd
	verbose bool
}

func NewClient(verbose bool) (*Client, error) {
	this := new(Client)
	q, err := lentil.Dial(Config.QueueAddr)
	if err != nil {
		return nil, err
	}
	this.q = q
	this.verbose = verbose
	return this, nil
}

func (this *Client) put(cmd, mailto, workdir, out, tube string, priority, delay int) error {
	msg := make(map[string]string)
	msg["cmd"] = cmd
	msg["mailto"] = mailto
	msg["tube"] = tube
	msg["pri"] = fmt.Sprintf("%d", priority) // Not used except for debugging
	if tube == "" {
		return errors.New("Missing required param -tube")
	}
	msg["delay"] = fmt.Sprintf("%d", delay) // Not used except for debugging
	workdir, e := filepath.Abs(workdir)
	if e != nil {
		return e
	}
	msg["workdir"] = workdir
	msg["out"] = out
	message, e := json.Marshal(msg)
	if this.verbose {
		log.Printf("QUEUEING UP: %s\n", message)
	}
	if e != nil {
		return e
	}
	if tube != "default" {
		e = this.q.Use(tube)
		if e != nil {
			return e
		}
	}
	_, e = this.q.Put(priority, delay, 60*60, message) // An hour TTR?
	return e
}

func (this *Client) putMany(input []byte) error {
	jobs := make([]map[string]string, 0)
	e := json.Unmarshal(input, &jobs)
	if e != nil {
		return e
	}
	for _, job := range jobs {
		priority := 0
		if priorityStr, exists := job["pri"]; exists {
			priority, e = strconv.Atoi(priorityStr)
			if e != nil {
				return e
			}
		}
		delay := 0
		if delayStr, exists := job["delay"]; exists {
			delay, e = strconv.Atoi(delayStr)
			if e != nil {
				return e
			}
		}
		out, exists := job["out"]
		if !exists {
			out = "/dev/null"
		}
		workdir := "/tmp"
		dir, exists := job["workdir"]
		if exists {
			workdir = dir
		}
		e = this.put(job["cmd"], job["mailto"], workdir, out, job["tube"], priority, delay)
		if e != nil {
			return e
		}
	}
	return nil
}

func (this *Client) stats() error {
	q := NewJobQueue(this.q, false, make([]string, 0))
	stats, err := json.Marshal(q)
	if err != nil {
		return err
	}
	buffer := bytes.NewBufferString("")
	err = json.Indent(buffer, stats, "", "\t")
	if err != nil {
		return err
	}
	fmt.Printf("%s\n", buffer.String())
	return nil
}

func (this *Client) drain(tube string) error {
	_, err := this.q.Watch(tube)
	if err != nil {
		return err
	}
	_, err = this.q.Ignore("default")
	if err != nil {
		return err
	}
	for job, err := this.q.ReserveWithTimeout(0); err == nil; job, err = this.q.ReserveWithTimeout(0) {
		this.q.Delete(job.Id)
	}
	return err
}
<|endoftext|>


<commit_before>package dns

// A concurrent client implementation.
// Client sends query to a channel which
// will then handle the query. Returned replys
// are return on another channel. Ready for handling --- same
// setup for server - a HANDLER function that gets run
// when the query returns.

import (
	"os"
	"io"
	"net"
)

type QueryHandler interface {
	QueryDNS(w RequestWriter, q *Msg)
}

// A RequestWriter interface is used by an DNS query handler to
// construct an DNS request.
type RequestWriter interface {
	WriteMessages([]*Msg)
	Write(*Msg)
	Send(*Msg) os.Error
	Receive() (*Msg, os.Error)
}

// hijacked connections...?
type reply struct {
	client         *Client
	addr           string
	req            *Msg
	conn           net.Conn
	tsigRequestMAC string
	tsigTimersOnly bool
}

type Request struct {
	Request *Msg
	Addr    string
	Client  *Client
}

// QueryMux is an DNS request multiplexer. It matches the
// zone name of each incoming request against a list of
// registered patterns add calls the handler for the pattern
// that most closely matches the zone name.
type QueryMux struct {
	m map[string]QueryHandler
}

// NewQueryMux allocates and returns a new QueryMux.
func NewQueryMux() *QueryMux { return &QueryMux{make(map[string]QueryHandler)} }

// DefaultQueryMux is the default QueryMux used by Query.
var DefaultQueryMux = NewQueryMux()

func newQueryChanSlice() chan []*Msg { return make(chan []*Msg) }
func newQueryChan() chan *Request    { return make(chan *Request) }

// Default channel to use for the resolver
var DefaultReplyChan = newQueryChanSlice()
var DefaultQueryChan = newQueryChan()

// The HandlerQueryFunc type is an adapter to allow the use of
// ordinary functions as DNS query handlers.  If f is a function
// with the appropriate signature, HandlerQueryFunc(f) is a
// QueryHandler object that calls f.
type HandlerQueryFunc func(RequestWriter, *Msg)

// QueryDNS calls f(w, reg)
func (f HandlerQueryFunc) QueryDNS(w RequestWriter, r *Msg) {
	go f(w, r)
}

func HandleQueryFunc(pattern string, handler func(RequestWriter, *Msg)) {
	DefaultQueryMux.HandleQueryFunc(pattern, handler)
}

// reusing zoneMatch from server.go
func (mux *QueryMux) match(zone string) QueryHandler {
	var h QueryHandler
	var n = 0
	for k, v := range mux.m {
		if !zoneMatch(k, zone) {
			continue
		}
		if h == nil || len(k) > n {
			n = len(k)
			h = v
		}
	}
	return h
}

func (mux *QueryMux) Handle(pattern string, handler QueryHandler) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	mux.m[pattern] = handler
}

func (mux *QueryMux) HandleQueryFunc(pattern string, handler func(RequestWriter, *Msg)) {
	mux.Handle(pattern, HandlerQueryFunc(handler))
}

func (mux *QueryMux) QueryDNS(w RequestWriter, r *Msg) {
	h := mux.match(r.Question[0].Name)
	if h == nil {
                panic("dns: no handler found for " + r.Question[0].Name)
	}
	h.QueryDNS(w, r)
}

// TODO add: LocalAddr
type Client struct {
	Net          string            // if "tcp" a TCP query will be initiated, otherwise an UDP one
	Addr         string            // address to call
	Attempts     int               // number of attempts
	Retry        bool              // retry with TCP
	ChannelQuery chan *Request     // read DNS request from this channel
	ChannelReply chan []*Msg       // read DNS request from this channel
	ReadTimeout  int64             // the net.Conn.SetReadTimeout value for new connections
	WriteTimeout int64             // the net.Conn.SetWriteTimeout value for new connections
	TsigSecret   map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>
}

func NewClient() *Client {
	c := new(Client)
	c.Net = "udp"
	c.Attempts = 1
	c.ChannelReply = DefaultReplyChan
	return c
}

type Query struct {
	ChannelQuery chan *Request // read DNS request from this channel
	Handler      QueryHandler  // handler to invoke, dns.DefaultQueryMux if nil
}

func (q *Query) Query() os.Error {
	handler := q.Handler
	if handler == nil {
		handler = DefaultQueryMux
	}
//forever:
	for {
		select {
		case in := <-q.ChannelQuery:
			w := new(reply)
			w.req = in.Request
			w.addr = in.Addr
			w.client = in.Client
			handler.QueryDNS(w, in.Request)
		}
	}
	return nil
}

func (q *Query) ListenAndQuery() os.Error {
	if q.ChannelQuery == nil {
		q.ChannelQuery = DefaultQueryChan
	}
	return q.Query()
}

// Start listener for firing off the queries. If
// c is nil DefaultQueryChan is used. If handler is nil
// DefaultQueryMux is used.
func ListenAndQuery(c chan *Request, handler QueryHandler) {
	q := &Query{ChannelQuery: c, Handler: handler}
	go q.ListenAndQuery()
}

func (w *reply) Write(m *Msg) {
	w.Client().ChannelReply <- []*Msg{w.req, m}
}

// Do performs an asynchronous query. The result is returned on the
// channel set in the c.
func (c *Client) Do(m *Msg, a string) {
	if c.ChannelQuery == nil {
		DefaultQueryChan <- &Request{Client: c, Addr: a, Request: m}
	} else {
		c.ChannelQuery <- &Request{Client: c, Addr: a, Request: m}
	}
}

// Perform an synchronous query. Send the message m to the address
// contained in a
func (c *Client) Exchange(m *Msg, a string) *Msg {
	w := new(reply)
	w.client = c
	w.addr = a
	out, ok := m.Pack()
	if !ok {
		//
	}
	_, err := w.writeClient(out)
	if err != nil {
		return nil
	}

	// udp / tcp
	p := make([]byte, DefaultMsgSize)
	n, err := w.readClient(p)
	if err != nil {
		return nil
	}
	p = p[:n]
        r := new(Msg)
	if ok := r.Unpack(p); !ok {
		return nil
	}
	return r
}

func (w *reply) WriteMessages(m []*Msg) {
	m1 := append([]*Msg{w.req}, m...)
	w.Client().ChannelReply <- m1
}

func (w *reply) Client() *Client {
	return w.client
}

func (w *reply) Request() *Msg {
	return w.req
}

func (w *reply) Receive() (*Msg, os.Error) {
	var p []byte
	m := new(Msg)
	switch w.Client().Net {
	case "tcp":
		p = make([]byte, MaxMsgSize)
	case "udp":
		p = make([]byte, DefaultMsgSize)
	}
	n, err := w.readClient(p)
	if err != nil {
		return nil, err
	}
	p = p[:n]
	if ok := m.Unpack(p); !ok {
		return nil, ErrUnpack
	}
	// Tsig
	if m.IsTsig() {
		secret := m.Extra[len(m.Extra)-1].(*RR_TSIG).Hdr.Name
		_, ok := w.Client().TsigSecret[secret]
		if !ok {
			return m, ErrNoSig
		}
		ok, err := TsigVerify(p, w.Client().TsigSecret[secret], w.tsigRequestMAC, w.tsigTimersOnly)
		if !ok {
			return m, err
		}
	}
	return m, nil
}

func (w *reply) readClient(p []byte) (n int, err os.Error) {
	if w.conn == nil {
		panic("no connection")
	}
	switch w.Client().Net {
	case "tcp":
		if len(p) < 1 {
			return 0, io.ErrShortBuffer
		}
		n, err = w.conn.(*net.TCPConn).Read(p[0:2])
		if err != nil || n != 2 {
			return n, err
		}
		l, _ := unpackUint16(p[0:2], 0)
		if l == 0 {
			return 0, ErrShortRead
		}
		if int(l) > len(p) {
			return int(l), io.ErrShortBuffer
		}
		n, err = w.conn.(*net.TCPConn).Read(p[:l])
		if err != nil {
			return n, err
		}
		i := n
		for i < int(l) {
			j, err := w.conn.(*net.TCPConn).Read(p[i:int(l)])
			if err != nil {
				return i, err
			}
			i += j
		}
		n = i
	case "udp":
		n, _, err = w.conn.(*net.UDPConn).ReadFromUDP(p)
		if err != nil {
			return n, err
		}
	}
	return
}

// Send a msg to the address specified in w.
// If the message m contains a TSIG record the transaction
// signature is calculated.
func (w *reply) Send(m *Msg) os.Error {
	if m.IsTsig() {
		secret := m.Extra[len(m.Extra)-1].(*RR_TSIG).Hdr.Name
		_, ok := w.Client().TsigSecret[secret]
		if !ok {
			return ErrNoSig
		}
		m, _ = TsigGenerate(m, w.Client().TsigSecret[secret], w.tsigRequestMAC, w.tsigTimersOnly)
		w.tsigRequestMAC = m.Extra[len(m.Extra)-1].(*RR_TSIG).MAC // Safe the requestMAC
	}
	out, ok := m.Pack()
	if !ok {
		return ErrPack
	}
	_, err := w.writeClient(out)
	if err != nil {
		return err
	}
	return nil
}

func (w *reply) writeClient(p []byte) (n int, err os.Error) {
	c := w.Client()
	if c.Attempts == 0 {
		panic("c.Attempts 0")
	}
	if c.Net == "" {
		panic("c.Net empty")
	}

	conn, err := net.Dial(c.Net, w.addr)
	if err != nil {
		return 0, err
	}
	w.conn = conn
	switch c.Net {
	case "tcp":
		if len(p) < 2 {
			return 0, io.ErrShortBuffer
		}
		for a := 0; a < c.Attempts; a++ {
			l := make([]byte, 2)
			l[0], l[1] = packUint16(uint16(len(p)))
			n, err = conn.Write(l)
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
			if n != 2 {
				return n, io.ErrShortWrite
			}
			n, err = conn.Write(p)
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
			i := n
			if i < len(p) {
				j, err := conn.Write(p[i:len(p)])
				if err != nil {
					if e, ok := err.(net.Error); ok && e.Timeout() {
						// We are half way in our write...
						continue
					}
					return i, err
				}
				i += j
			}
			n = i
		}
	case "udp":
		for a := 0; a < c.Attempts; a++ {
			n, err = conn.(*net.UDPConn).WriteTo(p, conn.RemoteAddr())
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return 0, err
			}
		}
	}
	return 0, nil
}
<commit_msg>support tcp4, tcp6/udp6<commit_after>package dns

// A concurrent client implementation.
// Client sends query to a channel which
// will then handle the query. Returned replys
// are return on another channel. Ready for handling --- same
// setup for server - a HANDLER function that gets run
// when the query returns.

import (
	"os"
	"io"
	"net"
)

type QueryHandler interface {
	QueryDNS(w RequestWriter, q *Msg)
}

// A RequestWriter interface is used by an DNS query handler to
// construct an DNS request.
type RequestWriter interface {
	WriteMessages([]*Msg)
	Write(*Msg)
	Send(*Msg) os.Error
	Receive() (*Msg, os.Error)
}

// hijacked connections...?
type reply struct {
	client         *Client
	addr           string
	req            *Msg
	conn           net.Conn
	tsigRequestMAC string
	tsigTimersOnly bool
}

type Request struct {
	Request *Msg
	Addr    string
	Client  *Client
}

// QueryMux is an DNS request multiplexer. It matches the
// zone name of each incoming request against a list of
// registered patterns add calls the handler for the pattern
// that most closely matches the zone name.
type QueryMux struct {
	m map[string]QueryHandler
}

// NewQueryMux allocates and returns a new QueryMux.
func NewQueryMux() *QueryMux { return &QueryMux{make(map[string]QueryHandler)} }

// DefaultQueryMux is the default QueryMux used by Query.
var DefaultQueryMux = NewQueryMux()

func newQueryChanSlice() chan []*Msg { return make(chan []*Msg) }
func newQueryChan() chan *Request    { return make(chan *Request) }

// Default channel to use for the resolver
var DefaultReplyChan = newQueryChanSlice()
var DefaultQueryChan = newQueryChan()

// The HandlerQueryFunc type is an adapter to allow the use of
// ordinary functions as DNS query handlers.  If f is a function
// with the appropriate signature, HandlerQueryFunc(f) is a
// QueryHandler object that calls f.
type HandlerQueryFunc func(RequestWriter, *Msg)

// QueryDNS calls f(w, reg)
func (f HandlerQueryFunc) QueryDNS(w RequestWriter, r *Msg) {
	go f(w, r)
}

func HandleQueryFunc(pattern string, handler func(RequestWriter, *Msg)) {
	DefaultQueryMux.HandleQueryFunc(pattern, handler)
}

// reusing zoneMatch from server.go
func (mux *QueryMux) match(zone string) QueryHandler {
	var h QueryHandler
	var n = 0
	for k, v := range mux.m {
		if !zoneMatch(k, zone) {
			continue
		}
		if h == nil || len(k) > n {
			n = len(k)
			h = v
		}
	}
	return h
}

func (mux *QueryMux) Handle(pattern string, handler QueryHandler) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	mux.m[pattern] = handler
}

func (mux *QueryMux) HandleQueryFunc(pattern string, handler func(RequestWriter, *Msg)) {
	mux.Handle(pattern, HandlerQueryFunc(handler))
}

func (mux *QueryMux) QueryDNS(w RequestWriter, r *Msg) {
	h := mux.match(r.Question[0].Name)
	if h == nil {
                panic("dns: no handler found for " + r.Question[0].Name)
	}
	h.QueryDNS(w, r)
}

// TODO add: LocalAddr
type Client struct {
	Net          string            // if "tcp" a TCP query will be initiated, otherwise an UDP one
	Addr         string            // address to call
	Attempts     int               // number of attempts
	Retry        bool              // retry with TCP
	ChannelQuery chan *Request     // read DNS request from this channel
	ChannelReply chan []*Msg       // read DNS request from this channel
	ReadTimeout  int64             // the net.Conn.SetReadTimeout value for new connections
	WriteTimeout int64             // the net.Conn.SetWriteTimeout value for new connections
	TsigSecret   map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>
}

func NewClient() *Client {
	c := new(Client)
	c.Net = "udp"
	c.Attempts = 1
	c.ChannelReply = DefaultReplyChan
	return c
}

type Query struct {
	ChannelQuery chan *Request // read DNS request from this channel
	Handler      QueryHandler  // handler to invoke, dns.DefaultQueryMux if nil
}

func (q *Query) Query() os.Error {
	handler := q.Handler
	if handler == nil {
		handler = DefaultQueryMux
	}
//forever:
	for {
		select {
		case in := <-q.ChannelQuery:
			w := new(reply)
			w.req = in.Request
			w.addr = in.Addr
			w.client = in.Client
			handler.QueryDNS(w, in.Request)
		}
	}
	return nil
}

func (q *Query) ListenAndQuery() os.Error {
	if q.ChannelQuery == nil {
		q.ChannelQuery = DefaultQueryChan
	}
	return q.Query()
}

// Start listener for firing off the queries. If
// c is nil DefaultQueryChan is used. If handler is nil
// DefaultQueryMux is used.
func ListenAndQuery(c chan *Request, handler QueryHandler) {
	q := &Query{ChannelQuery: c, Handler: handler}
	go q.ListenAndQuery()
}

func (w *reply) Write(m *Msg) {
	w.Client().ChannelReply <- []*Msg{w.req, m}
}

// Do performs an asynchronous query. The result is returned on the
// channel set in the c.
func (c *Client) Do(m *Msg, a string) {
	if c.ChannelQuery == nil {
		DefaultQueryChan <- &Request{Client: c, Addr: a, Request: m}
	} else {
		c.ChannelQuery <- &Request{Client: c, Addr: a, Request: m}
	}
}

// Perform an synchronous query. Send the message m to the address
// contained in a
func (c *Client) Exchange(m *Msg, a string) *Msg {
	w := new(reply)
	w.client = c
	w.addr = a
	out, ok := m.Pack()
	if !ok {
		//
	}
	_, err := w.writeClient(out)
	if err != nil {
		return nil
	}

	// udp / tcp
	p := make([]byte, DefaultMsgSize)
	n, err := w.readClient(p)
	if err != nil {
		return nil
	}
	p = p[:n]
        r := new(Msg)
	if ok := r.Unpack(p); !ok {
		return nil
	}
	return r
}

func (w *reply) WriteMessages(m []*Msg) {
	m1 := append([]*Msg{w.req}, m...)
	w.Client().ChannelReply <- m1
}

func (w *reply) Client() *Client {
	return w.client
}

func (w *reply) Request() *Msg {
	return w.req
}

func (w *reply) Receive() (*Msg, os.Error) {
	var p []byte
	m := new(Msg)
	switch w.Client().Net {
	case "tcp", "tcp4", "tcp6":
		p = make([]byte, MaxMsgSize)
	case "udp", "udp4", "udp6":
		p = make([]byte, DefaultMsgSize)
	}
	n, err := w.readClient(p)
	if err != nil {
		return nil, err
	}
	p = p[:n]
	if ok := m.Unpack(p); !ok {
		return nil, ErrUnpack
	}
	// Tsig
	if m.IsTsig() {
		secret := m.Extra[len(m.Extra)-1].(*RR_TSIG).Hdr.Name
		_, ok := w.Client().TsigSecret[secret]
		if !ok {
			return m, ErrNoSig
		}
		ok, err := TsigVerify(p, w.Client().TsigSecret[secret], w.tsigRequestMAC, w.tsigTimersOnly)
		if !ok {
			return m, err
		}
	}
	return m, nil
}

func (w *reply) readClient(p []byte) (n int, err os.Error) {
	if w.conn == nil {
		panic("no connection")
	}
	switch w.Client().Net {
	case "tcp", "tcp4", "tcp6":
		if len(p) < 1 {
			return 0, io.ErrShortBuffer
		}
		n, err = w.conn.(*net.TCPConn).Read(p[0:2])
		if err != nil || n != 2 {
			return n, err
		}
		l, _ := unpackUint16(p[0:2], 0)
		if l == 0 {
			return 0, ErrShortRead
		}
		if int(l) > len(p) {
			return int(l), io.ErrShortBuffer
		}
		n, err = w.conn.(*net.TCPConn).Read(p[:l])
		if err != nil {
			return n, err
		}
		i := n
		for i < int(l) {
			j, err := w.conn.(*net.TCPConn).Read(p[i:int(l)])
			if err != nil {
				return i, err
			}
			i += j
		}
		n = i
	case "udp", "udp4", "udp6":
		n, _, err = w.conn.(*net.UDPConn).ReadFromUDP(p)
		if err != nil {
			return n, err
		}
	}
	return
}

// Send a msg to the address specified in w.
// If the message m contains a TSIG record the transaction
// signature is calculated.
func (w *reply) Send(m *Msg) os.Error {
	if m.IsTsig() {
		secret := m.Extra[len(m.Extra)-1].(*RR_TSIG).Hdr.Name
		_, ok := w.Client().TsigSecret[secret]
		if !ok {
			return ErrNoSig
		}
		m, _ = TsigGenerate(m, w.Client().TsigSecret[secret], w.tsigRequestMAC, w.tsigTimersOnly)
		w.tsigRequestMAC = m.Extra[len(m.Extra)-1].(*RR_TSIG).MAC // Safe the requestMAC
	}
	out, ok := m.Pack()
	if !ok {
		return ErrPack
	}
	_, err := w.writeClient(out)
	if err != nil {
		return err
	}
	return nil
}

func (w *reply) writeClient(p []byte) (n int, err os.Error) {
	c := w.Client()
	if c.Attempts == 0 {
		panic("c.Attempts 0")
	}
	if c.Net == "" {
		panic("c.Net empty")
	}

	conn, err := net.Dial(c.Net, w.addr)
	if err != nil {
		return 0, err
	}
	w.conn = conn
	switch c.Net {
	case "tcp", "tcp4", "tcp6":
		if len(p) < 2 {
			return 0, io.ErrShortBuffer
		}
		for a := 0; a < c.Attempts; a++ {
			l := make([]byte, 2)
			l[0], l[1] = packUint16(uint16(len(p)))
			n, err = conn.Write(l)
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
			if n != 2 {
				return n, io.ErrShortWrite
			}
			n, err = conn.Write(p)
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
			i := n
			if i < len(p) {
				j, err := conn.Write(p[i:len(p)])
				if err != nil {
					if e, ok := err.(net.Error); ok && e.Timeout() {
						// We are half way in our write...
						continue
					}
					return i, err
				}
				i += j
			}
			n = i
		}
	case "udp", "udp4", "udp6":
		for a := 0; a < c.Attempts; a++ {
			n, err = conn.(*net.UDPConn).WriteTo(p, conn.RemoteAddr())
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return 0, err
			}
		}
	}
	return 0, nil
}
<|endoftext|>


<commit_before>package twitter

import (
	"context"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/garyburd/go-oauth/oauth"
)

// Client represents the client used to make requests to the Twitter API 1.1.
type Client struct {
	httpClient  HTTPClient
	oauthClient *oauth.Client
	accessCreds *oauth.Credentials

	gzipDisabled bool
}

// NewClient returns a new Client instance using the provided application
// credentials, optional default consumer credentials, and optional HTTPClient.
func NewClient(consumerCreds ConsumerCredentials, accessCreds AccessCredentials, httpClient HTTPClient) *Client {
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	oauthClient := oauth.Client{
		Credentials: oauth.Credentials{
			Token:  consumerCreds.Key,
			Secret: consumerCreds.Secret,
		},
	}
	return &Client{
		httpClient:  httpClient,
		oauthClient: &oauthClient,
		accessCreds: &oauth.Credentials{
			Token:  accessCreds.Token,
			Secret: accessCreds.Secret,
		},
	}
}

// WithGzipDisabled returns a new shallow copy of the Client with Gzip disabled.
func (c *Client) WithGzipDisabled() *Client {
	newC := *c
	newC.gzipDisabled = true
	return &newC
}

// do implements the OAuthClient interface. It is used to make an OAuth HTTP
// request with the provided context, HTTP method, OAuth credentials, URL
// string, and URL query parameters. It returns the corresponding HTTP response
// or error.
func (c *Client) do(ctx context.Context, method string, urlStr string, values url.Values) (*http.Response, error) {
	// Set credentials.
	accessCreds := c.accessCredentials(ctx)
	// Create HTTP request.
	req, err := http.NewRequest(method, urlStr, strings.NewReader(values.Encode()))
	if err != nil {
		return nil, err
	}
	req = req.WithContext(ctx)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	if !c.gzipDisabled {
		req.Header.Set("Accept-Encoding", "gzip")
	}
	err = c.oauthClient.SetAuthorizationHeader(req.Header, accessCreds, method, req.URL, values)
	if err != nil {
		return nil, err
	}
	// Make request.
	resp, err := c.httpClient.Do(req)
	if c.gzipDisabled || err != nil || !isGzipped(resp.Header) {
		return resp, err
	}
	return gzipResponse(resp)
}

// HTTPClient is the interface for making HTTP requests. It accepts an HTTP
// request and returns the corresponding HTTP response or error.
type HTTPClient interface {
	Do(*http.Request) (*http.Response, error)
}

// ConsumerCredentials represents a token/secret pair of auth credentials.
type ConsumerCredentials struct {
	Key    string
	Secret string
}

// oauthClient is the interface for making an OAuth HTTP request using the
// provided context, HTTP method, optional consumer credentials, URL string,
// and URL query parameters. It returns the corresponding HTTP response or
// error.
type oauthClient interface {
	do(ctx context.Context, method string, urlStr string, values url.Values) (*http.Response, error)
}

// RateLimit represents the Twitter rate limit for the credentials used in the
// associated request.
type RateLimit struct {
	Limit     int
	Remaining int
	Reset     int
}

func getRateLimit(h http.Header) RateLimit {
	limit, _ := strconv.Atoi(h.Get("X-Rate-Limit-Limit"))
	remaining, _ := strconv.Atoi(h.Get("X-Rate-Limit-Remaining"))
	reset, _ := strconv.Atoi(h.Get("X-Rate-Limit-Reset"))
	return RateLimit{
		Limit:     limit,
		Remaining: remaining,
		Reset:     reset,
	}
}

// AccessCredentials represents an OAuth token and secret pair.
type AccessCredentials struct {
	Token  string
	Secret string
}

type accCredsKeyType int

const accCredsKey = 0

func (c *Client) accessCredentials(ctx context.Context) *oauth.Credentials {
	if val := ctx.Value(accCredsKey); val != nil {
		return val.(*oauth.Credentials)
	}
	return c.accessCreds
}

// WithAccessCredentials returns a new context with the provided parent context
// and AccessCredentials.
func WithAccessCredentials(parent context.Context, creds AccessCredentials) context.Context {
	return context.WithValue(parent, accCredsKey, &oauth.Credentials{
		Token:  creds.Token,
		Secret: creds.Secret,
	})
}
<commit_msg>Use private type as key in context value<commit_after>package twitter

import (
	"context"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/garyburd/go-oauth/oauth"
)

// Client represents the client used to make requests to the Twitter API 1.1.
type Client struct {
	httpClient  HTTPClient
	oauthClient *oauth.Client
	accessCreds *oauth.Credentials

	gzipDisabled bool
}

// NewClient returns a new Client instance using the provided application
// credentials, optional default consumer credentials, and optional HTTPClient.
func NewClient(consumerCreds ConsumerCredentials, accessCreds AccessCredentials, httpClient HTTPClient) *Client {
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	oauthClient := oauth.Client{
		Credentials: oauth.Credentials{
			Token:  consumerCreds.Key,
			Secret: consumerCreds.Secret,
		},
	}
	return &Client{
		httpClient:  httpClient,
		oauthClient: &oauthClient,
		accessCreds: &oauth.Credentials{
			Token:  accessCreds.Token,
			Secret: accessCreds.Secret,
		},
	}
}

// WithGzipDisabled returns a new shallow copy of the Client with Gzip disabled.
func (c *Client) WithGzipDisabled() *Client {
	newC := *c
	newC.gzipDisabled = true
	return &newC
}

// do implements the OAuthClient interface. It is used to make an OAuth HTTP
// request with the provided context, HTTP method, OAuth credentials, URL
// string, and URL query parameters. It returns the corresponding HTTP response
// or error.
func (c *Client) do(ctx context.Context, method string, urlStr string, values url.Values) (*http.Response, error) {
	// Set credentials.
	accessCreds := c.accessCredentials(ctx)
	// Create HTTP request.
	req, err := http.NewRequest(method, urlStr, strings.NewReader(values.Encode()))
	if err != nil {
		return nil, err
	}
	req = req.WithContext(ctx)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	if !c.gzipDisabled {
		req.Header.Set("Accept-Encoding", "gzip")
	}
	err = c.oauthClient.SetAuthorizationHeader(req.Header, accessCreds, method, req.URL, values)
	if err != nil {
		return nil, err
	}
	// Make request.
	resp, err := c.httpClient.Do(req)
	if c.gzipDisabled || err != nil || !isGzipped(resp.Header) {
		return resp, err
	}
	return gzipResponse(resp)
}

// HTTPClient is the interface for making HTTP requests. It accepts an HTTP
// request and returns the corresponding HTTP response or error.
type HTTPClient interface {
	Do(*http.Request) (*http.Response, error)
}

// ConsumerCredentials represents a token/secret pair of auth credentials.
type ConsumerCredentials struct {
	Key    string
	Secret string
}

// oauthClient is the interface for making an OAuth HTTP request using the
// provided context, HTTP method, optional consumer credentials, URL string,
// and URL query parameters. It returns the corresponding HTTP response or
// error.
type oauthClient interface {
	do(ctx context.Context, method string, urlStr string, values url.Values) (*http.Response, error)
}

// RateLimit represents the Twitter rate limit for the credentials used in the
// associated request.
type RateLimit struct {
	Limit     int
	Remaining int
	Reset     int
}

func getRateLimit(h http.Header) RateLimit {
	limit, _ := strconv.Atoi(h.Get("X-Rate-Limit-Limit"))
	remaining, _ := strconv.Atoi(h.Get("X-Rate-Limit-Remaining"))
	reset, _ := strconv.Atoi(h.Get("X-Rate-Limit-Reset"))
	return RateLimit{
		Limit:     limit,
		Remaining: remaining,
		Reset:     reset,
	}
}

// AccessCredentials represents an OAuth token and secret pair.
type AccessCredentials struct {
	Token  string
	Secret string
}

type accCredsKeyType int

const accCredsKey accCredsKeyType = 0

func (c *Client) accessCredentials(ctx context.Context) *oauth.Credentials {
	if val := ctx.Value(accCredsKey); val != nil {
		return val.(*oauth.Credentials)
	}
	return c.accessCreds
}

// WithAccessCredentials returns a new context with the provided parent context
// and AccessCredentials.
func WithAccessCredentials(parent context.Context, creds AccessCredentials) context.Context {
	return context.WithValue(parent, accCredsKey, &oauth.Credentials{
		Token:  creds.Token,
		Secret: creds.Secret,
	})
}
<|endoftext|>


<commit_before>/*
A smart client for go.

Usage:

 client, err := couchbase.Connect("http://myserver:8091/")
 handleError(err)
 pool, err := client.GetPool("default")
 handleError(err)
 bucket, err := pool.GetBucket("MyAwesomeBucket")
 handleError(err)
 ...

or a shortcut for the bucket directly

 bucket, err := couchbase.GetBucket("http://myserver:8091/", "default", "default")
*/
package couchbase

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"sync/atomic"

	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

// Execute a function on a memcached connection to the node owning key "k"
//
// Note that this automatically handles transient errors by replaying
// your function on a "not-my-vbucket" error, so don't assume
// your command will only be executed only once.
func (b *Bucket) Do(k string, f func(mc *memcached.Client, vb uint16) error) error {
	vb := b.VBHash(k)
	for {
		masterId := b.VBucketServerMap.VBucketMap[vb][0]
		conn, err := b.connections[masterId].Get()
		defer b.connections[masterId].Return(conn)
		if err != nil {
			return err
		}

		err = f(conn, uint16(vb))
		switch err.(type) {
		default:
			return err
		case gomemcached.MCResponse:
			st := err.(gomemcached.MCResponse).Status
			atomic.AddUint64(&b.pool.client.Statuses[st], 1)
			if st == gomemcached.NOT_MY_VBUCKET {
				b.refresh()
			} else {
				return err
			}
		}
	}
}

type gathered_stats struct {
	sn   string
	vals map[string]string
}

func getStatsParallel(b *Bucket, offset int, which string, ch chan<- gathered_stats) {
	sn := b.VBucketServerMap.ServerList[offset]

	results := map[string]string{}
	conn, err := b.connections[offset].Get()
	defer b.connections[offset].Return(conn)
	if err != nil {
		ch <- gathered_stats{sn, results}
	} else {
		st, err := conn.StatsMap(which)
		if err == nil {
			ch <- gathered_stats{sn, st}
		} else {
			ch <- gathered_stats{sn, results}
		}
	}
}

// Get a set of stats from all servers.
//
// Returns a map of server ID -> map of stat key to map value.
func (b *Bucket) GetStats(which string) map[string]map[string]string {
	rv := map[string]map[string]string{}

	if b.VBucketServerMap.ServerList == nil {
		return rv
	}
	// Go grab all the things at once.
	todo := len(b.VBucketServerMap.ServerList)
	ch := make(chan gathered_stats, todo)

	for offset := range b.VBucketServerMap.ServerList {
		go getStatsParallel(b, offset, which, ch)
	}

	// Gather the results
	for i := 0; i < len(b.VBucketServerMap.ServerList); i++ {
		g := <-ch
		if len(g.vals) > 0 {
			rv[g.sn] = g.vals
		}
	}

	return rv
}

func (b *Bucket) doBulkGet(vb uint16, keys []string,
	ch chan<- map[string]*gomemcached.MCResponse) {

	masterId := b.VBucketServerMap.VBucketMap[vb][0]
	conn, err := b.connections[masterId].Get()
	if err != nil {
		ch <- map[string]*gomemcached.MCResponse{}
		return
	}
	defer b.connections[masterId].Return(conn)

	m, err := conn.GetBulk(vb, keys)
	switch err.(type) {
	default:
		ch <- m
	case *gomemcached.MCResponse:
		fmt.Printf("Got a memcached error")
		st := err.(*gomemcached.MCResponse).Status
		atomic.AddUint64(&b.pool.client.Statuses[st], 1)
		if st == gomemcached.NOT_MY_VBUCKET {
			b.refresh()
		}
		ch <- map[string]*gomemcached.MCResponse{}
	}
}

func (b *Bucket) processBulkGet(kdm map[uint16][]string,
	ch chan map[string]*gomemcached.MCResponse) {

	wch := make(chan uint16)

	worker := func() {
		for k := range wch {
			b.doBulkGet(k, kdm[k], ch)
		}
	}

	for i := 0; i < 4; i++ {
		go worker()
	}

	for k := range kdm {
		wch <- k
	}
	close(wch)

}
func (b *Bucket) GetBulk(keys []string) map[string]*gomemcached.MCResponse {
	// Organize by vbucket
	kdm := map[uint16][]string{}
	for _, k := range keys {
		vb := uint16(b.VBHash(k))
		a, ok := kdm[vb]
		if !ok {
			a = []string{}
		}
		kdm[vb] = append(a, k)
	}

	ch := make(chan map[string]*gomemcached.MCResponse)
	defer close(ch)

	go b.processBulkGet(kdm, ch)

	rv := map[string]*gomemcached.MCResponse{}
	for _ = range kdm {
		m := <-ch
		for k, v := range m {
			rv[k] = v
		}
	}

	return rv
}

// Set a value in this bucket.
// The value will be serialized into a JSON document.
func (b *Bucket) Set(k string, exp int, v interface{}) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return b.SetRaw(k, exp, data)
}

// Set a value in this bucket.
// The value will be stored as raw bytes.
func (b *Bucket) SetRaw(k string, exp int, v []byte) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		_, err := mc.Set(vb, k, 0, exp, v)
		if err != nil {
			return err
		}
		return nil
	})
}

// Adds a value to this bucket; like Set except that nothing happens if the key exists.
// The value will be serialized into a JSON document.
func (b *Bucket) Add(k string, exp int, v interface{}) (added bool, err error) {
	data, err := json.Marshal(v)
	if err != nil {
		return false, err
	}
	return b.AddRaw(k, exp, data)
}

// Adds a value to this bucket; like SetRaw except that nothing happens if the key exists.
// The value will be stored as raw bytes.
func (b *Bucket) AddRaw(k string, exp int, v []byte) (added bool, err error) {
	err = b.Do(k, func(mc *memcached.Client, vb uint16) error {
		switch res, err := memcached.UnwrapMemcachedError(mc.Add(vb, k, 0, exp, v)); {
		case err != nil:
			return err
		case res.Status == gomemcached.SUCCESS:
			added = true
		case res.Status != gomemcached.KEY_EEXISTS:
			return res
		}
		return nil
	})
	return
}

// Get a raw value from this bucket, including its CAS counter.
func (b *Bucket) GetsRaw(k string, cas *uint64) ([]byte, error) {
	var data []byte
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Get(vb, k)
		if err != nil {
			return err
		}
		if cas != nil {
			*cas = res.Cas
		}
		data = res.Body
		return nil
	})
	return data, err
}

// Get a value from this bucket, including its CAS counter.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Gets(k string, rv interface{}, cas *uint64) error {
	data, err := b.GetsRaw(k, cas)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, rv)
}

// Get a value from this bucket.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Get(k string, rv interface{}) error {
	return b.Gets(k, rv, nil)
}

// Get a raw value from this bucket.
func (b *Bucket) GetRaw(k string) ([]byte, error) {
	return b.GetsRaw(k, nil)
}

// Delete a key from this bucket.
func (b *Bucket) Delete(k string) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		_, err := mc.Del(vb, k)
		if err != nil {
			return err
		}
		return nil
	})
}

// Increment a key
func (b *Bucket) Incr(k string, amt, def uint64, exp int) (uint64, error) {
	var rv uint64
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Incr(vb, k, amt, def, exp)
		if err != nil {
			return err
		}
		rv = res
		return nil
	})
	return rv, err
}

// A callback function to update a document
type UpdateFunc func(current []byte) (updated []byte, err error)

// Return this as the error from an UpdateFunc to cancel the Update operation.
const UpdateCancel = memcached.CASQuit

// Safe update of a document, avoiding conflicts by using CAS.
// The callback function will be invoked with the current raw document contents (or nil if the
// document doesn't exist); it should return the updated raw contents (or nil to delete.)
// If it decides not to change anything it can return UpdateCancel as the error.
//
// If another writer modifies the document between the get and the set, the callback will be
// invoked again with the newer value.
func (b *Bucket) Update(k string, exp int, callback UpdateFunc) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		var callbackErr error
		casFunc := func(current []byte) ([]byte, memcached.CasOp) {
			var updated []byte
			updated, callbackErr = callback(current)
			if callbackErr != nil {
				return nil, memcached.CASQuit
			} else if updated == nil {
				return updated, memcached.CASDelete
			} else {
				return updated, memcached.CASStore
			}
		}
		_, err := mc.CAS(vb, k, casFunc, exp)
		if err == memcached.CASQuit {
			err = callbackErr
		}
		return err
	})
}

// Install a design document.
func (b *Bucket) PutDDoc(docname string, value interface{}) error {
	u, err := b.randomBaseURL()
	if err != nil {
		return err
	}

	j, err := json.Marshal(value)
	if err != nil {
		return err
	}

	u.Path = fmt.Sprintf("/%s/_design/%s", b.Name, docname)
	req, err := http.NewRequest("PUT", u.String(), bytes.NewReader(j))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")

	res, err := HttpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 201 {
		body, _ := ioutil.ReadAll(res.Body)
		return fmt.Errorf("Error installing view: %v / %s",
			res.Status, body)
	}

	return nil
}
<commit_msg>A bit of formatting.<commit_after>/*
A smart client for go.

Usage:

 client, err := couchbase.Connect("http://myserver:8091/")
 handleError(err)
 pool, err := client.GetPool("default")
 handleError(err)
 bucket, err := pool.GetBucket("MyAwesomeBucket")
 handleError(err)
 ...

or a shortcut for the bucket directly

 bucket, err := couchbase.GetBucket("http://myserver:8091/", "default", "default")
*/
package couchbase

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"sync/atomic"

	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

// Execute a function on a memcached connection to the node owning key "k"
//
// Note that this automatically handles transient errors by replaying
// your function on a "not-my-vbucket" error, so don't assume
// your command will only be executed only once.
func (b *Bucket) Do(k string, f func(mc *memcached.Client, vb uint16) error) error {
	vb := b.VBHash(k)
	for {
		masterId := b.VBucketServerMap.VBucketMap[vb][0]
		conn, err := b.connections[masterId].Get()
		defer b.connections[masterId].Return(conn)
		if err != nil {
			return err
		}

		err = f(conn, uint16(vb))
		switch err.(type) {
		default:
			return err
		case gomemcached.MCResponse:
			st := err.(gomemcached.MCResponse).Status
			atomic.AddUint64(&b.pool.client.Statuses[st], 1)
			if st == gomemcached.NOT_MY_VBUCKET {
				b.refresh()
			} else {
				return err
			}
		}
	}
}

type gathered_stats struct {
	sn   string
	vals map[string]string
}

func getStatsParallel(b *Bucket, offset int, which string,
	ch chan<- gathered_stats) {
	sn := b.VBucketServerMap.ServerList[offset]

	results := map[string]string{}
	conn, err := b.connections[offset].Get()
	defer b.connections[offset].Return(conn)
	if err != nil {
		ch <- gathered_stats{sn, results}
	} else {
		st, err := conn.StatsMap(which)
		if err == nil {
			ch <- gathered_stats{sn, st}
		} else {
			ch <- gathered_stats{sn, results}
		}
	}
}

// Get a set of stats from all servers.
//
// Returns a map of server ID -> map of stat key to map value.
func (b *Bucket) GetStats(which string) map[string]map[string]string {
	rv := map[string]map[string]string{}

	if b.VBucketServerMap.ServerList == nil {
		return rv
	}
	// Go grab all the things at once.
	todo := len(b.VBucketServerMap.ServerList)
	ch := make(chan gathered_stats, todo)

	for offset := range b.VBucketServerMap.ServerList {
		go getStatsParallel(b, offset, which, ch)
	}

	// Gather the results
	for i := 0; i < len(b.VBucketServerMap.ServerList); i++ {
		g := <-ch
		if len(g.vals) > 0 {
			rv[g.sn] = g.vals
		}
	}

	return rv
}

func (b *Bucket) doBulkGet(vb uint16, keys []string,
	ch chan<- map[string]*gomemcached.MCResponse) {

	masterId := b.VBucketServerMap.VBucketMap[vb][0]
	conn, err := b.connections[masterId].Get()
	if err != nil {
		ch <- map[string]*gomemcached.MCResponse{}
		return
	}
	defer b.connections[masterId].Return(conn)

	m, err := conn.GetBulk(vb, keys)
	switch err.(type) {
	default:
		ch <- m
	case *gomemcached.MCResponse:
		fmt.Printf("Got a memcached error")
		st := err.(*gomemcached.MCResponse).Status
		atomic.AddUint64(&b.pool.client.Statuses[st], 1)
		if st == gomemcached.NOT_MY_VBUCKET {
			b.refresh()
		}
		ch <- map[string]*gomemcached.MCResponse{}
	}
}

func (b *Bucket) processBulkGet(kdm map[uint16][]string,
	ch chan map[string]*gomemcached.MCResponse) {

	wch := make(chan uint16)

	worker := func() {
		for k := range wch {
			b.doBulkGet(k, kdm[k], ch)
		}
	}

	for i := 0; i < 4; i++ {
		go worker()
	}

	for k := range kdm {
		wch <- k
	}
	close(wch)

}
func (b *Bucket) GetBulk(keys []string) map[string]*gomemcached.MCResponse {
	// Organize by vbucket
	kdm := map[uint16][]string{}
	for _, k := range keys {
		vb := uint16(b.VBHash(k))
		a, ok := kdm[vb]
		if !ok {
			a = []string{}
		}
		kdm[vb] = append(a, k)
	}

	ch := make(chan map[string]*gomemcached.MCResponse)
	defer close(ch)

	go b.processBulkGet(kdm, ch)

	rv := map[string]*gomemcached.MCResponse{}
	for _ = range kdm {
		m := <-ch
		for k, v := range m {
			rv[k] = v
		}
	}

	return rv
}

// Set a value in this bucket.
// The value will be serialized into a JSON document.
func (b *Bucket) Set(k string, exp int, v interface{}) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return b.SetRaw(k, exp, data)
}

// Set a value in this bucket.
// The value will be stored as raw bytes.
func (b *Bucket) SetRaw(k string, exp int, v []byte) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		_, err := mc.Set(vb, k, 0, exp, v)
		if err != nil {
			return err
		}
		return nil
	})
}

// Adds a value to this bucket; like Set except that nothing happens
// if the key exists.  The value will be serialized into a JSON
// document.
func (b *Bucket) Add(k string, exp int, v interface{}) (added bool, err error) {
	data, err := json.Marshal(v)
	if err != nil {
		return false, err
	}
	return b.AddRaw(k, exp, data)
}

// Adds a value to this bucket; like SetRaw except that nothing
// happens if the key exists.  The value will be stored as raw bytes.
func (b *Bucket) AddRaw(k string, exp int, v []byte) (added bool, err error) {
	err = b.Do(k, func(mc *memcached.Client, vb uint16) error {
		switch res, err := memcached.UnwrapMemcachedError(mc.Add(vb, k, 0, exp, v)); {
		case err != nil:
			return err
		case res.Status == gomemcached.SUCCESS:
			added = true
		case res.Status != gomemcached.KEY_EEXISTS:
			return res
		}
		return nil
	})
	return
}

// Get a raw value from this bucket, including its CAS counter.
func (b *Bucket) GetsRaw(k string, cas *uint64) ([]byte, error) {
	var data []byte
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Get(vb, k)
		if err != nil {
			return err
		}
		if cas != nil {
			*cas = res.Cas
		}
		data = res.Body
		return nil
	})
	return data, err
}

// Get a value from this bucket, including its CAS counter.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Gets(k string, rv interface{}, cas *uint64) error {
	data, err := b.GetsRaw(k, cas)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, rv)
}

// Get a value from this bucket.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Get(k string, rv interface{}) error {
	return b.Gets(k, rv, nil)
}

// Get a raw value from this bucket.
func (b *Bucket) GetRaw(k string) ([]byte, error) {
	return b.GetsRaw(k, nil)
}

// Delete a key from this bucket.
func (b *Bucket) Delete(k string) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		_, err := mc.Del(vb, k)
		if err != nil {
			return err
		}
		return nil
	})
}

// Increment a key
func (b *Bucket) Incr(k string, amt, def uint64, exp int) (uint64, error) {
	var rv uint64
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Incr(vb, k, amt, def, exp)
		if err != nil {
			return err
		}
		rv = res
		return nil
	})
	return rv, err
}

// A callback function to update a document
type UpdateFunc func(current []byte) (updated []byte, err error)

// Return this as the error from an UpdateFunc to cancel the Update
// operation.
const UpdateCancel = memcached.CASQuit

// Safe update of a document, avoiding conflicts by using CAS.
//
// The callback function will be invoked with the current raw document
// contents (or nil if the document doesn't exist); it should return
// the updated raw contents (or nil to delete.)  If it decides not to
// change anything it can return UpdateCancel as the error.
//
// If another writer modifies the document between the get and the
// set, the callback will be invoked again with the newer value.
func (b *Bucket) Update(k string, exp int, callback UpdateFunc) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		var callbackErr error
		casFunc := func(current []byte) ([]byte, memcached.CasOp) {
			var updated []byte
			updated, callbackErr = callback(current)
			if callbackErr != nil {
				return nil, memcached.CASQuit
			} else if updated == nil {
				return updated, memcached.CASDelete
			} else {
				return updated, memcached.CASStore
			}
		}
		_, err := mc.CAS(vb, k, casFunc, exp)
		if err == memcached.CASQuit {
			err = callbackErr
		}
		return err
	})
}

// Install a design document.
func (b *Bucket) PutDDoc(docname string, value interface{}) error {
	u, err := b.randomBaseURL()
	if err != nil {
		return err
	}

	j, err := json.Marshal(value)
	if err != nil {
		return err
	}

	u.Path = fmt.Sprintf("/%s/_design/%s", b.Name, docname)
	req, err := http.NewRequest("PUT", u.String(), bytes.NewReader(j))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")

	res, err := HttpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 201 {
		body, _ := ioutil.ReadAll(res.Body)
		return fmt.Errorf("Error installing view: %v / %s",
			res.Status, body)
	}

	return nil
}
<|endoftext|>


<commit_before>package twiliolo

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

// ROOT is the baseURL of the Twilio API
const ROOT = "https://api.twilio.com"

// VERSION is the version of the Twilio API
const VERSION = "2010-04-01"

// Client is the interface of a Twilio Client
type Client interface {
	AccountSid() string
	AuthToken() string
	RootURL() string
	Get(string, []RequestOption) ([]byte, error)
	Post(string, []RequestOption, url.Values) ([]byte, error)
	Delete(string, []RequestOption) error
}

// HTTPClient is the interface of an HTTP client making a request
type HTTPClient interface {
	Do(req *http.Request) (*http.Response, error)
}

// TwilioClient is the Twilio API client
type TwilioClient struct {
	accountSid string
	authToken  string
	rootURL    string
	httpClient HTTPClient
}

var _ Client = &TwilioClient{}

// NewClient instanciates a new TwilioClient
func NewClient(accountSid, authToken string, httpClient HTTPClient) *TwilioClient {
	rootURL := ROOT + "/" + VERSION + "/Accounts/" + accountSid
	return &TwilioClient{accountSid, authToken, rootURL, httpClient}
}

func (c *TwilioClient) Post(uri string, requestOptions []RequestOption, values url.Values) ([]byte, error) {
	uri, err := c.buildURI(uri, requestOptions)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest("POST", uri, strings.NewReader(values.Encode()))
	if err != nil {
		return nil, err
	}

	req.SetBasicAuth(c.AccountSid(), c.AuthToken())
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")

	res, err := c.httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return body, err
	}

	if res.StatusCode != 200 && res.StatusCode != 201 {
		if res.StatusCode == 500 {
			return body, ErrTwilioServer
		}

		twilioError := new(TwilioError)
		json.Unmarshal(body, twilioError)
		return body, twilioError

	}

	return body, err
}

func (c *TwilioClient) Get(uri string, requestOptions []RequestOption) ([]byte, error) {
	uri, err := c.buildURI(uri, requestOptions)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest("GET", uri, nil)
	if err != nil {
		return nil, err
	}

	req.SetBasicAuth(c.AccountSid(), c.AuthToken())

	res, err := c.httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return body, err
	}

	if res.StatusCode != 200 && res.StatusCode != 201 {
		if res.StatusCode == 500 {
			return body, ErrTwilioServer
		}
		twilioError := new(TwilioError)
		err := json.Unmarshal(body, &twilioError)
		if err != nil {
			return body, err
		}

		return body, twilioError
	}

	return body, err
}

func (c *TwilioClient) Delete(uri string, requestOptions []RequestOption) error {
	uri, err := c.buildURI(uri, requestOptions)
	if err != nil {
		return err
	}

	req, err := http.NewRequest("DELETE", uri, nil)
	if err != nil {
		return err
	}

	req.SetBasicAuth(c.AccountSid(), c.AuthToken())

	res, err := c.httpClient.Do(req)
	if err != nil {
		return err
	}

	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}

	if res.StatusCode != 204 {
		if res.StatusCode == 500 {
			return ErrTwilioServer
		}

		twilioError := new(TwilioError)
		err := json.Unmarshal(body, &twilioError)
		if err != nil {
			return err
		}

		return twilioError
	}

	return nil
}

// AccountSid returns the AccountSid of the client
func (c *TwilioClient) AccountSid() string {
	return c.accountSid
}

// AuthToken returns the AuthToken of the client
func (c *TwilioClient) AuthToken() string {
	return c.authToken
}

// RootURL returns the RootURL of the client
func (c *TwilioClient) RootURL() string {
	return c.rootURL
}

func (c *TwilioClient) buildURI(uri string, requestOptions []RequestOption) (string, error) {
	uri = strings.Trim(uri, "/")
	if uri == "" {
		return "", errors.New("Empty URI")
	}

	parts := make([]string, 0)

	// Check for "http" because sometimes we get raw URLs from following the metadata.
	if !strings.HasPrefix(uri, "http") {
		parts = append(parts, c.RootURL())
	}

	parts = append(parts, uri)

	u, err := url.Parse(strings.Join(parts, "/"))
	if err != nil {
		return "", err
	}

	q := u.Query()
	for _, option := range requestOptions {
		key, value := option.GetValue()
		q.Add(key, value)
	}
	u.RawQuery = q.Encode()

	return u.String(), nil
}
<commit_msg>Change method BuildURI name to BuildURL<commit_after>package twiliolo

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

// ROOT is the baseURL of the Twilio API
const ROOT = "https://api.twilio.com"

// VERSION is the version of the Twilio API
const VERSION = "2010-04-01"

// Client is the interface of a Twilio Client
type Client interface {
	AccountSid() string
	AuthToken() string
	RootURL() string
	Get(string, []RequestOption) ([]byte, error)
	Post(string, []RequestOption, url.Values) ([]byte, error)
	Delete(string, []RequestOption) error
}

// HTTPClient is the interface of an HTTP client making a request
type HTTPClient interface {
	Do(req *http.Request) (*http.Response, error)
}

// TwilioClient is the Twilio API client
type TwilioClient struct {
	accountSid string
	authToken  string
	rootURL    string
	httpClient HTTPClient
}

var _ Client = &TwilioClient{}

// NewClient instanciates a new TwilioClient
func NewClient(accountSid, authToken string, httpClient HTTPClient) *TwilioClient {
	rootURL := ROOT + "/" + VERSION + "/Accounts/" + accountSid
	return &TwilioClient{accountSid, authToken, rootURL, httpClient}
}

func (c *TwilioClient) Post(uri string, requestOptions []RequestOption, values url.Values) ([]byte, error) {
	uri, err := c.buildURL(uri, requestOptions)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest("POST", uri, strings.NewReader(values.Encode()))
	if err != nil {
		return nil, err
	}

	req.SetBasicAuth(c.AccountSid(), c.AuthToken())
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")

	res, err := c.httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return body, err
	}

	if res.StatusCode != 200 && res.StatusCode != 201 {
		if res.StatusCode == 500 {
			return body, ErrTwilioServer
		}

		twilioError := new(TwilioError)
		json.Unmarshal(body, twilioError)
		return body, twilioError

	}

	return body, err
}

func (c *TwilioClient) Get(uri string, requestOptions []RequestOption) ([]byte, error) {
	uri, err := c.buildURL(uri, requestOptions)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest("GET", uri, nil)
	if err != nil {
		return nil, err
	}

	req.SetBasicAuth(c.AccountSid(), c.AuthToken())

	res, err := c.httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return body, err
	}

	if res.StatusCode != 200 && res.StatusCode != 201 {
		if res.StatusCode == 500 {
			return body, ErrTwilioServer
		}
		twilioError := new(TwilioError)
		err := json.Unmarshal(body, &twilioError)
		if err != nil {
			return body, err
		}

		return body, twilioError
	}

	return body, err
}

func (c *TwilioClient) Delete(uri string, requestOptions []RequestOption) error {
	uri, err := c.buildURL(uri, requestOptions)
	if err != nil {
		return err
	}

	req, err := http.NewRequest("DELETE", uri, nil)
	if err != nil {
		return err
	}

	req.SetBasicAuth(c.AccountSid(), c.AuthToken())

	res, err := c.httpClient.Do(req)
	if err != nil {
		return err
	}

	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}

	if res.StatusCode != 204 {
		if res.StatusCode == 500 {
			return ErrTwilioServer
		}

		twilioError := new(TwilioError)
		err := json.Unmarshal(body, &twilioError)
		if err != nil {
			return err
		}

		return twilioError
	}

	return nil
}

// AccountSid returns the AccountSid of the client
func (c *TwilioClient) AccountSid() string {
	return c.accountSid
}

// AuthToken returns the AuthToken of the client
func (c *TwilioClient) AuthToken() string {
	return c.authToken
}

// RootURL returns the RootURL of the client
func (c *TwilioClient) RootURL() string {
	return c.rootURL
}

func (c *TwilioClient) buildURL(uri string, requestOptions []RequestOption) (string, error) {
	uri = strings.Trim(uri, "/")
	if uri == "" {
		return "", errors.New("Empty URI")
	}

	var urlStr string
	// Check for "http" because sometimes we get raw URLs from following the metadata.
	if !strings.HasPrefix(uri, "http") {
		urlStr = c.RootURL() + "/" + uri
	} else {
		urlStr = uri
	}

	u, err := url.Parse(urlStr)
	if err != nil {
		return "", err
	}

	q := u.Query()
	for _, option := range requestOptions {
		key, value := option.GetValue()
		q.Add(key, value)
	}
	u.RawQuery = q.Encode()

	return u.String(), nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"log"
	"net/url"
	"reflect"
	"strconv"

	"github.com/gorilla/websocket"
	"bytes"
	"fmt"
	"time"
)

func execAction1(param float64) string {
	time.Sleep(1000 * time.Millisecond)
	return "result from action1 param=" + strconv.FormatFloat(param, 'f', -1, 64)
}

func execAction2(param string) string {
	time.Sleep(2000 * time.Millisecond)
	return "result from action2 param=" + param
}

func execAction3(param1 float64, param2 bool) string {
	time.Sleep(4000 * time.Millisecond)
	return "result from action3 param1=" + strconv.FormatFloat(param1, 'f', -1, 64) + " param2=" + strconv.FormatBool(param2)
}

func execAction(actionMap map[string]interface{},
actionName string,
actionParameters []ActionParameter) (string, error) {

	// Retrieve the function from the map
	actionFunc := actionMap[actionName]
	// Test if the function was retrieved
	if actionFunc == nil {
		return "", errors.New("Unknown action name")
	}

	// func ValueOf(i interface{}) Value
	// ValueOf returns a new Value initialized to the concrete value stored in the interface i.
	f := reflect.ValueOf(actionFunc)

	// NumIn returns a function type's input parameter count.
	// It panics if the type's Kind is not Func.
	// TODO: check if f is a function
	if len(actionParameters) != f.Type().NumIn() {
		return "", errors.New("Wrong number of parameters")
	}

	var buffer bytes.Buffer
	inputParameters := make([]reflect.Value, len(actionParameters))
	for k, param := range actionParameters {
		buffer.WriteString(fmt.Sprintf("  parameter[%d]=%s", k, reflect.TypeOf(param.Value)))
		inputParameters[k] = reflect.ValueOf(param.Value)
	}
	//log.Printf("action.name=%s %s\n", actionName, buffer.String())

	// func (v Value) Call(in []Value) []Value
	// Cf. https://golang.org/pkg/reflect/#Value.Call
	results := f.Call(inputParameters)

	// TODO: check if len(results) == 1, and if results[0] is really a string
	result := results[0].String()

	return result, nil
}

var addr = flag.String("addr", "localhost:8080", "http service address")
var token = flag.String("token", "1234", "token to identify this client")

/**
 * createWebSocketConnection establishes a WebSocket to a server with a specific url.
 * But the server can be down, so we have to retry to create the WebSocket every x second.
 * TODO: find if it's the best solution to create a WebSocket with retries
 */
func createWebSocketConnection(u url.URL) *websocket.Conn {

	for {
		c, httpResponse, err := websocket.DefaultDialer.Dial(u.String(), nil)
		if err == nil {
			return c

		}
		log.Println("dial: err=", err, "httpResponse=", httpResponse)
		time.Sleep(2 * time.Second)
	}
}

func main() {

	flag.Parse()
	log.SetFlags(log.Ldate | log.Ltime)

	actionMap := map[string]interface{}{
		"action1": execAction1,
		"action2": execAction2,
		"action3": execAction3,
	}

	u := url.URL{Scheme: "ws", Host: *addr, Path: "/action", RawQuery: "token=" + *token}
	log.Printf("connecting to %s", u.String())

	// Create a WebSocket to a server with a specific url.
	c := createWebSocketConnection(u)

	defer c.Close()

	for {
		messageType, message, err := c.ReadMessage()
		if err != nil {
			log.Println("read: err=", err)
			c = createWebSocketConnection(u)
			continue
		}

		if messageType != websocket.TextMessage {
			log.Println("received a non text message")
			continue
		}

		go func() {
			var action Action
			json.Unmarshal(message, &action)

			log.Printf("received action[%s] name=%s\n", action.ID, action.Name)

			result, err := execAction(actionMap, action.Name, action.Parameters)
			if err != nil {
				log.Println("Error:  [", err, "]")
				action.Result = "action not found"
			} else {
				action.Result = result
			}

			jsonAction, err := json.Marshal(action)
			if err != nil {
				log.Println("marshal:", err)
				return
			}

			err = c.WriteMessage(websocket.TextMessage, []byte(jsonAction))
			if err != nil {
				log.Println("write:", err)
				return
			}
			log.Printf("sent action[%s] name=%s  result=%s\n", action.ID, action.Name, action.Result)
		}()
	}
}
<commit_msg>Try websocket.IsUnexpectedCloseError() to identify an error from the server<commit_after>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"log"
	"net/url"
	"reflect"
	"strconv"

	"github.com/gorilla/websocket"
	"bytes"
	"fmt"
	"time"
)

func execAction1(param float64) string {
	time.Sleep(1000 * time.Millisecond)
	return "result from action1 param=" + strconv.FormatFloat(param, 'f', -1, 64)
}

func execAction2(param string) string {
	time.Sleep(2000 * time.Millisecond)
	return "result from action2 param=" + param
}

func execAction3(param1 float64, param2 bool) string {
	time.Sleep(4000 * time.Millisecond)
	return "result from action3 param1=" + strconv.FormatFloat(param1, 'f', -1, 64) + " param2=" + strconv.FormatBool(param2)
}

func execAction(actionMap map[string]interface{},
actionName string,
actionParameters []ActionParameter) (string, error) {

	// Retrieve the function from the map
	actionFunc := actionMap[actionName]
	// Test if the function was retrieved
	if actionFunc == nil {
		return "", errors.New("Unknown action name")
	}

	// func ValueOf(i interface{}) Value
	// ValueOf returns a new Value initialized to the concrete value stored in the interface i.
	f := reflect.ValueOf(actionFunc)

	// NumIn returns a function type's input parameter count.
	// It panics if the type's Kind is not Func.
	// TODO: check if f is a function
	if len(actionParameters) != f.Type().NumIn() {
		return "", errors.New("Wrong number of parameters")
	}

	var buffer bytes.Buffer
	inputParameters := make([]reflect.Value, len(actionParameters))
	for k, param := range actionParameters {
		buffer.WriteString(fmt.Sprintf("  parameter[%d]=%s", k, reflect.TypeOf(param.Value)))
		inputParameters[k] = reflect.ValueOf(param.Value)
	}

	// func (v Value) Call(in []Value) []Value
	// Cf. https://golang.org/pkg/reflect/#Value.Call
	results := f.Call(inputParameters)

	// TODO: check if len(results) == 1, and if results[0] is really a string
	result := results[0].String()

	return result, nil
}

var addr = flag.String("addr", "localhost:8080", "http service address")
var token = flag.String("token", "1234", "token to identify this client")

/**
 * createWebSocketConnection establishes a WebSocket to a server with a specific url.
 * But the server can be down, so we have to retry to create the WebSocket every x second.
 * TODO: find if it's the best solution to create a WebSocket with retries
 */
func createWebSocketConnection(u url.URL) *websocket.Conn {

	for {
		c, httpResponse, err := websocket.DefaultDialer.Dial(u.String(), nil)
		if err == nil {
			return c

		}
		log.Println("dial: err=", err, "httpResponse=", httpResponse)
		time.Sleep(2 * time.Second)
	}
}

func main() {

	flag.Parse()
	log.SetFlags(log.Ldate | log.Ltime)

	actionMap := map[string]interface{}{
		"action1": execAction1,
		"action2": execAction2,
		"action3": execAction3,
	}

	u := url.URL{Scheme: "ws", Host: *addr, Path: "/action", RawQuery: "token=" + *token}
	log.Printf("connecting to %s", u.String())

	// Create a WebSocket to a server with a specific url.
	c := createWebSocketConnection(u)

	defer c.Close()

	for {
		messageType, message, err := c.ReadMessage()
		if err != nil {
			log.Println("read: err=", err)
			// TODO: manage the different possible errors
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				log.Printf("IsUnexpectedCloseError: %v\n", err)
				c = createWebSocketConnection(u)
			}
			continue
		}

		if messageType != websocket.TextMessage {
			log.Println("received a non text message")
			continue
		}

		go func() {
			var action Action
			json.Unmarshal(message, &action)

			log.Printf("received action[%s] name=%s\n", action.ID, action.Name)

			result, err := execAction(actionMap, action.Name, action.Parameters)
			if err != nil {
				log.Println("Error:  [", err, "]")
				action.Result = "action not found"
			} else {
				action.Result = result
			}

			jsonAction, err := json.Marshal(action)
			if err != nil {
				log.Println("marshal:", err)
				return
			}

			err = c.WriteMessage(websocket.TextMessage, []byte(jsonAction))
			if err != nil {
				log.Println("write:", err)
				return
			}
			log.Printf("sent action[%s] name=%s  result=%s\n", action.ID, action.Name, action.Result)
		}()
	}
}
<|endoftext|>


<commit_before>/*
client.go implements RFC 6143 §7.5 Client-to-Server Messages.
See http://tools.ietf.org/html/rfc6143#section-7.5 for more info.
*/
package vnc

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"unicode"
)

const (
	setPixelFormatMsg = uint8(iota)
	_
	setEncodingsMsg
	framebufferUpdateRequestMsg
	keyEventMsg
	pointerEventMsg
	clientCutTextMsg
)

// SetPixelFormat sets the format in which pixel values should be sent
// in FramebufferUpdate messages from the server.
//
// See RFC 6143 Section 7.5.1
func (c *ClientConn) SetPixelFormat(pf PixelFormat) error {
	var buf bytes.Buffer

	// message-type
	if err := binary.Write(&buf, binary.BigEndian, setPixelFormatMsg); err != nil {
		return err
	}
	//padding
	padding := [3]byte{}
	buf.Write(padding[:])
	// pixel-format
	pfBytes, err := pf.Bytes()
	if err != nil {
		return err
	}
	buf.Write(pfBytes)

	// Send the data down the connection.
	if _, err := c.c.Write(buf.Bytes()); err != nil {
		return err
	}

	// Invalidate the color map.
	if pf.TrueColor == RFBFalse {
		c.colorMap = [256]Color{}
	}

	c.pixelFormat = pf
	return nil
}

type SetEncodingsMessage struct {
	Msg     uint8   // message-type
	Pad     [1]byte // padding
	NumEncs uint16  // number-of-encodings
	Encs    []int32 // encoding-type
}

// SetEncodings sets the encoding types in which the pixel data can be sent
// from the server. After calling this method, the encs slice given should not
// be modified.
//
// See RFC 6143 Section 7.5.2
func (c *ClientConn) SetEncodings(e []Encoding) error {
	// Prepare message.
	msg := SetEncodingsMessage{
		Msg:     setEncodingsMsg,
		NumEncs: uint16(len(e)),
	}
	for _, v := range e {
		msg.Encs = append(msg.Encs, int32(v.Type()))
	}

	// Send message.
	if err := binary.Write(c.c, binary.BigEndian, &msg.Msg); err != nil {
		return err
	}
	for i := range msg.Pad {
		if err := binary.Write(c.c, binary.BigEndian, &msg.Pad[i]); err != nil {
			return err
		}
	}
	if err := binary.Write(c.c, binary.BigEndian, &msg.NumEncs); err != nil {
		return err
	}
	for i := range msg.Encs {
		if err := binary.Write(c.c, binary.BigEndian, &msg.Encs[i]); err != nil {
			return err
		}
	}
	c.encodings = e

	return nil
}

// FramebufferUpdateRequestMessage defines a FramebufferUpdateRequest message.
type FramebufferUpdateRequestMessage struct {
	Msg                 uint8
	Inc                 uint8
	X, Y, Width, Height uint16
}

// Requests a framebuffer update from the server. There may be an indefinite
// time between the request and the actual framebuffer update being received.
//
// See RFC 6143 Section 7.5.3
func (c *ClientConn) FramebufferUpdateRequest(inc uint8, x, y, w, h uint16) error {
	msg := FramebufferUpdateRequestMessage{framebufferUpdateRequestMsg, inc, x, y, w, h}
	if err := binary.Write(c.c, binary.BigEndian, &msg); err != nil {
		return err
	}
	return nil
}

// KeyEventMessage holds the wire format message.
type KeyEventMessage struct {
	Msg      uint8
	DownFlag uint8
	_        [2]byte
	Key      uint32
}

const (
	PressKey   = true
	ReleaseKey = false
)

// KeyEvent indicates a key press or release and sends it to the server.
// The key is indicated using the X Window System "keysym" value. Use
// Google to find a reference of these values. To simulate a key press,
// you must send a key with both a down event, and a non-down event.
//
// See RFC 6143 Section 7.5.4.
func (c *ClientConn) KeyEvent(keysym uint32, down bool) error {
	var downFlag uint8 = RFBFalse
	if down {
		downFlag = RFBTrue
	}

	msg := KeyEventMessage{keyEventMsg, downFlag, [2]byte{}, keysym}
	if err := binary.Write(c.c, binary.BigEndian, msg); err != nil {
		return err
	}

	settleUI()
	return nil
}

// ButtonMask represents a mask of pointer presses/releases.
type ButtonMask uint8

// All available button mask components.
const (
	ButtonLeft ButtonMask = 1 << iota
	ButtonMiddle
	ButtonRight
	Button4
	Button5
	Button6
	Button7
	Button8
	ButtonNone = ButtonMask(0)
)

// PointerEventMessage holds the wire format message.
type PointerEventMessage struct {
	Msg  uint8
	Mask uint8
	X, Y uint16
}

// PointerEvent indicates that pointer movement or a pointer button
// press or release.
//
// The mask is a bitwise mask of various ButtonMask values. When a button
// is set, it is pressed, when it is unset, it is released.
//
// See RFC 6143 Section 7.5.5
func (c *ClientConn) PointerEvent(mask ButtonMask, x, y uint16) error {
	msg := PointerEventMessage{pointerEventMsg, uint8(mask), x, y}
	if err := binary.Write(c.c, binary.BigEndian, msg); err != nil {
		return err
	}

	settleUI()
	return nil
}

// ClientCutText tells the server that the client has new text in its cut buffer.
// The text string MUST only contain Latin-1 characters. This encoding
// is compatible with Go's native string format, but can only use up to
// unicode.MaxLatin values.
//
// See RFC 6143 Section 7.5.6
func (c *ClientConn) ClientCutText(text string) error {
	var buf bytes.Buffer

	// This is the fixed size data we'll send
	fixedData := []interface{}{
		clientCutTextMsg,
		uint8(0),
		uint8(0),
		uint8(0),
		uint32(len(text)),
	}

	for _, val := range fixedData {
		if err := binary.Write(&buf, binary.BigEndian, val); err != nil {
			return err
		}
	}

	for _, char := range text {
		if char > unicode.MaxLatin1 {
			return fmt.Errorf("Character '%s' is not valid Latin-1", char)
		}

		if err := binary.Write(&buf, binary.BigEndian, uint8(char)); err != nil {
			return err
		}
	}

	dataLength := 8 + len(text)
	if _, err := c.c.Write(buf.Bytes()[0:dataLength]); err != nil {
		return err
	}

	return nil
}

//
// Constants to use for KeyEvents PointerEvents.
//

// Latin 1 (byte 3 = 0)
// ISO/IEC 8859-1 = Unicode U+0020..U+00FF
const (
	KeySpace = iota + 0x0020
	KeyExclam
	KeyQuoteDbl
	KeyNumberSign
	KeyDollar
	KeyPercent
	KeyAmpersand
	KeyApostrophe
	KeyParenLeft
	KeyParenRight
	KeyAsterisk
	KeyPlus
	KeyComma
	KeyMinus
	KeyPeriod
	KeySlash
	Key0
	Key1
	Key2
	Key3
	Key4
	Key5
	Key6
	Key7
	Key8
	Key9
	KeyColon
	KeySemicolon
	KeyLess
	KeyEqual
	KeyGreater
	KeyQuestion
	KeyAt
	KeyA
	KeyB
	KeyC
	KeyD
	KeyE
	KeyF
	KeyG
	KeyH
	KeyI
	KeyJ
	KeyK
	KeyL
	KeyM
	KeyN
	KeyO
	KeyP
	KeyQ
	KeyR
	KeyS
	KeyT
	KeyU
	KeyV
	KeyW
	KeyX
	KeyY
	KeyZ
	KeyBracketLeft
	KeyBackslash
	KeyBracketRight
	KeyAsciiCircum
	KeyUnderscore
	KeyGrave
	Keya
	Keyb
	Keyc
	Keyd
	Keye
	Keyf
	Keyg
	Keyh
	Keyi
	Keyj
	Keyk
	Keyl
	Keym
	Keyn
	Keyo
	Keyp
	Keyq
	Keyr
	Keys
	Keyt
	Keyu
	Keyv
	Keyw
	Keyx
	Keyy
	Keyz
	KeyBraceLeft
	KeyBar
	KeyBraceRight
	KeyAsciiTilde
)
const (
	KeyBackspace = iota + 0xff08
	KeyTab
	KeyLinefeed
	KeyClear
	_
	KeyReturn
)
const (
	KeyPause      = 0xff13
	KeyScrollLock = 0xff14
	KeySysReq     = 0xff15
	KeyEscape     = 0xff1b
)
const (
	KeyHome = iota + 0xff50
	KeyLeft
	KeyUp
	KeyRight
	KeyDown
	KeyPageUp
	KeyPageDown
	KeyEnd
	KeyInsert = 0xff63
)
const (
	KeyF1 = iota + 0xffbe
	KeyF2
	KeyF3
	KeyF4
	KeyF5
	KeyF6
	KeyF7
	KeyF8
	KeyF9
	KeyF10
	KeyF11
	KeyF12
)
const (
	KeyShiftLeft = iota + 0xffe1
	KeyShiftRight
	KeyControlLeft
	KeyControlRight
	KeyCapsLock
	_
	_
	_
	KeyAltLeft
	KeyAltRight

	KeyDelete = 0xffff
)
<commit_msg>removed extra whitespace<commit_after>/*
client.go implements RFC 6143 §7.5 Client-to-Server Messages.
See http://tools.ietf.org/html/rfc6143#section-7.5 for more info.
*/
package vnc

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"unicode"
)

const (
	setPixelFormatMsg = uint8(iota)
	_
	setEncodingsMsg
	framebufferUpdateRequestMsg
	keyEventMsg
	pointerEventMsg
	clientCutTextMsg
)

// SetPixelFormat sets the format in which pixel values should be sent
// in FramebufferUpdate messages from the server.
//
// See RFC 6143 Section 7.5.1
func (c *ClientConn) SetPixelFormat(pf PixelFormat) error {
	var buf bytes.Buffer

	// message-type
	if err := binary.Write(&buf, binary.BigEndian, setPixelFormatMsg); err != nil {
		return err
	}
	//padding
	padding := [3]byte{}
	buf.Write(padding[:])
	// pixel-format
	pfBytes, err := pf.Bytes()
	if err != nil {
		return err
	}
	buf.Write(pfBytes)

	// Send the data down the connection.
	if _, err := c.c.Write(buf.Bytes()); err != nil {
		return err
	}

	// Invalidate the color map.
	if pf.TrueColor == RFBFalse {
		c.colorMap = [256]Color{}
	}

	c.pixelFormat = pf
	return nil
}

type SetEncodingsMessage struct {
	Msg     uint8   // message-type
	Pad     [1]byte // padding
	NumEncs uint16  // number-of-encodings
	Encs    []int32 // encoding-type
}

// SetEncodings sets the encoding types in which the pixel data can be sent
// from the server. After calling this method, the encs slice given should not
// be modified.
//
// See RFC 6143 Section 7.5.2
func (c *ClientConn) SetEncodings(e []Encoding) error {
	// Prepare message.
	msg := SetEncodingsMessage{
		Msg:     setEncodingsMsg,
		NumEncs: uint16(len(e)),
	}
	for _, v := range e {
		msg.Encs = append(msg.Encs, int32(v.Type()))
	}

	// Send message.
	if err := binary.Write(c.c, binary.BigEndian, &msg.Msg); err != nil {
		return err
	}
	for i := range msg.Pad {
		if err := binary.Write(c.c, binary.BigEndian, &msg.Pad[i]); err != nil {
			return err
		}
	}
	if err := binary.Write(c.c, binary.BigEndian, &msg.NumEncs); err != nil {
		return err
	}
	for i := range msg.Encs {
		if err := binary.Write(c.c, binary.BigEndian, &msg.Encs[i]); err != nil {
			return err
		}
	}
	c.encodings = e

	return nil
}

// FramebufferUpdateRequestMessage defines a FramebufferUpdateRequest message.
type FramebufferUpdateRequestMessage struct {
	Msg                 uint8
	Inc                 uint8
	X, Y, Width, Height uint16
}

// Requests a framebuffer update from the server. There may be an indefinite
// time between the request and the actual framebuffer update being received.
//
// See RFC 6143 Section 7.5.3
func (c *ClientConn) FramebufferUpdateRequest(inc uint8, x, y, w, h uint16) error {
	msg := FramebufferUpdateRequestMessage{framebufferUpdateRequestMsg, inc, x, y, w, h}
	if err := binary.Write(c.c, binary.BigEndian, &msg); err != nil {
		return err
	}
	return nil
}

// KeyEventMessage holds the wire format message.
type KeyEventMessage struct {
	Msg      uint8
	DownFlag uint8
	_        [2]byte
	Key      uint32
}

const (
	PressKey   = true
	ReleaseKey = false
)

// KeyEvent indicates a key press or release and sends it to the server.
// The key is indicated using the X Window System "keysym" value. Use
// Google to find a reference of these values. To simulate a key press,
// you must send a key with both a down event, and a non-down event.
//
// See RFC 6143 Section 7.5.4.
func (c *ClientConn) KeyEvent(keysym uint32, down bool) error {
	var downFlag uint8 = RFBFalse
	if down {
		downFlag = RFBTrue
	}

	msg := KeyEventMessage{keyEventMsg, downFlag, [2]byte{}, keysym}
	if err := binary.Write(c.c, binary.BigEndian, msg); err != nil {
		return err
	}

	settleUI()
	return nil
}

// ButtonMask represents a mask of pointer presses/releases.
type ButtonMask uint8

// All available button mask components.
const (
	ButtonLeft ButtonMask = 1 << iota
	ButtonMiddle
	ButtonRight
	Button4
	Button5
	Button6
	Button7
	Button8
	ButtonNone = ButtonMask(0)
)

// PointerEventMessage holds the wire format message.
type PointerEventMessage struct {
	Msg  uint8
	Mask uint8
	X, Y uint16
}

// PointerEvent indicates that pointer movement or a pointer button
// press or release.
//
// The mask is a bitwise mask of various ButtonMask values. When a button
// is set, it is pressed, when it is unset, it is released.
//
// See RFC 6143 Section 7.5.5
func (c *ClientConn) PointerEvent(mask ButtonMask, x, y uint16) error {
	msg := PointerEventMessage{pointerEventMsg, uint8(mask), x, y}
	if err := binary.Write(c.c, binary.BigEndian, msg); err != nil {
		return err
	}

	settleUI()
	return nil
}

// ClientCutText tells the server that the client has new text in its cut buffer.
// The text string MUST only contain Latin-1 characters. This encoding
// is compatible with Go's native string format, but can only use up to
// unicode.MaxLatin values.
//
// See RFC 6143 Section 7.5.6
func (c *ClientConn) ClientCutText(text string) error {
	var buf bytes.Buffer

	// This is the fixed size data we'll send
	fixedData := []interface{}{
		clientCutTextMsg,
		uint8(0),
		uint8(0),
		uint8(0),
		uint32(len(text)),
	}

	for _, val := range fixedData {
		if err := binary.Write(&buf, binary.BigEndian, val); err != nil {
			return err
		}
	}

	for _, char := range text {
		if char > unicode.MaxLatin1 {
			return fmt.Errorf("Character '%s' is not valid Latin-1", char)
		}

		if err := binary.Write(&buf, binary.BigEndian, uint8(char)); err != nil {
			return err
		}
	}

	dataLength := 8 + len(text)
	if _, err := c.c.Write(buf.Bytes()[0:dataLength]); err != nil {
		return err
	}

	return nil
}

//
// Constants to use for KeyEvents PointerEvents.
//

// Latin 1 (byte 3 = 0)
// ISO/IEC 8859-1 = Unicode U+0020..U+00FF
const (
	KeySpace = iota + 0x0020
	KeyExclam
	KeyQuoteDbl
	KeyNumberSign
	KeyDollar
	KeyPercent
	KeyAmpersand
	KeyApostrophe
	KeyParenLeft
	KeyParenRight
	KeyAsterisk
	KeyPlus
	KeyComma
	KeyMinus
	KeyPeriod
	KeySlash
	Key0
	Key1
	Key2
	Key3
	Key4
	Key5
	Key6
	Key7
	Key8
	Key9
	KeyColon
	KeySemicolon
	KeyLess
	KeyEqual
	KeyGreater
	KeyQuestion
	KeyAt
	KeyA
	KeyB
	KeyC
	KeyD
	KeyE
	KeyF
	KeyG
	KeyH
	KeyI
	KeyJ
	KeyK
	KeyL
	KeyM
	KeyN
	KeyO
	KeyP
	KeyQ
	KeyR
	KeyS
	KeyT
	KeyU
	KeyV
	KeyW
	KeyX
	KeyY
	KeyZ
	KeyBracketLeft
	KeyBackslash
	KeyBracketRight
	KeyAsciiCircum
	KeyUnderscore
	KeyGrave
	Keya
	Keyb
	Keyc
	Keyd
	Keye
	Keyf
	Keyg
	Keyh
	Keyi
	Keyj
	Keyk
	Keyl
	Keym
	Keyn
	Keyo
	Keyp
	Keyq
	Keyr
	Keys
	Keyt
	Keyu
	Keyv
	Keyw
	Keyx
	Keyy
	Keyz
	KeyBraceLeft
	KeyBar
	KeyBraceRight
	KeyAsciiTilde
)
const (
	KeyBackspace = iota + 0xff08
	KeyTab
	KeyLinefeed
	KeyClear
	_
	KeyReturn
)
const (
	KeyPause      = 0xff13
	KeyScrollLock = 0xff14
	KeySysReq     = 0xff15
	KeyEscape     = 0xff1b
)
const (
	KeyHome = iota + 0xff50
	KeyLeft
	KeyUp
	KeyRight
	KeyDown
	KeyPageUp
	KeyPageDown
	KeyEnd
	KeyInsert = 0xff63
)
const (
	KeyF1 = iota + 0xffbe
	KeyF2
	KeyF3
	KeyF4
	KeyF5
	KeyF6
	KeyF7
	KeyF8
	KeyF9
	KeyF10
	KeyF11
	KeyF12
)
const (
	KeyShiftLeft = iota + 0xffe1
	KeyShiftRight
	KeyControlLeft
	KeyControlRight
	KeyCapsLock
	_
	_
	_
	KeyAltLeft
	KeyAltRight
	KeyDelete = 0xffff
)
<|endoftext|>


<commit_before>package gremgo

import (
	"io/ioutil"
	"sync"
)

// Client is a container for the gremgo client.
type Client struct {
	conn             dialer
	requests         chan []byte
	responses        chan []byte
	results          map[string][]interface{}
	responseNotifyer map[string]chan int // responseNotifyer notifies the requester that a response has arrived for the request
	respMutex        *sync.Mutex
}

// NewDialer returns a WebSocket dialer to use when connecting to Gremlin Server
func NewDialer(host string) (dialer *Ws) {
	dialer = new(Ws)
	dialer.host = "ws://" + host
	return dialer
}

func newClient() (c Client) {
	c.requests = make(chan []byte, 3)  // c.requests takes any request and delivers it to the WriteWorker for dispatch to Gremlin Server
	c.responses = make(chan []byte, 3) // c.responses takes raw responses from ReadWorker and delivers it for sorting to handelResponse
	c.results = make(map[string][]interface{})
	c.responseNotifyer = make(map[string]chan int)
	c.respMutex = &sync.Mutex{} // c.mutex ensures that sorting is thread safe
	return
}

// Dial returns a gremgo client for interaction with the Gremlin Server specified in the host IP.
func Dial(conn dialer) (c Client, err error) {

	c = newClient()
	c.conn = conn

	// Connects to Gremlin Server
	err = conn.connect()
	if err != nil {
		return
	}

	go c.writeWorker()
	go c.readWorker()

	return
}

func (c *Client) executeRequest(query string, bindings map[string]string) (resp interface{}, err error) {
	req, id := prepareRequest(query, bindings)
	msg, err := packageRequest(req)
	if err != nil {
		return // TODO: Fix error handling
	}
	c.responseNotifyer[id] = make(chan int, 1)
	c.dispatchRequest(msg)
	resp = c.retrieveResponse(id)
	return
}

// Execute formats a raw Gremlin query, sends it to Gremlin Server, and returns the result.
func (c *Client) Execute(query string, bindings map[string]string) (resp interface{}, err error) {
	resp, err = c.executeRequest(query, bindings)
	return
}

// ExecuteFile takes a file path to a Gremlin script, sends it to Gremlin Server, and returns the result.
func (c *Client) ExecuteFile(path string, bindings map[string]string) (resp interface{}, err error) {
	d, err := ioutil.ReadFile(path) // Read script from file
	if err != nil {
		return
	}
	query := string(d)
	resp, err = c.executeRequest(query, bindings)
	return
}
<commit_msg>Added some logging for error events. Added close method to Client so that hte underlying socket connection can be closed if needed.<commit_after>package gremgo

import (
	"io/ioutil"
	"log"
	"sync"
)

// Client is a container for the gremgo client.
type Client struct {
	conn             dialer
	requests         chan []byte
	responses        chan []byte
	results          map[string][]interface{}
	responseNotifyer map[string]chan int // responseNotifyer notifies the requester that a response has arrived for the request
	respMutex        *sync.Mutex
	Errored          bool
}

// NewDialer returns a WebSocket dialer to use when connecting to Gremlin Server
func NewDialer(host string) (dialer *Ws) {
	dialer = new(Ws)
	dialer.host = "ws://" + host
	return dialer
}

func newClient() (c Client) {
	c.requests = make(chan []byte, 3)  // c.requests takes any request and delivers it to the WriteWorker for dispatch to Gremlin Server
	c.responses = make(chan []byte, 3) // c.responses takes raw responses from ReadWorker and delivers it for sorting to handelResponse
	c.results = make(map[string][]interface{})
	c.responseNotifyer = make(map[string]chan int)
	c.respMutex = &sync.Mutex{} // c.mutex ensures that sorting is thread safe
	return
}

// Dial returns a gremgo client for interaction with the Gremlin Server specified in the host IP.
func Dial(conn dialer) (c Client, err error) {

	c = newClient()
	c.conn = conn

	// Connects to Gremlin Server
	err = conn.connect()
	if err != nil {
		return
	}

	go c.writeWorker()
	go c.readWorker()

	return
}

func (c *Client) executeRequest(query string, bindings map[string]string) (resp interface{}, err error) {
	req, id := prepareRequest(query, bindings)
	msg, err := packageRequest(req)
	if err != nil {
		log.Println(err)
		return
	}
	c.responseNotifyer[id] = make(chan int, 1)
	c.dispatchRequest(msg)
	resp = c.retrieveResponse(id)
	return
}

// Execute formats a raw Gremlin query, sends it to Gremlin Server, and returns the result.
func (c *Client) Execute(query string, bindings map[string]string) (resp interface{}, err error) {
	resp, err = c.executeRequest(query, bindings)
	return
}

// ExecuteFile takes a file path to a Gremlin script, sends it to Gremlin Server, and returns the result.
func (c *Client) ExecuteFile(path string, bindings map[string]string) (resp interface{}, err error) {
	d, err := ioutil.ReadFile(path) // Read script from file
	if err != nil {
		log.Println(err)
		return
	}
	query := string(d)
	resp, err = c.executeRequest(query, bindings)
	return
}

// Close closes the underlying connection and marks the client as closed.
func (c *Client) Close() {
	if c.conn != nil {
		c.conn.close()
	}
}
<|endoftext|>