<commit_before>package mdns

import (
	"fmt"
	"log"
	"net"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
	"github.com/miekg/dns"
)

// ServiceEntry is returned after we query for a service
type ServiceEntry struct {
	Name   string
	Host   string
	AddrV4 net.IP
	AddrV6 net.IP
	Port   int
	Info   string

	Addr net.IP // @Deprecated

	hasTXT bool
	sent   bool
}

// complete is used to check if we have all the info we need
func (s *ServiceEntry) complete() bool {
	return (s.AddrV4 != nil || s.AddrV6 != nil || s.Addr != nil) && s.Port != 0 && s.hasTXT
}

// QueryParam is used to customize how a Lookup is performed
type QueryParam struct {
	Service             string               // Service to lookup
	Domain              string               // Lookup domain, default "local"
	Timeout             time.Duration        // Lookup timeout, default 1 second
	Interface           *net.Interface       // Multicast interface to use
	Entries             chan<- *ServiceEntry // Entries Channel
	WantUnicastResponse bool                 // Unicast response desired, as per 5.4 in RFC
}

// DefaultParams is used to return a default set of QueryParam's
func DefaultParams(service string) *QueryParam {
	return &QueryParam{
		Service:             service,
		Domain:              "local",
		Timeout:             time.Second,
		Entries:             make(chan *ServiceEntry),
		WantUnicastResponse: false, // TODO(reddaly): Change this default.
	}
}

// Query looks up a given service, in a domain, waiting at most
// for a timeout before finishing the query. The results are streamed
// to a channel. Sends will not block, so clients should make sure to
// either read or buffer.
func Query(params *QueryParam) error {
	// Create a new client
	client, err := newClient()
	if err != nil {
		return err
	}
	defer client.Close()

	// Set the multicast interface
	if params.Interface != nil {
		if err := client.setInterface(params.Interface); err != nil {
			return err
		}
	}

	// Ensure defaults are set
	if params.Domain == "" {
		params.Domain = "local"
	}
	if params.Timeout == 0 {
		params.Timeout = time.Second
	}

	// Run the query
	return client.query(params)
}

// Lookup is the same as Query, however it uses all the default parameters
func Lookup(service string, entries chan<- *ServiceEntry) error {
	params := DefaultParams(service)
	params.Entries = entries
	return Query(params)
}

// Client provides a query interface that can be used to
// search for service providers using mDNS
type client struct {
	ipv4UnicastConn *net.UDPConn
	ipv6UnicastConn *net.UDPConn

	ipv4MulticastConn *net.UDPConn
	ipv6MulticastConn *net.UDPConn

	closed    bool
	closedCh  chan struct{} // TODO(reddaly): This doesn't appear to be used.
	closeLock sync.Mutex
}

// NewClient creates a new mdns Client that can be used to query
// for records
func newClient() (*client, error) {
	// TODO(reddaly): At least attempt to bind to the port required in the spec.
	// Create a IPv4 listener
	uconn4, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
	if err != nil {
		//log.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
	}
	uconn6, err := net.ListenUDP("udp6", &net.UDPAddr{IP: net.IPv6zero, Port: 0})
	if err != nil {
		//log.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
	}

	if uconn4 == nil && uconn6 == nil {
		return nil, fmt.Errorf("failed to bind to any unicast udp port")
	}

	mconn4, err := net.ListenMulticastUDP("udp4", nil, ipv4Addr)
	if err != nil {
		//log.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
	}
	mconn6, err := net.ListenMulticastUDP("udp6", nil, ipv6Addr)
	if err != nil {
		//log.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
	}

	if mconn4 == nil && mconn6 == nil {
		return nil, fmt.Errorf("failed to bind to any multicast udp port")
	}

	c := &client{
		ipv4MulticastConn: mconn4,
		ipv6MulticastConn: mconn6,
		ipv4UnicastConn:   uconn4,
		ipv6UnicastConn:   uconn6,
		closedCh:          make(chan struct{}),
	}
	return c, nil
}

// Close is used to cleanup the client
func (c *client) Close() error {
	c.closeLock.Lock()
	defer c.closeLock.Unlock()

	if c.closed {
		return nil
	}
	c.closed = true

	close(c.closedCh)

	if c.ipv4UnicastConn != nil {
		c.ipv4UnicastConn.Close()
	}
	if c.ipv6UnicastConn != nil {
		c.ipv6UnicastConn.Close()
	}
	if c.ipv4MulticastConn != nil {
		c.ipv4MulticastConn.Close()
	}
	if c.ipv6MulticastConn != nil {
		c.ipv6MulticastConn.Close()
	}

	return nil
}

// setInterface is used to set the query interface, uses sytem
// default if not provided
func (c *client) setInterface(iface *net.Interface) error {
	p := ipv4.NewPacketConn(c.ipv4UnicastConn)
	if err := p.SetMulticastInterface(iface); err != nil {
		return err
	}
	p2 := ipv6.NewPacketConn(c.ipv6UnicastConn)
	if err := p2.SetMulticastInterface(iface); err != nil {
		return err
	}
	p = ipv4.NewPacketConn(c.ipv4MulticastConn)
	if err := p.SetMulticastInterface(iface); err != nil {
		return err
	}
	p2 = ipv6.NewPacketConn(c.ipv6MulticastConn)
	if err := p2.SetMulticastInterface(iface); err != nil {
		return err
	}
	return nil
}

// query is used to perform a lookup and stream results
func (c *client) query(params *QueryParam) error {
	// Create the service name
	serviceAddr := fmt.Sprintf("%s.%s.", trimDot(params.Service), trimDot(params.Domain))

	// Start listening for response packets
	msgCh := make(chan *dns.Msg, 32)
	go c.recv(c.ipv4UnicastConn, msgCh)
	go c.recv(c.ipv6UnicastConn, msgCh)
	go c.recv(c.ipv4MulticastConn, msgCh)
	go c.recv(c.ipv6MulticastConn, msgCh)

	// Send the query
	m := new(dns.Msg)
	m.SetQuestion(serviceAddr, dns.TypePTR)
	// RFC 6762, section 18.12.  Repurposing of Top Bit of qclass in Question
	// Section
	//
	// In the Question Section of a Multicast DNS query, the top bit of the qclass
	// field is used to indicate that unicast responses are preferred for this
	// particular question.  (See Section 5.4.)
	if params.WantUnicastResponse {
		m.Question[0].Qclass |= 1 << 15
	}
	m.RecursionDesired = false
	if err := c.sendQuery(m); err != nil {
		return err
	}

	// Map the in-progress responses
	inprogress := make(map[string]*ServiceEntry)

	// Listen until we reach the timeout
	finish := time.After(params.Timeout)
	for {
		select {
		case resp := <-msgCh:
			var inp *ServiceEntry
			for _, answer := range append(resp.Answer, resp.Extra...) {
				// TODO(reddaly): Check that response corresponds to serviceAddr?
				switch rr := answer.(type) {
				case *dns.PTR:
					// Create new entry for this
					inp = ensureName(inprogress, rr.Ptr)

				case *dns.SRV:
					// Check for a target mismatch
					if rr.Target != rr.Hdr.Name {
						alias(inprogress, rr.Hdr.Name, rr.Target)
					}

					// Get the port
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Host = rr.Target
					inp.Port = int(rr.Port)

				case *dns.TXT:
					// Pull out the txt
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Info = strings.Join(rr.Txt, "|")
					inp.hasTXT = true

				case *dns.A:
					// Pull out the IP
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Addr = rr.A // @Deprecated
					inp.AddrV4 = rr.A

				case *dns.AAAA:
					// Pull out the IP
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Addr = rr.AAAA // @Deprecated
					inp.AddrV6 = rr.AAAA
				}
			}

			if inp == nil {
				continue
			}

			// Check if this entry is complete
			if inp.complete() {
				if inp.sent {
					continue
				}
				inp.sent = true
				select {
				case params.Entries <- inp:
				default:
				}
			} else {
				// Fire off a node specific query
				m := new(dns.Msg)
				m.SetQuestion(inp.Name, dns.TypePTR)
				m.RecursionDesired = false
				if err := c.sendQuery(m); err != nil {
					log.Printf("[ERR] mdns: Failed to query instance %s: %v", inp.Name, err)
				}
			}
		case <-finish:
			return nil
		}
	}
}

// sendQuery is used to multicast a query out
func (c *client) sendQuery(q *dns.Msg) error {
	buf, err := q.Pack()
	if err != nil {
		return err
	}
	if c.ipv4UnicastConn != nil {
		c.ipv4UnicastConn.WriteToUDP(buf, ipv4Addr)
	}
	if c.ipv6UnicastConn != nil {
		c.ipv6UnicastConn.WriteToUDP(buf, ipv6Addr)
	}
	return nil
}

// recv is used to receive until we get a shutdown
func (c *client) recv(l *net.UDPConn, msgCh chan *dns.Msg) {
	if l == nil {
		return
	}
	buf := make([]byte, 65536)
	for !c.closed {
		n, err := l.Read(buf)
		if err != nil {
			log.Printf("[ERR] mdns: Failed to read packet: %v", err)
			continue
		}
		msg := new(dns.Msg)
		if err := msg.Unpack(buf[:n]); err != nil {
			log.Printf("[ERR] mdns: Failed to unpack packet: %v", err)
			continue
		}
		select {
		case msgCh <- msg:
		case <-c.closedCh:
			return
		}
	}
}

// ensureName is used to ensure the named node is in progress
func ensureName(inprogress map[string]*ServiceEntry, name string) *ServiceEntry {
	if inp, ok := inprogress[name]; ok {
		return inp
	}
	inp := &ServiceEntry{
		Name: name,
	}
	inprogress[name] = inp
	return inp
}

// alias is used to setup an alias between two entries
func alias(inprogress map[string]*ServiceEntry, src, dst string) {
	srcEntry := ensureName(inprogress, src)
	inprogress[dst] = srcEntry
}
<commit_msg>Fix implementation of the Interface config parameter<commit_after>package mdns

import (
	"fmt"
	"log"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/miekg/dns"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

// ServiceEntry is returned after we query for a service
type ServiceEntry struct {
	Name   string
	Host   string
	AddrV4 net.IP
	AddrV6 net.IP
	Port   int
	Info   string

	Addr net.IP // @Deprecated

	hasTXT bool
	sent   bool
}

// complete is used to check if we have all the info we need
func (s *ServiceEntry) complete() bool {
	return (s.AddrV4 != nil || s.AddrV6 != nil || s.Addr != nil) && s.Port != 0 && s.hasTXT
}

// QueryParam is used to customize how a Lookup is performed
type QueryParam struct {
	Service             string               // Service to lookup
	Domain              string               // Lookup domain, default "local"
	Timeout             time.Duration        // Lookup timeout, default 1 second
	Interface           *net.Interface       // Multicast interface to use
	Entries             chan<- *ServiceEntry // Entries Channel
	WantUnicastResponse bool                 // Unicast response desired, as per 5.4 in RFC
}

// DefaultParams is used to return a default set of QueryParam's
func DefaultParams(service string) *QueryParam {
	return &QueryParam{
		Service:             service,
		Domain:              "local",
		Timeout:             time.Second,
		Entries:             make(chan *ServiceEntry),
		WantUnicastResponse: false, // TODO(reddaly): Change this default.
	}
}

// Query looks up a given service, in a domain, waiting at most
// for a timeout before finishing the query. The results are streamed
// to a channel. Sends will not block, so clients should make sure to
// either read or buffer.
func Query(params *QueryParam) error {
	// Create a new client
	client, err := newClient()
	if err != nil {
		return err
	}
	defer client.Close()

	// Ensure defaults are set
	if params.Domain == "" {
		params.Domain = "local"
	}
	if params.Timeout == 0 {
		params.Timeout = time.Second
	}

	// Run the query
	return client.query(params)
}

// Lookup is the same as Query, however it uses all the default parameters
func Lookup(service string, entries chan<- *ServiceEntry) error {
	params := DefaultParams(service)
	params.Entries = entries
	return Query(params)
}

// Client provides a query interface that can be used to
// search for service providers using mDNS
type client struct {
	ipv4UnicastConn *net.UDPConn
	ipv6UnicastConn *net.UDPConn

	ipv4MulticastConn *net.UDPConn
	ipv6MulticastConn *net.UDPConn

	closed    bool
	closedCh  chan struct{} // TODO(reddaly): This doesn't appear to be used.
	closeLock sync.Mutex
}

// NewClient creates a new mdns Client that can be used to query
// for records
func newClient() (*client, error) {
	// TODO(reddaly): At least attempt to bind to the port required in the spec.
	// Create a IPv4 listener
	uconn4, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
	if err != nil {
		//log.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
	}
	uconn6, err := net.ListenUDP("udp6", &net.UDPAddr{IP: net.IPv6zero, Port: 0})
	if err != nil {
		//log.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
	}

	if uconn4 == nil && uconn6 == nil {
		return nil, fmt.Errorf("failed to bind to any unicast udp port")
	}

	mconn4, err := net.ListenMulticastUDP("udp4", nil, ipv4Addr)
	if err != nil {
		//log.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
	}
	mconn6, err := net.ListenMulticastUDP("udp6", nil, ipv6Addr)
	if err != nil {
		//log.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
	}

	if mconn4 == nil && mconn6 == nil {
		return nil, fmt.Errorf("failed to bind to any multicast udp port")
	}

	c := &client{
		ipv4MulticastConn: mconn4,
		ipv6MulticastConn: mconn6,
		ipv4UnicastConn:   uconn4,
		ipv6UnicastConn:   uconn6,
		closedCh:          make(chan struct{}),
	}
	return c, nil
}

// Close is used to cleanup the client
func (c *client) Close() error {
	c.closeLock.Lock()
	defer c.closeLock.Unlock()

	if c.closed {
		return nil
	}
	c.closed = true

	close(c.closedCh)

	if c.ipv4UnicastConn != nil {
		c.ipv4UnicastConn.Close()
	}
	if c.ipv6UnicastConn != nil {
		c.ipv6UnicastConn.Close()
	}
	if c.ipv4MulticastConn != nil {
		c.ipv4MulticastConn.Close()
	}
	if c.ipv6MulticastConn != nil {
		c.ipv6MulticastConn.Close()
	}

	return nil
}

// setInterface is used to set the query interface, uses sytem
// default if not provided
func (c *client) setInterface(iface *net.Interface) error {
	p := ipv4.NewPacketConn(c.ipv4UnicastConn)
	if err := p.SetMulticastInterface(iface); err != nil {
		return err
	}
	p2 := ipv6.NewPacketConn(c.ipv6UnicastConn)
	if err := p2.SetMulticastInterface(iface); err != nil {
		return err
	}
	p = ipv4.NewPacketConn(c.ipv4MulticastConn)
	if err := p.SetMulticastInterface(iface); err != nil {
		return err
	}
	p2 = ipv6.NewPacketConn(c.ipv6MulticastConn)
	if err := p2.SetMulticastInterface(iface); err != nil {
		return err
	}
	return nil
}

// query is used to perform a lookup and stream results
func (c *client) query(params *QueryParam) error {
	// Create the service name
	serviceAddr := fmt.Sprintf("%s.%s.", trimDot(params.Service), trimDot(params.Domain))

	// Start listening for response packets
	msgCh := make(chan *dns.Msg, 32)
	go c.recv(c.ipv4UnicastConn, msgCh)
	go c.recv(c.ipv6UnicastConn, msgCh)
	go c.recv(c.ipv4MulticastConn, msgCh)
	go c.recv(c.ipv6MulticastConn, msgCh)

	// Send the query
	m := new(dns.Msg)
	m.SetQuestion(serviceAddr, dns.TypePTR)
	// RFC 6762, section 18.12.  Repurposing of Top Bit of qclass in Question
	// Section
	//
	// In the Question Section of a Multicast DNS query, the top bit of the qclass
	// field is used to indicate that unicast responses are preferred for this
	// particular question.  (See Section 5.4.)
	if params.WantUnicastResponse {
		m.Question[0].Qclass |= 1 << 15
	}
	m.RecursionDesired = false
	if err := c.sendQuery(m, params.Interface); err != nil {
		return err
	}

	// Map the in-progress responses
	inprogress := make(map[string]*ServiceEntry)

	// Listen until we reach the timeout
	finish := time.After(params.Timeout)
	for {
		select {
		case resp := <-msgCh:
			var inp *ServiceEntry
			for _, answer := range append(resp.Answer, resp.Extra...) {
				// TODO(reddaly): Check that response corresponds to serviceAddr?
				switch rr := answer.(type) {
				case *dns.PTR:
					// Create new entry for this
					inp = ensureName(inprogress, rr.Ptr)

				case *dns.SRV:
					// Check for a target mismatch
					if rr.Target != rr.Hdr.Name {
						alias(inprogress, rr.Hdr.Name, rr.Target)
					}

					// Get the port
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Host = rr.Target
					inp.Port = int(rr.Port)

				case *dns.TXT:
					// Pull out the txt
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Info = strings.Join(rr.Txt, "|")
					inp.hasTXT = true

				case *dns.A:
					// Pull out the IP
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Addr = rr.A // @Deprecated
					inp.AddrV4 = rr.A

				case *dns.AAAA:
					// Pull out the IP
					inp = ensureName(inprogress, rr.Hdr.Name)
					inp.Addr = rr.AAAA // @Deprecated
					inp.AddrV6 = rr.AAAA
				}
			}

			if inp == nil {
				continue
			}

			// Check if this entry is complete
			if inp.complete() {
				if inp.sent {
					continue
				}
				inp.sent = true
				select {
				case params.Entries <- inp:
				default:
				}
			} else {
				// Fire off a node specific query
				m := new(dns.Msg)
				m.SetQuestion(inp.Name, dns.TypePTR)
				m.RecursionDesired = false
				if err := c.sendQuery(m, params.Interface); err != nil {
					log.Printf("[ERR] mdns: Failed to query instance %s: %v", inp.Name, err)
				}
			}
		case <-finish:
			return nil
		}
	}
}

// sendQuery is used to multicast a query out
func (c *client) sendQuery(q *dns.Msg, iface *net.Interface) error {
	buf, err := q.Pack()
	if err != nil {
		return err
	}
	if c.ipv4UnicastConn != nil {
		p := ipv4.NewPacketConn(c.ipv4UnicastConn)
		if iface != nil {
			p.SetMulticastInterface(iface)
		}
		p.WriteTo(buf, nil, ipv4Addr)
	}
	if c.ipv6UnicastConn != nil {
		p := ipv6.NewPacketConn(c.ipv6UnicastConn)
		if iface != nil {
			p.SetMulticastInterface(iface)
		}
		p.WriteTo(buf, nil, ipv6Addr)
	}
	return nil
}

// recv is used to receive until we get a shutdown
func (c *client) recv(l *net.UDPConn, msgCh chan *dns.Msg) {
	if l == nil {
		return
	}
	buf := make([]byte, 65536)
	for !c.closed {
		n, err := l.Read(buf)
		if err != nil {
			// log.Printf("[ERR] mdns: Failed to read packet: %v", err)
			continue
		}
		msg := new(dns.Msg)
		if err := msg.Unpack(buf[:n]); err != nil {
			log.Printf("[ERR] mdns: Failed to unpack packet: %v", err)
			continue
		}
		select {
		case msgCh <- msg:
		case <-c.closedCh:
			return
		}
	}
}

// ensureName is used to ensure the named node is in progress
func ensureName(inprogress map[string]*ServiceEntry, name string) *ServiceEntry {
	if inp, ok := inprogress[name]; ok {
		return inp
	}
	inp := &ServiceEntry{
		Name: name,
	}
	inprogress[name] = inp
	return inp
}

// alias is used to setup an alias between two entries
func alias(inprogress map[string]*ServiceEntry, src, dst string) {
	srcEntry := ensureName(inprogress, src)
	inprogress[dst] = srcEntry
}
<|endoftext|>


<commit_before>package prisclient

import (
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/priscillachat/prislog"
	"io"
	"net"
	"time"
)

type Client struct {
	raw      net.Conn
	decoder  *json.Decoder
	encoder  *json.Encoder
	SourceId string
	Logger   *prislog.PrisLog
}

type CommandBlock struct {
	Id      string   `json:"id,omitempty"`
	Action  string   `json:"action,omitempty"`
	Type    string   `json:"type,omitempty"`
	Time    int64    `json:"time,omitempty"`
	Data    string   `json:"data,omitempty"`
	Array   []string `json:"array,omitempty"`
	Options []string `json:"options,omitempty"`
}

type MessageBlock struct {
	Message       string   `json:"message,omitempty"`
	From          string   `json:"from,omitempty"`
	Room          string   `json:"room,omitempty"`
	Mentioned     bool     `json:"mentioned,omitempty"`
	Stripped      string   `json:"stripped,omitempty"`
	MentionNotify []string `json:"mentionnotify,omitempty"`
}

type Query struct {
	Type    string        `json:"type,omitempty"`
	Source  string        `json:"source,omitempty"`
	To      string        `json:"to,omitempty"`
	Command *CommandBlock `json:"command,omitempty"`
	Message *MessageBlock `json:"message,omitempty"`
}

func RandomId() string {
	b := make([]byte, 8)
	io.ReadFull(rand.Reader, b)
	return fmt.Sprintf("%x", b)
}

func NewClient(host, port string,
	logger *prislog.PrisLog) (*Client, error) {

	pris := new(Client)

	pris.Logger = logger

	conn, err := net.Dial("tcp", host+":"+port)

	if err != nil {
		return pris, err
	}

	pris.raw = conn
	pris.decoder = json.NewDecoder(pris.raw)
	pris.encoder = json.NewEncoder(pris.raw)

	return pris, nil
}

func (pris *Client) Engage(clienttype, sourceid string) error {
	if clienttype != "adapter" && clienttype != "responder" {
		return errors.New("client type has to be adapter or responder")
	}

	q := Query{
		Type:   "command",
		Source: sourceid,
		Command: &CommandBlock{
			Id:     RandomId(),
			Action: "engage",
			Type:   clienttype,
			Time:   time.Now().Unix(),
		},
	}

	pris.SourceId = sourceid

	return pris.encoder.Encode(&q)
}

func (pris *Client) listen(out chan<- *Query) {
	var q *Query
	for {
		q = new(Query)
		err := pris.decoder.Decode(q)

		if err != nil {
			fmt.Println(err)
			if err.Error() == "EOF" {
				out <- &Query{
					Type:   "command",
					Source: "pris",
					Command: &CommandBlock{
						Action: "disengage",
					},
				}
				break
			}
		}

		if pris.ValidateQuery(q) {
			pris.Logger.Debug.Println("Query received:", *q)
			if q.Type == "message" {
				pris.Logger.Debug.Println("Message:", q.Message.Message)
				pris.Logger.Debug.Println("From:", q.Message.From)
				pris.Logger.Debug.Println("Room:", q.Message.Room)
			}
			out <- q
		} else {
			pris.Logger.Error.Println("Invalid query from server(?!!):", q)
		}
	}
}

func (pris *Client) Run(toPris <-chan *Query, fromPris chan<- *Query) {

	go pris.listen(fromPris)
	var q *Query
	for {
		q = <-toPris
		if pris.ValidateQuery(q) {
			pris.encoder.Encode(q)
		} else {
			pris.Logger.Error.Println("Invalid query:", q)
		}
	}
}

func (pris *Client) ValidateQuery(q *Query) bool {
	switch {
	case q.Type == "command" && q.Command != nil:
		return true
	case q.Type == "message" && q.Message != nil && q.Message.Room != "":
		return true
	}
	return false
}
<commit_msg>don't litter variables all over the place<commit_after>package prisclient

import (
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/priscillachat/prislog"
	"io"
	"net"
	"time"
)

type Client struct {
	raw      net.Conn
	decoder  *json.Decoder
	encoder  *json.Encoder
	SourceId string
	Type     string
	Logger   *prislog.PrisLog
}

type CommandBlock struct {
	Id      string   `json:"id,omitempty"`
	Action  string   `json:"action,omitempty"`
	Type    string   `json:"type,omitempty"`
	Time    int64    `json:"time,omitempty"`
	Data    string   `json:"data,omitempty"`
	Array   []string `json:"array,omitempty"`
	Options []string `json:"options,omitempty"`
}

type MessageBlock struct {
	Message       string   `json:"message,omitempty"`
	From          string   `json:"from,omitempty"`
	Room          string   `json:"room,omitempty"`
	Mentioned     bool     `json:"mentioned,omitempty"`
	Stripped      string   `json:"stripped,omitempty"`
	MentionNotify []string `json:"mentionnotify,omitempty"`
}

type Query struct {
	Type    string        `json:"type,omitempty"`
	Source  string        `json:"source,omitempty"`
	To      string        `json:"to,omitempty"`
	Command *CommandBlock `json:"command,omitempty"`
	Message *MessageBlock `json:"message,omitempty"`
}

func RandomId() string {
	b := make([]byte, 8)
	io.ReadFull(rand.Reader, b)
	return fmt.Sprintf("%x", b)
}

func NewClient(host, port, clientType, sourceid string,
	logger *prislog.PrisLog) (*Client, error) {

	if clientType != "adapter" && clientType != "responder" {
		return nil, errors.New("client type has to be adapter or responder")
	}

	pris := new(Client)

	pris.Logger = logger
	pris.SourceId = sourceid
	pris.Type = clientType

	conn, err := net.Dial("tcp", host+":"+port)

	if err != nil {
		return pris, err
	}

	pris.raw = conn
	pris.decoder = json.NewDecoder(pris.raw)
	pris.encoder = json.NewEncoder(pris.raw)

	return pris, nil
}

func (pris *Client) Engage() error {
	q := Query{
		Type:   "command",
		Source: pris.SourceId,
		Command: &CommandBlock{
			Id:     RandomId(),
			Action: "engage",
			Type:   pris.Type,
			Time:   time.Now().Unix(),
		},
	}

	return pris.encoder.Encode(&q)
}

func (pris *Client) listen(out chan<- *Query) {
	var q *Query
	for {
		q = new(Query)
		err := pris.decoder.Decode(q)

		if err != nil {
			fmt.Println(err)
			if err.Error() == "EOF" {
				out <- &Query{
					Type:   "command",
					Source: "pris",
					Command: &CommandBlock{
						Action: "disengage",
					},
				}
				break
			}
		}

		if pris.ValidateQuery(q) {
			pris.Logger.Debug.Println("Query received:", *q)
			if q.Type == "message" {
				pris.Logger.Debug.Println("Message:", q.Message.Message)
				pris.Logger.Debug.Println("From:", q.Message.From)
				pris.Logger.Debug.Println("Room:", q.Message.Room)
			}
			out <- q
		} else {
			pris.Logger.Error.Println("Invalid query from server(?!!):", q)
		}
	}
}

func (pris *Client) Run(toPris <-chan *Query, fromPris chan<- *Query) {

	go pris.listen(fromPris)
	var q *Query
	for {
		q = <-toPris
		if pris.ValidateQuery(q) {
			pris.encoder.Encode(q)
		} else {
			pris.Logger.Error.Println("Invalid query:", q)
		}
	}
}

func (pris *Client) ValidateQuery(q *Query) bool {
	switch {
	case q.Type == "command" && q.Command != nil:
		return true
	case q.Type == "message" && q.Message != nil && q.Message.Room != "":
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package pusher

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	u "net/url"
	"os"
	"regexp"
	"strings"
	"time"
)

var pusherPathRegex = regexp.MustCompile("^/apps/([0-9]+)$")

// Client to the HTTP API of Pusher
type Client struct {
	AppId   string
	Key     string
	Secret  string
	Host    string        // host or host:port pair
	Secure  bool          // true for HTTPS
	Timeout time.Duration // Request timeout for HTTP requests
}

/*
Clients can be instantiated from a specially-crafted Pusher URL. For example:

	client := pusher.ClientFromURL("http://key:secret@api.pusherapp.com/apps/app_id")

*/
func ClientFromURL(url string) (*Client, error) {
	url2, err := u.Parse(url)
	if err != nil {
		return nil, err
	}

	c := Client{
		Host: url2.Host,
	}

	matches := pusherPathRegex.FindStringSubmatch(url2.Path)
	if len(matches) == 0 {
		return nil, errors.New("No app ID found")
	}
	c.AppId = matches[1]

	if url2.User == nil {
		return nil, errors.New("Missing <key>:<secret>")
	}
	c.Key = url2.User.Username()
	var isSet bool
	c.Secret, isSet = url2.User.Password()
	if !isSet {
		return nil, errors.New("Missing <secret>")
	}

	if url2.Scheme == "https" {
		c.Secure = true
	}

	return &c, nil
}

/*
Alternatively, you can instantiate a client from an environment variable.
This is particularly relevant if you are using Pusher as a Heroku add-on,
which stores credentials in a `"PUSHER_URL"` environment variable. For example:
	client := pusher.ClientFromEnv("PUSHER_URL")

*/
func ClientFromEnv(key string) (*Client, error) {
	url := os.Getenv(key)
	return ClientFromURL(url)
}

// triggerMulti
// socketId pointer to string

func (c *Client) Trigger(channel string, event string, data interface{}) (*BufferedEvents, error) {
	return c.trigger([]string{channel}, event, data, nil)
}

func (c *Client) TriggerMulti(channels []string, event string, data interface{}) (*BufferedEvents, error) {
	return c.trigger(channels, event, data, nil)
}

func (c *Client) TriggerExclusive(channel string, event string, data interface{}, socketID string) (*BufferedEvents, error) {
	return c.trigger([]string{channel}, event, data, &socketID)
}

func (c *Client) TriggerMultiExclusive(channels []string, event string, data interface{}, socketID string) (*BufferedEvents, error) {
	return c.trigger(channels, event, data, &socketID)
}

func (c *Client) trigger(channels []string, event string, data interface{}, socketId *string) (*BufferedEvents, error) {
	if len(channels) > 10 {
		return nil, errors.New("You cannot trigger on more than 10 channels at once")
	}

	if !channelsAreValid(channels) {
		return nil, errors.New("At least one of your channels' names are invalid")
	}

	payload, err := createTriggerPayload(channels, event, data, socketId)
	if err != nil {
		return nil, err
	}

	path := fmt.Sprintf("/apps/%s/events", c.AppId)
	u := createRequestUrl("POST", c.Host, path, c.Key, c.Secret, authTimestamp(), c.Secure, payload, nil)
	response, err := request("POST", u, payload, c.Timeout)
	if err != nil {
		return nil, err
	}

	return unmarshalledBufferedEvents(response)
}

func (c *Client) Channels(additionalQueries map[string]string) (*ChannelsList, error) {
	path := fmt.Sprintf("/apps/%s/channels", c.AppId)
	u := createRequestUrl("GET", c.Host, path, c.Key, c.Secret, authTimestamp(), c.Secure, nil, additionalQueries)
	response, err := request("GET", u, nil, c.Timeout)
	if err != nil {
		return nil, err
	}
	return unmarshalledChannelsList(response)
}

func (c *Client) Channel(name string, additionalQueries map[string]string) (*Channel, error) {
	path := fmt.Sprintf("/apps/%s/channels/%s", c.AppId, name)
	u := createRequestUrl("GET", c.Host, path, c.Key, c.Secret, authTimestamp(), c.Secure, nil, additionalQueries)
	response, err := request("GET", u, nil, c.Timeout)
	if err != nil {
		return nil, err
	}
	return unmarshalledChannel(response, name)
}

func (c *Client) GetChannelUsers(name string) (*Users, error) {
	path := fmt.Sprintf("/apps/%s/channels/%s/users", c.AppId, name)
	u := createRequestUrl("GET", c.Host, path, c.Key, c.Secret, authTimestamp(), c.Secure, nil, nil)
	response, err := request("GET", u, nil, c.Timeout)
	if err != nil {
		return nil, err
	}
	return unmarshalledChannelUsers(response)
}

func (c *Client) AuthenticatePrivateChannel(params []byte) (response []byte, err error) {
	return c.authenticateChannel(params, nil)
}

func (c *Client) AuthenticatePresenceChannel(params []byte, member MemberData) (response []byte, err error) {
	return c.authenticateChannel(params, &member)
}

func (c *Client) authenticateChannel(params []byte, member *MemberData) (response []byte, err error) {
	channelName, socketId, err := parseAuthRequestParams(params)

	if err != nil {
		return
	}

	stringToSign := strings.Join([]string{socketId, channelName}, ":")

	var jsonUserData string

	if member != nil {
		var _jsonUserData []byte
		_jsonUserData, err = json.Marshal(member)
		if err != nil {
			return
		}

		jsonUserData = string(_jsonUserData)
		stringToSign = strings.Join([]string{stringToSign, jsonUserData}, ":")
	}

	_response := createAuthMap(c.Key, c.Secret, stringToSign)

	if member != nil {
		_response["channel_data"] = jsonUserData
	}

	response, err = json.Marshal(_response)
	return
}

func (c *Client) Webhook(header http.Header, body []byte) (*Webhook, error) {

	for _, token := range header["X-Pusher-Key"] {
		if token == c.Key && checkSignature(header.Get("X-Pusher-Signature"), string(body), c.Secret) {
			return unmarshalledWebhook(body)
		}
	}
	return nil, errors.New("Invalid webhook")
}
<commit_msg>Added more godoc<commit_after>package pusher

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	u "net/url"
	"os"
	"regexp"
	"strings"
	"time"
)

var pusherPathRegex = regexp.MustCompile("^/apps/([0-9]+)$")

// Client to the HTTP API of Pusher.
type Client struct {
	AppId   string
	Key     string
	Secret  string
	Host    string        // host or host:port pair
	Secure  bool          // true for HTTPS
	Timeout time.Duration // Request timeout for HTTP requests
}

/*
Clients can be instantiated from a specially-crafted Pusher URL. For example:

	client := pusher.ClientFromURL("http://key:secret@api.pusherapp.com/apps/app_id")

*/
func ClientFromURL(url string) (*Client, error) {
	url2, err := u.Parse(url)
	if err != nil {
		return nil, err
	}

	c := Client{
		Host: url2.Host,
	}

	matches := pusherPathRegex.FindStringSubmatch(url2.Path)
	if len(matches) == 0 {
		return nil, errors.New("No app ID found")
	}
	c.AppId = matches[1]

	if url2.User == nil {
		return nil, errors.New("Missing <key>:<secret>")
	}
	c.Key = url2.User.Username()
	var isSet bool
	c.Secret, isSet = url2.User.Password()
	if !isSet {
		return nil, errors.New("Missing <secret>")
	}

	if url2.Scheme == "https" {
		c.Secure = true
	}

	return &c, nil
}

/*
Alternatively, you can instantiate a client from an environment variable.
This is particularly relevant if you are using Pusher as a Heroku add-on,
which stores credentials in a `"PUSHER_URL"` environment variable. For example:
	client := pusher.ClientFromEnv("PUSHER_URL")

*/
func ClientFromEnv(key string) (*Client, error) {
	url := os.Getenv(key)
	return ClientFromURL(url)
}

/*
It is possible to trigger an event on one or more channels. Channel names can
contain only characters which are alphanumeric, `_` or `-`` and have
to be at most 200 characters long. Event name can be at most 200 characters long too.

Pass in the channel's name, the event's name, and a data payload. The data payload must
be marshallable into JSON.

	data := map[string]string{"hello": "world"}
	client.Trigger("greeting_channel", "say_hello", data)

*/
func (c *Client) Trigger(channel string, event string, data interface{}) (*BufferedEvents, error) {
	return c.trigger([]string{channel}, event, data, nil)
}

/*
The same as `client.Trigger`, except one passes in a slice of `channels` as the first parameter.
The maximum length of channels is 10.
	client.TriggerMulti([]string{"a_channel", "another_channel"}, "event", data)
*/
func (c *Client) TriggerMulti(channels []string, event string, data interface{}) (*BufferedEvents, error) {
	return c.trigger(channels, event, data, nil)
}

/*
This method allow you to exclude a recipient whose connection has that
`socket_id` from receiving the event. You can read more here: http://pusher.com/docs/duplicates.

	client.TriggerExclusive("a_channel", "event", data, "123.12")

*/
func (c *Client) TriggerExclusive(channel string, event string, data interface{}, socketID string) (*BufferedEvents, error) {
	return c.trigger([]string{channel}, event, data, &socketID)
}

/*
Excluding a recipient on a trigger to multiple channels.

	client.TriggerMultiExclusive([]string{"a_channel", "another_channel"}, "event", data, "123.12")

*/
func (c *Client) TriggerMultiExclusive(channels []string, event string, data interface{}, socketID string) (*BufferedEvents, error) {
	return c.trigger(channels, event, data, &socketID)
}

func (c *Client) trigger(channels []string, event string, data interface{}, socketId *string) (*BufferedEvents, error) {
	if len(channels) > 10 {
		return nil, errors.New("You cannot trigger on more than 10 channels at once")
	}

	if !channelsAreValid(channels) {
		return nil, errors.New("At least one of your channels' names are invalid")
	}

	payload, err := createTriggerPayload(channels, event, data, socketId)
	if err != nil {
		return nil, err
	}

	path := fmt.Sprintf("/apps/%s/events", c.AppId)
	u := createRequestUrl("POST", c.Host, path, c.Key, c.Secret, authTimestamp(), c.Secure, payload, nil)
	response, err := request("POST", u, payload, c.Timeout)
	if err != nil {
		return nil, err
	}

	return unmarshalledBufferedEvents(response)
}

/*
One can use this method to get a list of all the channels in an applicaiton.

The parameter `additionalQueries` is a map with query options. A key with `"filter_by_prefix"`
will filter the returned channels. To get number of users subscribed to a presence-channel,
specify an `"info"` key with value `"user_count"`. Pass in `nil` if you do not wish to specify any query attributes.

	channelsParams := map[string]string{
	    "filter_by_prefix": "presence-",
	    "info":             "user_count",
	}

	channels, err := client.Channels(channelsParams)

	//channels=> &{Channels:map[presence-chatroom:{UserCount:4} presence-notifications:{UserCount:31}  ]}

*/
func (c *Client) Channels(additionalQueries map[string]string) (*ChannelsList, error) {
	path := fmt.Sprintf("/apps/%s/channels", c.AppId)
	u := createRequestUrl("GET", c.Host, path, c.Key, c.Secret, authTimestamp(), c.Secure, nil, additionalQueries)
	response, err := request("GET", u, nil, c.Timeout)
	if err != nil {
		return nil, err
	}
	return unmarshalledChannelsList(response)
}

/*
Get the state of a single channel.

The parameter `additionalQueries` is a map with query options. An `"info"` key can
have comma-separated vales of `"user_count"`, for presence-channels,
and `"subscription_count"`, for all-channels. Note that the subscription
count is not allowed by default. Please contact us at http://support.pusher.com
if you wish to enable this. Pass in `nil` if you do not wish to specify any query attributes.

	channelParams := map[string]string{
		"info": "user_count,subscription_count",
	}

	channel, err := client.Channel("presence-chatroom", channelParams)

	//channel=> &{Name:presence-chatroom Occupied:true UserCount:42 SubscriptionCount:42}
*/
func (c *Client) Channel(name string, additionalQueries map[string]string) (*Channel, error) {
	path := fmt.Sprintf("/apps/%s/channels/%s", c.AppId, name)
	u := createRequestUrl("GET", c.Host, path, c.Key, c.Secret, authTimestamp(), c.Secure, nil, additionalQueries)
	response, err := request("GET", u, nil, c.Timeout)
	if err != nil {
		return nil, err
	}
	return unmarshalledChannel(response, name)
}

/*
Get a list of users in a presence-channel by passing to this method the channel name.

	users, err := client.GetChannelUsers("presence-chatroom")

	//users=> &{List:[{Id:13} {Id:90}]}

*/
func (c *Client) GetChannelUsers(name string) (*Users, error) {
	path := fmt.Sprintf("/apps/%s/channels/%s/users", c.AppId, name)
	u := createRequestUrl("GET", c.Host, path, c.Key, c.Secret, authTimestamp(), c.Secure, nil, nil)
	response, err := request("GET", u, nil, c.Timeout)
	if err != nil {
		return nil, err
	}
	return unmarshalledChannelUsers(response)
}

/*
Application security is very important so Pusher provides a mechanism
for authenticating a user’s access to a channel at the point of subscription.

This can be used both to restrict access to private channels,
and in the case of presence channels notify subscribers of
who else is also subscribed via presence events.

This library provides a mechanism for generating an
authentication signature to send back to the client
and authorize them.

For more information see our docs:http://pusher.com/docs/authenticating_users.

This is an example of authenticating a private-channel, using the built-in
Golang HTTP library to start a server.

	func pusherAuth(res http.ResponseWriter, req *http.Request) {

		params, _ := ioutil.ReadAll(req.Body)
		response, err := client.AuthenticatePrivateChannel(params)

		if err != nil {
			panic(err)
		}

		fmt.Fprintf(res, string(response))

	}

	func main() {
		http.HandleFunc("/pusher/auth", pusherAuth)
		http.ListenAndServe(":5000", nil)
	}

*/
func (c *Client) AuthenticatePrivateChannel(params []byte) (response []byte, err error) {
	return c.authenticateChannel(params, nil)
}

/*
Using presence channels is similar to private channels,
but in order to identify a user, clients are
sent a user_id and, optionally, custom data.

	params, _ := ioutil.ReadAll(req.Body)

	presenceData := pusher.MemberData{
		UserId: "1",
		UserInfo: map[string]string{
			"twitter": "jamiepatel",
		},
	}

	response, err := client.AuthenticatePresenceChannel(params, presenceData)

	if err != nil {
		panic(err)
	}

	fmt.Fprintf(res, response)

*/
func (c *Client) AuthenticatePresenceChannel(params []byte, member MemberData) (response []byte, err error) {
	return c.authenticateChannel(params, &member)
}

func (c *Client) authenticateChannel(params []byte, member *MemberData) (response []byte, err error) {
	channelName, socketId, err := parseAuthRequestParams(params)

	if err != nil {
		return
	}

	stringToSign := strings.Join([]string{socketId, channelName}, ":")

	var jsonUserData string

	if member != nil {
		var _jsonUserData []byte
		_jsonUserData, err = json.Marshal(member)
		if err != nil {
			return
		}

		jsonUserData = string(_jsonUserData)
		stringToSign = strings.Join([]string{stringToSign, jsonUserData}, ":")
	}

	_response := createAuthMap(c.Key, c.Secret, stringToSign)

	if member != nil {
		_response["channel_data"] = jsonUserData
	}

	response, err = json.Marshal(_response)
	return
}

/*
On your [dashboard](http://app.pusher.com), you can set up webhooks to POST
a payload to your server after certain events. Such events include channels
being occupied or vacated, members being added or removed in presence-channels,
or after client-originated events. For more information see https://pusher.com/docs/webhooks.

This library provides a mechanism for checking that these POST
requests are indeed from Pusher, by checking the token and
authentication signature in the header of the request.

If the webhook is valid, a `*pusher.Webhook* will be returned, and the `err` value will be nil.
If it is invalid, the first return value will be nil, and an error will be passed.

	func pusherWebhook(res http.ResponseWriter, req *http.Request) {

		body, _ := ioutil.ReadAll(req.Body)
		webhook, err := client.Webhook(req.Header, body)
	  	if err != nil {
	      fmt.Println("Webhook is invalid :(")
	  	} else {
	      fmt.Printf("%+v\n", webhook.Events)
	  	}

	}
*/
func (c *Client) Webhook(header http.Header, body []byte) (*Webhook, error) {

	for _, token := range header["X-Pusher-Key"] {
		if token == c.Key && checkSignature(header.Get("X-Pusher-Signature"), string(body), c.Secret) {
			return unmarshalledWebhook(body)
		}
	}
	return nil, errors.New("Invalid webhook")
}
<|endoftext|>


<commit_before>package whois

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"time"
)

const (
	// DefaultTimeout sets the maximum lifetime of whois requests.
	DefaultTimeout = 30 * time.Second

	// DefaultReadLimit sets the maximum bytes a client will attempt to read from a connection.
	DefaultReadLimit = 1 << 20 // 1 MB
)

// Client represents a whois client. It contains an http.Client, for executing
// some whois Requests.
type Client struct {
	Dial        func(string, string) (net.Conn, error)                  // Deprecated, use DialContext instead
	DialContext func(context.Context, string, string) (net.Conn, error) // Only used for port 43 (whois) requests, not HTTP(S)
	HTTPClient  *http.Client                                            // If nil, http.DefaultClient will be used
	Timeout     time.Duration                                           // Deprecated (use a Context instead)
}

// DefaultClient represents a shared whois client with a default timeout, HTTP
// transport, and dialer.
var DefaultClient = NewClient(DefaultTimeout)

// NewClient creates and initializes a new Client with the specified timeout.
func NewClient(timeout time.Duration) *Client {
	return &Client{
		Timeout: timeout,
	}
}

func (c *Client) dialContext(ctx context.Context, network, address string) (net.Conn, error) {
	var conn net.Conn
	var err error
	switch {
	case c.DialContext != nil:
		conn, err = c.DialContext(ctx, network, address)
	case c.Dial != nil:
		conn, err = c.Dial(network, address)
	default:
		conn, err = defaultDialer.DialContext(ctx, network, address)
	}
	if err != nil {
		return nil, err
	}
	if deadline, ok := ctx.Deadline(); ok {
		err = conn.SetDeadline(deadline)
	}
	return conn, err
}

var defaultDialer = &net.Dialer{}

// FetchError reports the underlying error and includes the target host of the fetch operation.
type FetchError struct {
	Err  error
	Host string
}

// Error implements the error interface.
func (f *FetchError) Error() string {
	return f.Err.Error()
}

// Fetch sends the Request to a whois server.
func (c *Client) Fetch(req *Request) (*Response, error) {
	return c.FetchContext(context.Background(), req)
}

// FetchContext sends the Request to a whois server.
// If ctx cancels or times out before the request completes, it will return an error.
func (c *Client) FetchContext(ctx context.Context, req *Request) (*Response, error) {
	if c.Timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(context.Background(), c.Timeout)
		defer cancel()
	}
	if req.URL != "" {
		return c.fetchHTTP(ctx, req)
	}
	return c.fetchWhois(ctx, req)
}

func (c *Client) fetchWhois(ctx context.Context, req *Request) (*Response, error) {
	if req.Host == "" {
		return nil, &FetchError{fmt.Errorf("no request host for %s", req.Query), "unknown"}
	}
	conn, err := c.dialContext(ctx, "tcp", req.Host+":43")
	if err != nil {
		return nil, &FetchError{err, req.Host}
	}
	defer conn.Close()
	if _, err = conn.Write(req.Body); err != nil {
		logError(err)
		return nil, &FetchError{err, req.Host}
	}
	res := NewResponse(req.Query, req.Host)
	if res.Body, err = ioutil.ReadAll(io.LimitReader(conn, DefaultReadLimit)); err != nil {
		logError(err)
		return nil, &FetchError{err, req.Host}
	}
	res.DetectContentType("")
	return res, nil
}

func (c *Client) fetchHTTP(ctx context.Context, req *Request) (*Response, error) {
	hreq, err := httpRequest(ctx, req)
	if err != nil {
		return nil, &FetchError{err, req.Host}
	}
	hc := c.HTTPClient
	if hc == nil {
		hc = http.DefaultClient
	}
	hres, err := hc.Do(hreq)
	if err != nil {
		return nil, &FetchError{err, req.Host}
	}
	res := NewResponse(req.Query, req.Host)
	if res.Body, err = ioutil.ReadAll(io.LimitReader(hres.Body, DefaultReadLimit)); err != nil {
		logError(err)
		return nil, &FetchError{err, req.Host}
	}
	res.DetectContentType(hres.Header.Get("Content-Type"))
	return res, nil
}

func httpRequest(ctx context.Context, req *Request) (*http.Request, error) {
	var hreq *http.Request
	var err error
	// POST if non-zero Request.Body
	if len(req.Body) > 0 {
		hreq, err = http.NewRequest("POST", req.URL, bytes.NewReader(req.Body))
	} else {
		hreq, err = http.NewRequest("GET", req.URL, nil)
	}
	if err != nil {
		return nil, err
	}
	// Some web whois servers require a Referer header
	hreq.Header.Add("Referer", req.URL)
	return hreq.WithContext(ctx), nil
}

func logError(err error) {
	switch t := err.(type) {
	case net.Error:
		fmt.Fprintf(os.Stderr, "net.Error timeout=%t, temp=%t: %s\n", t.Timeout(), t.Temporary(), err.Error())
	default:
		fmt.Fprintf(os.Stderr, "Unknown error %v: %s\n", t, err.Error())
	}
}
<commit_msg>use the passed Context (FIXME: add tests)<commit_after>package whois

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"time"
)

const (
	// DefaultTimeout sets the maximum lifetime of whois requests.
	DefaultTimeout = 30 * time.Second

	// DefaultReadLimit sets the maximum bytes a client will attempt to read from a connection.
	DefaultReadLimit = 1 << 20 // 1 MB
)

// Client represents a whois client. It contains an http.Client, for executing
// some whois Requests.
type Client struct {
	Dial        func(string, string) (net.Conn, error)                  // Deprecated, use DialContext instead
	DialContext func(context.Context, string, string) (net.Conn, error) // Only used for port 43 (whois) requests, not HTTP(S)
	HTTPClient  *http.Client                                            // If nil, http.DefaultClient will be used
	Timeout     time.Duration                                           // Deprecated (use a Context instead)
}

// DefaultClient represents a shared whois client with a default timeout, HTTP
// transport, and dialer.
var DefaultClient = NewClient(DefaultTimeout)

// NewClient creates and initializes a new Client with the specified timeout.
func NewClient(timeout time.Duration) *Client {
	return &Client{
		Timeout: timeout,
	}
}

func (c *Client) dialContext(ctx context.Context, network, address string) (net.Conn, error) {
	var conn net.Conn
	var err error
	switch {
	case c.DialContext != nil:
		conn, err = c.DialContext(ctx, network, address)
	case c.Dial != nil:
		conn, err = c.Dial(network, address)
	default:
		conn, err = defaultDialer.DialContext(ctx, network, address)
	}
	if err != nil {
		return nil, err
	}
	if deadline, ok := ctx.Deadline(); ok {
		err = conn.SetDeadline(deadline)
	}
	return conn, err
}

var defaultDialer = &net.Dialer{}

// FetchError reports the underlying error and includes the target host of the fetch operation.
type FetchError struct {
	Err  error
	Host string
}

// Error implements the error interface.
func (f *FetchError) Error() string {
	return f.Err.Error()
}

// Fetch sends the Request to a whois server.
func (c *Client) Fetch(req *Request) (*Response, error) {
	return c.FetchContext(context.Background(), req)
}

// FetchContext sends the Request to a whois server.
// If ctx cancels or times out before the request completes, it will return an error.
func (c *Client) FetchContext(ctx context.Context, req *Request) (*Response, error) {
	if c.Timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, c.Timeout)
		defer cancel()
	}
	if req.URL != "" {
		return c.fetchHTTP(ctx, req)
	}
	return c.fetchWhois(ctx, req)
}

func (c *Client) fetchWhois(ctx context.Context, req *Request) (*Response, error) {
	if req.Host == "" {
		return nil, &FetchError{fmt.Errorf("no request host for %s", req.Query), "unknown"}
	}
	conn, err := c.dialContext(ctx, "tcp", req.Host+":43")
	if err != nil {
		return nil, &FetchError{err, req.Host}
	}
	defer conn.Close()
	if _, err = conn.Write(req.Body); err != nil {
		logError(err)
		return nil, &FetchError{err, req.Host}
	}
	res := NewResponse(req.Query, req.Host)
	if res.Body, err = ioutil.ReadAll(io.LimitReader(conn, DefaultReadLimit)); err != nil {
		logError(err)
		return nil, &FetchError{err, req.Host}
	}
	res.DetectContentType("")
	return res, nil
}

func (c *Client) fetchHTTP(ctx context.Context, req *Request) (*Response, error) {
	hreq, err := httpRequest(ctx, req)
	if err != nil {
		return nil, &FetchError{err, req.Host}
	}
	hc := c.HTTPClient
	if hc == nil {
		hc = http.DefaultClient
	}
	hres, err := hc.Do(hreq)
	if err != nil {
		return nil, &FetchError{err, req.Host}
	}
	res := NewResponse(req.Query, req.Host)
	if res.Body, err = ioutil.ReadAll(io.LimitReader(hres.Body, DefaultReadLimit)); err != nil {
		logError(err)
		return nil, &FetchError{err, req.Host}
	}
	res.DetectContentType(hres.Header.Get("Content-Type"))
	return res, nil
}

func httpRequest(ctx context.Context, req *Request) (*http.Request, error) {
	var hreq *http.Request
	var err error
	// POST if non-zero Request.Body
	if len(req.Body) > 0 {
		hreq, err = http.NewRequest("POST", req.URL, bytes.NewReader(req.Body))
	} else {
		hreq, err = http.NewRequest("GET", req.URL, nil)
	}
	if err != nil {
		return nil, err
	}
	// Some web whois servers require a Referer header
	hreq.Header.Add("Referer", req.URL)
	return hreq.WithContext(ctx), nil
}

func logError(err error) {
	switch t := err.(type) {
	case net.Error:
		fmt.Fprintf(os.Stderr, "net.Error timeout=%t, temp=%t: %s\n", t.Timeout(), t.Temporary(), err.Error())
	default:
		fmt.Fprintf(os.Stderr, "Unknown error %v: %s\n", t, err.Error())
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net"
	"time"
)

const pingCount = 10
const protoFmt = "ntwrk%s :%s\r\n"
const timeout = time.Duration(15) * time.Second

// testContext holds a test function, action name, and address to connect to.
type testContext struct {
	Action string
	Fn     func(net.Conn, time.Duration) (int64, error)
	Addr   string
}

// startClient starts the network test suite.
func startClient(host string, port int) {
	addr := fmt.Sprintf("%s:%d", host, port)
	ping(addr)
	perform(testContext{"download", download, addr})
	perform(testContext{"upload", upload, addr})
}

// perform runs a network test and prints the recorded bandwidth.
func perform(ctx testContext) {
	conn, err := openConn(ctx.Addr, ctx.Action)
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return
	}

	since := time.Now()
	ticker := time.NewTicker(time.Millisecond * 150)
	go func() {
		for t := range ticker.C {
			elapsed := t.Sub(since)
			progress := formatProgress(elapsed, timeout)
			fmt.Printf("\r %9s: %s", ctx.Action, progress)
		}
	}()
	defer ticker.Stop()

	bytes, err := ctx.Fn(conn, timeout)
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return
	}

	elapsed := time.Since(since).Seconds()
	fmt.Printf(" %s\n", formatBytes(bytes, elapsed))
}

// ping performs a network latency test.
func ping(addr string) {
	conn, err := openConn(addr, "echo")
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return
	}
	resp := make([]byte, 6)
	since := time.Now()
	for i := 0; i < pingCount; i++ {
		conn.Write([]byte("echo\r\n"))
		conn.Read(resp)
		if string(resp) != "echo\r\n" {
			panic("error: invalid echo")
		}
	}
	elapsed := time.Since(since) / pingCount
	fmt.Printf("\r %9s: %s\n", "latency", elapsed)
}

// whoami requests the client's external IP address from `host` and prints it.
func whoami(host string, port int) {
	addr := fmt.Sprintf("%s:%d", host, port)
	resp := make([]byte, 40)

	conn, err := openConn(addr, "whoami")
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return
	}

	conn.Read(resp)
	fmt.Print(string(resp))
}

// openConn opens a connection to `host` and writes a formatted message to it.
func openConn(host string, action string) (conn net.Conn, err error) {
	conn, err = net.Dial("tcp", host)
	if err != nil {
		return
	}

	fmt.Fprintf(conn, protoFmt, proto, action)
	return
}
<commit_msg>Remove panic from echo function<commit_after>package main

import (
	"fmt"
	"net"
	"time"
)

const pingCount = 10
const protoFmt = "ntwrk%s :%s\r\n"
const timeout = time.Duration(15) * time.Second

// testContext holds a test function, action name, and address to connect to.
type testContext struct {
	Action string
	Fn     func(net.Conn, time.Duration) (int64, error)
	Addr   string
}

// startClient starts the network test suite.
func startClient(host string, port int) {
	addr := fmt.Sprintf("%s:%d", host, port)
	ping(addr)
	perform(testContext{"download", download, addr})
	perform(testContext{"upload", upload, addr})
}

// perform runs a network test and prints the recorded bandwidth.
func perform(ctx testContext) {
	conn, err := openConn(ctx.Addr, ctx.Action)
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return
	}

	since := time.Now()
	ticker := time.NewTicker(time.Millisecond * 150)
	go func() {
		for t := range ticker.C {
			elapsed := t.Sub(since)
			progress := formatProgress(elapsed, timeout)
			fmt.Printf("\r %9s: %s", ctx.Action, progress)
		}
	}()
	defer ticker.Stop()

	bytes, err := ctx.Fn(conn, timeout)
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return
	}

	elapsed := time.Since(since).Seconds()
	fmt.Printf(" %s\n", formatBytes(bytes, elapsed))
}

// ping performs a network latency test.
func ping(addr string) {
	conn, err := openConn(addr, "echo")
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return
	}
	resp := make([]byte, 6)
	since := time.Now()
	for i := 0; i < pingCount; i++ {
		conn.Write([]byte("echo\r\n"))
		conn.Read(resp)
		if string(resp) != "echo\r\n" {
			fmt.Println("error: invalid echo reply")
			return
		}
	}
	elapsed := time.Since(since) / pingCount
	fmt.Printf("\r %9s: %s\n", "latency", elapsed)
}

// whoami requests the client's external IP address from `host` and prints it.
func whoami(host string, port int) {
	addr := fmt.Sprintf("%s:%d", host, port)
	resp := make([]byte, 40)

	conn, err := openConn(addr, "whoami")
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return
	}

	conn.Read(resp)
	fmt.Print(string(resp))
}

// openConn opens a connection to `host` and writes a formatted message to it.
func openConn(host string, action string) (conn net.Conn, err error) {
	conn, err = net.Dial("tcp", host)
	if err != nil {
		return
	}

	fmt.Fprintf(conn, protoFmt, proto, action)
	return
}
<|endoftext|>


<commit_before>package libovsdb

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net"
	"reflect"
	"sync"

	"os"

	"github.com/cenkalti/rpc2"
	"github.com/cenkalti/rpc2/jsonrpc"
)

// OvsdbClient is an OVSDB client
type OvsdbClient struct {
	rpcClient     *rpc2.Client
	Schema        map[string]DatabaseSchema
	handlers      []NotificationHandler
	handlersMutex *sync.Mutex
}

func newOvsdbClient(c *rpc2.Client) *OvsdbClient {
	ovs := &OvsdbClient{
		rpcClient:     c,
		Schema:        make(map[string]DatabaseSchema),
		handlersMutex: &sync.Mutex{},
	}
	connectionsMutex.Lock()
	defer connectionsMutex.Unlock()
	if connections == nil {
		connections = make(map[*rpc2.Client]*OvsdbClient)
	}
	connections[c] = ovs
	return ovs
}

// Would rather replace this connection map with an OvsdbClient Receiver scoped method
// Unfortunately rpc2 package acts wierd with a receiver scoped method and needs some investigation.
var (
	connections      map[*rpc2.Client]*OvsdbClient
	connectionsMutex = &sync.RWMutex{}
)

const (
	// DefaultAddress is the default IPV4 address that is used for a connection
	DefaultAddress = "127.0.0.1"
	// DefaultPort is the default port used for a connection
	DefaultPort     = 6640
	UNIX            = "unix"
	TCP             = "tcp"
	SSL             = "ssl"
	SKIP_TLS_VERIFY = true
)

// ConnectUsingTCP creates an OVSDB connection using TCP and returns and OvsdbClient
func ConnectUsingTCP(protocol string, target string) (*OvsdbClient, error) {
	conn, err := net.Dial(protocol, target)

	if err != nil {
		return nil, err
	}

	c := rpc2.NewClientWithCodec(jsonrpc.NewJSONCodec(conn))
	c.Handle("echo", echo)
	c.Handle("update", update)
	go c.Run()
	go handleDisconnectNotification(c)

	ovs := newOvsdbClient(c)

	// Process Async Notifications
	dbs, err := ovs.ListDbs()
	if err == nil {
		for _, db := range dbs {
			schema, err := ovs.GetSchema(db)
			if err == nil {
				ovs.Schema[db] = *schema
			} else {
				return nil, err
			}
		}
	}
	return ovs, nil
}

// ConnectUsingSSL creates an OVSDB connection using SSL and returns and OvsdbClient
func ConnectUsingSSL(protocol string, target string) (*OvsdbClient, error) {
	cert, err := tls.LoadX509KeyPair(os.Getenv("CLIENT_CERT_CA_CERT"),
		os.Getenv("CLIENT_PRIVKEY"))
	if err != nil {
		log.Fatalf("client: loadkeys: %s", err)
		return nil, err
	}
	if len(cert.Certificate) != 2 {
		log.Fatal("client.crt should have 2 concatenated certificates: client + CA")
		return nil, err
	}
	ca, err := x509.ParseCertificate(cert.Certificate[1])
	if err != nil {
		log.Fatal(err)
		return nil, err
	}
	certPool := x509.NewCertPool()
	certPool.AddCert(ca)
	config := tls.Config{
		Certificates:       []tls.Certificate{cert},
		RootCAs:            certPool,
		InsecureSkipVerify: SKIP_TLS_VERIFY,
	}
	conn, err := tls.Dial(TCP, target, &config)
	if err != nil {
		log.Fatalf("client: dial: %s", err)
		return nil, err
	}
	log.Println("client: connected to: ", conn.RemoteAddr())
	c := rpc2.NewClientWithCodec(jsonrpc.NewJSONCodec(conn))
	c.SetBlocking(true)
	c.Handle("echo", echo)
	c.Handle("update", update)
	go c.Run()
	go handleDisconnectNotification(c)

	ovs := newOvsdbClient(c)

	// Process Async Notifications
	dbs, err := ovs.ListDbs()
	if err == nil {
		for _, db := range dbs {
			schema, err := ovs.GetSchema(db)
			if err == nil {
				ovs.Schema[db] = *schema
			} else {
				return nil, err
			}
		}
	}
	return ovs, nil
}

// Connect creates an OVSDB connection and returns and OvsdbClient
func Connect(ipAddr string, port int, protocol string) (*OvsdbClient, error) {
	if ipAddr == "" {
		ipAddr = DefaultAddress
	}

	if port <= 0 {
		port = DefaultPort
	}

	target := fmt.Sprintf("%s:%d", ipAddr, port)
	switch protocol {
	case TCP, UNIX:
		return ConnectUsingTCP(protocol, target)
	case SSL:
		return ConnectUsingSSL(protocol, target)
	default:
		return nil, errors.New("Supported protocols are TCP, UNIX and SSL")
	}
}

// ConnectWithUnixSocket makes a OVSDB Connection via a Unix Socket
func ConnectWithUnixSocket(socketFile string) (*OvsdbClient, error) {

	if _, err := os.Stat(socketFile); os.IsNotExist(err) {
		return nil, errors.New("Invalid socket file")
	}

	return ConnectUsingTCP(UNIX, socketFile)
}

// Register registers the supplied NotificationHandler to recieve OVSDB Notifications
func (ovs *OvsdbClient) Register(handler NotificationHandler) {
	ovs.handlersMutex.Lock()
	defer ovs.handlersMutex.Unlock()
	ovs.handlers = append(ovs.handlers, handler)
}

//Get Handler by index
func getHandlerIndex(handler NotificationHandler, handlers []NotificationHandler) (int, error) {
	for i, h := range handlers {
		if reflect.DeepEqual(h, handler) {
			return i, nil
		}
	}
	return -1, errors.New("Handler not found")
}

// Unregister the supplied NotificationHandler to not recieve OVSDB Notifications anymore
func (ovs *OvsdbClient) Unregister(handler NotificationHandler) error {
	ovs.handlersMutex.Lock()
	defer ovs.handlersMutex.Unlock()
	i, err := getHandlerIndex(handler, ovs.handlers)
	if err != nil {
		return err
	}
	ovs.handlers = append(ovs.handlers[:i], ovs.handlers[i+1:]...)
	return nil
}

// NotificationHandler is the interface that must be implemented to receive notifcations
type NotificationHandler interface {
	// RFC 7047 section 4.1.6 Update Notification
	Update(context interface{}, tableUpdates TableUpdates)

	// RFC 7047 section 4.1.9 Locked Notification
	Locked([]interface{})

	// RFC 7047 section 4.1.10 Stolen Notification
	Stolen([]interface{})

	// RFC 7047 section 4.1.11 Echo Notification
	Echo([]interface{})

	Disconnected(*OvsdbClient)
}

// RFC 7047 : Section 4.1.6 : Echo
func echo(client *rpc2.Client, args []interface{}, reply *[]interface{}) error {
	*reply = args
	connectionsMutex.RLock()
	defer connectionsMutex.RUnlock()
	if _, ok := connections[client]; ok {
		connections[client].handlersMutex.Lock()
		defer connections[client].handlersMutex.Unlock()
		for _, handler := range connections[client].handlers {
			handler.Echo(nil)
		}
	}
	return nil
}

// RFC 7047 : Update Notification Section 4.1.6
// Processing "params": [<json-value>, <table-updates>]
func update(client *rpc2.Client, params []interface{}, reply *interface{}) error {
	if len(params) < 2 {
		return errors.New("Invalid Update message")
	}
	// Ignore params[0] as we dont use the <json-value> currently for comparison

	raw, ok := params[1].(map[string]interface{})
	if !ok {
		return errors.New("Invalid Update message")
	}
	var rowUpdates map[string]map[string]RowUpdate

	b, err := json.Marshal(raw)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b, &rowUpdates)
	if err != nil {
		return err
	}

	// Update the local DB cache with the tableUpdates
	tableUpdates := getTableUpdatesFromRawUnmarshal(rowUpdates)
	connectionsMutex.RLock()
	defer connectionsMutex.RUnlock()
	if _, ok := connections[client]; ok {
		connections[client].handlersMutex.Lock()
		defer connections[client].handlersMutex.Unlock()
		for _, handler := range connections[client].handlers {
			handler.Update(params, tableUpdates)
		}
	}

	return nil
}

// GetSchema returns the schema in use for the provided database name
// RFC 7047 : get_schema
func (ovs OvsdbClient) GetSchema(dbName string) (*DatabaseSchema, error) {
	args := NewGetSchemaArgs(dbName)
	var reply DatabaseSchema
	err := ovs.rpcClient.Call("get_schema", args, &reply)
	if err != nil {
		return nil, err
	}
	ovs.Schema[dbName] = reply
	return &reply, err
}

// ListDbs returns the list of databases on the server
// RFC 7047 : list_dbs
func (ovs OvsdbClient) ListDbs() ([]string, error) {
	var dbs []string
	err := ovs.rpcClient.Call("list_dbs", nil, &dbs)
	if err != nil {
		log.Fatal("ListDbs failure", err)
	}
	return dbs, err
}

// Transact performs the provided Operation's on the database
// RFC 7047 : transact
func (ovs OvsdbClient) Transact(database string, operation ...Operation) ([]OperationResult, error) {
	var reply []OperationResult
	db, ok := ovs.Schema[database]
	if !ok {
		return nil, errors.New("invalid Database Schema")
	}

	if ok := db.validateOperations(operation...); !ok {
		return nil, errors.New("Validation failed for the operation")
	}

	args := NewTransactArgs(database, operation...)
	err := ovs.rpcClient.Call("transact", args, &reply)
	if err != nil {
		return nil, err
	}
	return reply, nil
}

// MonitorAll is a convenience method to monitor every table/column
func (ovs OvsdbClient) MonitorAll(database string, jsonContext interface{}) (*TableUpdates, error) {
	schema, ok := ovs.Schema[database]
	if !ok {
		return nil, errors.New("invalid Database Schema")
	}

	requests := make(map[string]MonitorRequest)
	for table, tableSchema := range schema.Tables {
		var columns []string
		for column := range tableSchema.Columns {
			columns = append(columns, column)
		}
		requests[table] = MonitorRequest{
			Columns: columns,
			Select: MonitorSelect{
				Initial: true,
				Insert:  true,
				Delete:  true,
				Modify:  true,
			}}
	}
	return ovs.Monitor(database, jsonContext, requests)
}

// MonitorCancel will request cancel a previously issued monitor request
// RFC 7047 : monitor_cancel
func (ovs OvsdbClient) MonitorCancel(database string, jsonContext interface{}) error {
	var reply OperationResult

	args := NewMonitorCancelArgs(jsonContext)

	err := ovs.rpcClient.Call("monitor_cancel", args, &reply)
	if err != nil {
		return err
	}
	if reply.Error != "" {
		return fmt.Errorf("Error while executing transaction: %s", reply.Error)
	}
	return nil
}

// Monitor will provide updates for a given table/column
// RFC 7047 : monitor
func (ovs OvsdbClient) Monitor(database string, jsonContext interface{}, requests map[string]MonitorRequest) (*TableUpdates, error) {
	var reply TableUpdates

	args := NewMonitorArgs(database, jsonContext, requests)

	// This totally sucks. Refer to golang JSON issue #6213
	var response map[string]map[string]RowUpdate
	err := ovs.rpcClient.Call("monitor", args, &response)
	reply = getTableUpdatesFromRawUnmarshal(response)
	if err != nil {
		return nil, err
	}
	return &reply, err
}

func getTableUpdatesFromRawUnmarshal(raw map[string]map[string]RowUpdate) TableUpdates {
	var tableUpdates TableUpdates
	tableUpdates.Updates = make(map[string]TableUpdate)
	for table, update := range raw {
		tableUpdate := TableUpdate{update}
		tableUpdates.Updates[table] = tableUpdate
	}
	return tableUpdates
}

func clearConnection(c *rpc2.Client) {
	connectionsMutex.Lock()
	defer connectionsMutex.Unlock()
	if _, ok := connections[c]; ok {
		for _, handler := range connections[c].handlers {
			if handler != nil {
				handler.Disconnected(connections[c])
			}
		}
	}
	delete(connections, c)
}

func handleDisconnectNotification(c *rpc2.Client) {
	disconnected := c.DisconnectNotify()
	select {
	case <-disconnected:
		clearConnection(c)
	}
}

// Disconnect will close the OVSDB connection
func (ovs OvsdbClient) Disconnect() {
	ovs.rpcClient.Close()
	clearConnection(ovs.rpcClient)
}
<commit_msg>Pass correct context to update notification callbacks<commit_after>package libovsdb

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net"
	"reflect"
	"sync"

	"os"

	"github.com/cenkalti/rpc2"
	"github.com/cenkalti/rpc2/jsonrpc"
)

// OvsdbClient is an OVSDB client
type OvsdbClient struct {
	rpcClient     *rpc2.Client
	Schema        map[string]DatabaseSchema
	handlers      []NotificationHandler
	handlersMutex *sync.Mutex
}

func newOvsdbClient(c *rpc2.Client) *OvsdbClient {
	ovs := &OvsdbClient{
		rpcClient:     c,
		Schema:        make(map[string]DatabaseSchema),
		handlersMutex: &sync.Mutex{},
	}
	connectionsMutex.Lock()
	defer connectionsMutex.Unlock()
	if connections == nil {
		connections = make(map[*rpc2.Client]*OvsdbClient)
	}
	connections[c] = ovs
	return ovs
}

// Would rather replace this connection map with an OvsdbClient Receiver scoped method
// Unfortunately rpc2 package acts wierd with a receiver scoped method and needs some investigation.
var (
	connections      map[*rpc2.Client]*OvsdbClient
	connectionsMutex = &sync.RWMutex{}
)

const (
	// DefaultAddress is the default IPV4 address that is used for a connection
	DefaultAddress = "127.0.0.1"
	// DefaultPort is the default port used for a connection
	DefaultPort     = 6640
	UNIX            = "unix"
	TCP             = "tcp"
	SSL             = "ssl"
	SKIP_TLS_VERIFY = true
)

// ConnectUsingTCP creates an OVSDB connection using TCP and returns and OvsdbClient
func ConnectUsingTCP(protocol string, target string) (*OvsdbClient, error) {
	conn, err := net.Dial(protocol, target)

	if err != nil {
		return nil, err
	}

	c := rpc2.NewClientWithCodec(jsonrpc.NewJSONCodec(conn))
	c.Handle("echo", echo)
	c.Handle("update", update)
	go c.Run()
	go handleDisconnectNotification(c)

	ovs := newOvsdbClient(c)

	// Process Async Notifications
	dbs, err := ovs.ListDbs()
	if err == nil {
		for _, db := range dbs {
			schema, err := ovs.GetSchema(db)
			if err == nil {
				ovs.Schema[db] = *schema
			} else {
				return nil, err
			}
		}
	}
	return ovs, nil
}

// ConnectUsingSSL creates an OVSDB connection using SSL and returns and OvsdbClient
func ConnectUsingSSL(protocol string, target string) (*OvsdbClient, error) {
	cert, err := tls.LoadX509KeyPair(os.Getenv("CLIENT_CERT_CA_CERT"),
		os.Getenv("CLIENT_PRIVKEY"))
	if err != nil {
		log.Fatalf("client: loadkeys: %s", err)
		return nil, err
	}
	if len(cert.Certificate) != 2 {
		log.Fatal("client.crt should have 2 concatenated certificates: client + CA")
		return nil, err
	}
	ca, err := x509.ParseCertificate(cert.Certificate[1])
	if err != nil {
		log.Fatal(err)
		return nil, err
	}
	certPool := x509.NewCertPool()
	certPool.AddCert(ca)
	config := tls.Config{
		Certificates:       []tls.Certificate{cert},
		RootCAs:            certPool,
		InsecureSkipVerify: SKIP_TLS_VERIFY,
	}
	conn, err := tls.Dial(TCP, target, &config)
	if err != nil {
		log.Fatalf("client: dial: %s", err)
		return nil, err
	}
	log.Println("client: connected to: ", conn.RemoteAddr())
	c := rpc2.NewClientWithCodec(jsonrpc.NewJSONCodec(conn))
	c.SetBlocking(true)
	c.Handle("echo", echo)
	c.Handle("update", update)
	go c.Run()
	go handleDisconnectNotification(c)

	ovs := newOvsdbClient(c)

	// Process Async Notifications
	dbs, err := ovs.ListDbs()
	if err == nil {
		for _, db := range dbs {
			schema, err := ovs.GetSchema(db)
			if err == nil {
				ovs.Schema[db] = *schema
			} else {
				return nil, err
			}
		}
	}
	return ovs, nil
}

// Connect creates an OVSDB connection and returns and OvsdbClient
func Connect(ipAddr string, port int, protocol string) (*OvsdbClient, error) {
	if ipAddr == "" {
		ipAddr = DefaultAddress
	}

	if port <= 0 {
		port = DefaultPort
	}

	target := fmt.Sprintf("%s:%d", ipAddr, port)
	switch protocol {
	case TCP, UNIX:
		return ConnectUsingTCP(protocol, target)
	case SSL:
		return ConnectUsingSSL(protocol, target)
	default:
		return nil, errors.New("Supported protocols are TCP, UNIX and SSL")
	}
}

// ConnectWithUnixSocket makes a OVSDB Connection via a Unix Socket
func ConnectWithUnixSocket(socketFile string) (*OvsdbClient, error) {

	if _, err := os.Stat(socketFile); os.IsNotExist(err) {
		return nil, errors.New("Invalid socket file")
	}

	return ConnectUsingTCP(UNIX, socketFile)
}

// Register registers the supplied NotificationHandler to recieve OVSDB Notifications
func (ovs *OvsdbClient) Register(handler NotificationHandler) {
	ovs.handlersMutex.Lock()
	defer ovs.handlersMutex.Unlock()
	ovs.handlers = append(ovs.handlers, handler)
}

//Get Handler by index
func getHandlerIndex(handler NotificationHandler, handlers []NotificationHandler) (int, error) {
	for i, h := range handlers {
		if reflect.DeepEqual(h, handler) {
			return i, nil
		}
	}
	return -1, errors.New("Handler not found")
}

// Unregister the supplied NotificationHandler to not recieve OVSDB Notifications anymore
func (ovs *OvsdbClient) Unregister(handler NotificationHandler) error {
	ovs.handlersMutex.Lock()
	defer ovs.handlersMutex.Unlock()
	i, err := getHandlerIndex(handler, ovs.handlers)
	if err != nil {
		return err
	}
	ovs.handlers = append(ovs.handlers[:i], ovs.handlers[i+1:]...)
	return nil
}

// NotificationHandler is the interface that must be implemented to receive notifcations
type NotificationHandler interface {
	// RFC 7047 section 4.1.6 Update Notification
	Update(context interface{}, tableUpdates TableUpdates)

	// RFC 7047 section 4.1.9 Locked Notification
	Locked([]interface{})

	// RFC 7047 section 4.1.10 Stolen Notification
	Stolen([]interface{})

	// RFC 7047 section 4.1.11 Echo Notification
	Echo([]interface{})

	Disconnected(*OvsdbClient)
}

// RFC 7047 : Section 4.1.6 : Echo
func echo(client *rpc2.Client, args []interface{}, reply *[]interface{}) error {
	*reply = args
	connectionsMutex.RLock()
	defer connectionsMutex.RUnlock()
	if _, ok := connections[client]; ok {
		connections[client].handlersMutex.Lock()
		defer connections[client].handlersMutex.Unlock()
		for _, handler := range connections[client].handlers {
			handler.Echo(nil)
		}
	}
	return nil
}

// RFC 7047 : Update Notification Section 4.1.6
// Processing "params": [<json-value>, <table-updates>]
func update(client *rpc2.Client, params []interface{}, reply *interface{}) error {
	if len(params) < 2 {
		return errors.New("Invalid Update message")
	}
	// Ignore params[0] as we dont use the <json-value> currently for comparison

	raw, ok := params[1].(map[string]interface{})
	if !ok {
		return errors.New("Invalid Update message")
	}
	var rowUpdates map[string]map[string]RowUpdate

	b, err := json.Marshal(raw)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b, &rowUpdates)
	if err != nil {
		return err
	}

	// Update the local DB cache with the tableUpdates
	tableUpdates := getTableUpdatesFromRawUnmarshal(rowUpdates)
	connectionsMutex.RLock()
	defer connectionsMutex.RUnlock()
	if _, ok := connections[client]; ok {
		connections[client].handlersMutex.Lock()
		defer connections[client].handlersMutex.Unlock()
		for _, handler := range connections[client].handlers {
			handler.Update(params[0], tableUpdates)
		}
	}

	return nil
}

// GetSchema returns the schema in use for the provided database name
// RFC 7047 : get_schema
func (ovs OvsdbClient) GetSchema(dbName string) (*DatabaseSchema, error) {
	args := NewGetSchemaArgs(dbName)
	var reply DatabaseSchema
	err := ovs.rpcClient.Call("get_schema", args, &reply)
	if err != nil {
		return nil, err
	}
	ovs.Schema[dbName] = reply
	return &reply, err
}

// ListDbs returns the list of databases on the server
// RFC 7047 : list_dbs
func (ovs OvsdbClient) ListDbs() ([]string, error) {
	var dbs []string
	err := ovs.rpcClient.Call("list_dbs", nil, &dbs)
	if err != nil {
		log.Fatal("ListDbs failure", err)
	}
	return dbs, err
}

// Transact performs the provided Operation's on the database
// RFC 7047 : transact
func (ovs OvsdbClient) Transact(database string, operation ...Operation) ([]OperationResult, error) {
	var reply []OperationResult
	db, ok := ovs.Schema[database]
	if !ok {
		return nil, errors.New("invalid Database Schema")
	}

	if ok := db.validateOperations(operation...); !ok {
		return nil, errors.New("Validation failed for the operation")
	}

	args := NewTransactArgs(database, operation...)
	err := ovs.rpcClient.Call("transact", args, &reply)
	if err != nil {
		return nil, err
	}
	return reply, nil
}

// MonitorAll is a convenience method to monitor every table/column
func (ovs OvsdbClient) MonitorAll(database string, jsonContext interface{}) (*TableUpdates, error) {
	schema, ok := ovs.Schema[database]
	if !ok {
		return nil, errors.New("invalid Database Schema")
	}

	requests := make(map[string]MonitorRequest)
	for table, tableSchema := range schema.Tables {
		var columns []string
		for column := range tableSchema.Columns {
			columns = append(columns, column)
		}
		requests[table] = MonitorRequest{
			Columns: columns,
			Select: MonitorSelect{
				Initial: true,
				Insert:  true,
				Delete:  true,
				Modify:  true,
			}}
	}
	return ovs.Monitor(database, jsonContext, requests)
}

// MonitorCancel will request cancel a previously issued monitor request
// RFC 7047 : monitor_cancel
func (ovs OvsdbClient) MonitorCancel(database string, jsonContext interface{}) error {
	var reply OperationResult

	args := NewMonitorCancelArgs(jsonContext)

	err := ovs.rpcClient.Call("monitor_cancel", args, &reply)
	if err != nil {
		return err
	}
	if reply.Error != "" {
		return fmt.Errorf("Error while executing transaction: %s", reply.Error)
	}
	return nil
}

// Monitor will provide updates for a given table/column
// RFC 7047 : monitor
func (ovs OvsdbClient) Monitor(database string, jsonContext interface{}, requests map[string]MonitorRequest) (*TableUpdates, error) {
	var reply TableUpdates

	args := NewMonitorArgs(database, jsonContext, requests)

	// This totally sucks. Refer to golang JSON issue #6213
	var response map[string]map[string]RowUpdate
	err := ovs.rpcClient.Call("monitor", args, &response)
	reply = getTableUpdatesFromRawUnmarshal(response)
	if err != nil {
		return nil, err
	}
	return &reply, err
}

func getTableUpdatesFromRawUnmarshal(raw map[string]map[string]RowUpdate) TableUpdates {
	var tableUpdates TableUpdates
	tableUpdates.Updates = make(map[string]TableUpdate)
	for table, update := range raw {
		tableUpdate := TableUpdate{update}
		tableUpdates.Updates[table] = tableUpdate
	}
	return tableUpdates
}

func clearConnection(c *rpc2.Client) {
	connectionsMutex.Lock()
	defer connectionsMutex.Unlock()
	if _, ok := connections[c]; ok {
		for _, handler := range connections[c].handlers {
			if handler != nil {
				handler.Disconnected(connections[c])
			}
		}
	}
	delete(connections, c)
}

func handleDisconnectNotification(c *rpc2.Client) {
	disconnected := c.DisconnectNotify()
	select {
	case <-disconnected:
		clearConnection(c)
	}
}

// Disconnect will close the OVSDB connection
func (ovs OvsdbClient) Disconnect() {
	ovs.rpcClient.Close()
	clearConnection(ovs.rpcClient)
}
<|endoftext|>


<commit_before>package loggly

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
)

//Does not parse json properly when using bulk endpoint
const api = "https://logs-01.loggly.com/inputs/"
const Version = "0.0.1"

type Client interface {
	Send([]byte)
}

type client struct {
	token string
	tags  []string
}

func NewClient(token string, tags ...string) Client {
	c := &client{
		token: token,
		tags:  tags,
	}
	return c
}

func (c *client) Send(body []byte) {
	defer func() {
		if err := recover(); err != nil {
			fmt.Printf("Panic writing to loggly: %s", err)
		}
	}()
	req, err := http.NewRequest("POST", api+c.token, bytes.NewBuffer(body))
	if err != nil {
		fmt.Printf("Error writing to loggly: %v\n", err)
		return
	}

	req.Header.Add("User-Agent", "tt: go-loggly ("+Version+")")
	req.Header.Add("Content-Type", "text/plain")
	req.Header.Add("Content-Length", string(len(body)))

	if len(c.tags) != 0 {
		req.Header.Add("X-Loggly-Tag", strings.Join(c.tags, ","))
	}
	client := &http.Client{}
	res, err := client.Do(req)
	if err != nil {
		fmt.Printf("Error writing to loggly: %v\n", err)
		return
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		resp, _ := ioutil.ReadAll(res.Body)
		fmt.Printf("Error writing to loggly: %s (%v)\n", string(resp), res.StatusCode)
	}
}
<commit_msg>API-1698 fix Content-Length header<commit_after>package loggly

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"
)

//Does not parse json properly when using bulk endpoint
const api = "https://logs-01.loggly.com/inputs/"
const Version = "0.0.1"

type Client interface {
	Send([]byte)
}

type client struct {
	token string
	tags  []string
}

func NewClient(token string, tags ...string) Client {
	c := &client{
		token: token,
		tags:  tags,
	}
	return c
}

func (c *client) Send(body []byte) {
	defer func() {
		if err := recover(); err != nil {
			fmt.Printf("Panic writing to loggly: %s", err)
		}
	}()
	req, err := http.NewRequest("POST", api+c.token, bytes.NewBuffer(body))
	if err != nil {
		fmt.Printf("Error creating loggly request: %v\n", err)
		return
	}

	req.Header.Add("User-Agent", "tt: go-loggly ("+Version+")")
	req.Header.Add("Content-Type", "text/plain")
	req.Header.Add("Content-Length", strconv.Itoa(len(body)))

	if len(c.tags) != 0 {
		req.Header.Add("X-Loggly-Tag", strings.Join(c.tags, ","))
	}
	client := &http.Client{}
	res, err := client.Do(req)
	if err != nil {
		fmt.Printf("Error sending to loggly: %v\n", err)
		return
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		resp, _ := ioutil.ReadAll(res.Body)
		fmt.Printf("Error response from loggly: %s (%v)\n", string(resp), res.StatusCode)
	}
}
<|endoftext|>