<commit_before>// Standardize the mac address.
package mac

import (
	"fmt"
	"strconv"
	"strings"
)

type Standard struct {
	// If true, output the upper character, such as AA:BB:11:22:33:44.
	// Or, output the lower character.
	Upper bool

	// If true, pad with leading zero, such as 01:02:03:04:05:06.
	Unified bool
}

var (
	StandardUU = NewStandard(true, true)
	StandardUu = NewStandard(true, false)
	StandarduU = NewStandard(false, true)
	Standarduu = NewStandard(false, false)
)

func NewStandard(upper, unified bool) Standard {
	return Standard{Upper: upper, Unified: unified}
}

// Convert the argument of mac to the specifical standard mac address.
//
// Return the empty string if the argument of mac is not the legal mac address.
func (m Standard) Standardize(mac string) string {
	macs := strings.Split(mac, ":")
	if len(macs) != 6 {
		return ""
	}

	width := ""
	upper := "x"
	if m.Upper {
		upper = "X"
	}
	if m.Unified {
		width = "2"
	}
	formatter := fmt.Sprintf("%%0%s%s", width, upper)

	for i := 0; i < 6; i++ {
		if _v, err := strconv.ParseUint(macs[i], 16, 64); err != nil {
			return ""
		} else {
			macs[i] = fmt.Sprintf(formatter, _v)
		}
	}

	return strings.Join(macs, ":")
}

// Same as NewStandard(true, true).Standardize(mac)
func StandardizeUU(mac string) string {
	return StandardUU.Standardize(mac)
}

// Same as NewStandard(true, false).Standardize(mac)
func StandardizeUu(mac string) string {
	return StandardUu.Standardize(mac)
}

// Same as NewStandard(false, true).Standardize(mac)
func StandardizeuU(mac string) string {
	return StandarduU.Standardize(mac)
}

// Same as NewStandard(false, false).Standardize(mac)
func Standardizeuu(mac string) string {
	return Standarduu.Standardize(mac)
}
<commit_msg>Add the default standardize and comment<commit_after>// Package mac standardize the mac address.
package mac

import (
	"fmt"
	"strconv"
	"strings"
)

// Standard is a struct to standardize MAC
type Standard struct {
	// If true, output the upper character, such as AA:BB:11:22:33:44.
	// Or, output the lower character.
	Upper bool

	// If true, pad with leading zero, such as 01:02:03:04:05:06.
	Unified bool
}

var (
	// StandardUU standardize MAC such as "".
	StandardUU = NewStandard(true, true)

	// StandardUu standardize MAC such as "".
	StandardUu = NewStandard(true, false)

	// StandarduU standardize MAC such as "".
	StandarduU = NewStandard(false, true)

	// Standarduu standardize MAC such as "".
	Standarduu = NewStandard(false, false)

	// Default is alias of StandardizeUU.
	Default = StandardUU
)

// NewStandard returns a new Standard.
func NewStandard(upper, unified bool) Standard {
	return Standard{Upper: upper, Unified: unified}
}

// Standardize converts the argument of mac to the specifical standard mac address.
//
// Return the empty string if the argument of mac is not the legal mac address.
func (m Standard) Standardize(mac string) string {
	macs := strings.Split(mac, ":")
	if len(macs) != 6 {
		return ""
	}

	width := ""
	upper := "x"
	if m.Upper {
		upper = "X"
	}
	if m.Unified {
		width = "2"
	}
	formatter := fmt.Sprintf("%%0%s%s", width, upper)

	for i := 0; i < 6; i++ {
		if _v, err := strconv.ParseUint(macs[i], 16, 64); err != nil {
			return ""
		} else {
			macs[i] = fmt.Sprintf(formatter, _v)
		}
	}

	return strings.Join(macs, ":")
}

// StandardizeUU is the same as NewStandard(true, true).Standardize(mac)
func StandardizeUU(mac string) string {
	return StandardUU.Standardize(mac)
}

// StandardizeUu is the same as NewStandard(true, false).Standardize(mac)
func StandardizeUu(mac string) string {
	return StandardUu.Standardize(mac)
}

// StandardizeuU is the same as NewStandard(false, true).Standardize(mac)
func StandardizeuU(mac string) string {
	return StandarduU.Standardize(mac)
}

// Standardizeuu is the same as NewStandard(false, false).Standardize(mac)
func Standardizeuu(mac string) string {
	return Standarduu.Standardize(mac)
}

// StandardizeDefault is the same as Default.Standardize(mac)
func StandardizeDefault(mac string) string {
	return Default.Standardize(mac)
}
<|endoftext|>


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

package main

import (
	"bufio"
	"code.google.com/p/gopacket"
	"code.google.com/p/gopacket/layers"
	"code.google.com/p/gopacket/pcap"
	"code.google.com/p/gopacket/tcpassembly"
	"code.google.com/p/gopacket/tcpassembly/tcpreader"
	"io"
	"log"
	"net/http"
	"time"
)

type PacketLayer struct {
	assembler     *tcpassembly.Assembler
	streamFactory *httpStreamFactory
	storage       *Storage
}

func NewPacketLayer(s *Storage) *PacketLayer {
	pl := &PacketLayer{storage: s}

	// Set up assembly
	pl.streamFactory = newStreamFactory(pl.storage)
	streamPool := tcpassembly.NewStreamPool(pl.streamFactory)
	pl.assembler = tcpassembly.NewAssembler(streamPool)

	return pl
}

func (pl *PacketLayer) CreatePacketsChannel() (packets chan gopacket.Packet) {
	var handle *pcap.Handle
	var err error

	if *inputfile != "" {
		handle, err = pcap.OpenOffline(*inputfile)
	} else {
		log.Printf("starting capture on interface %q", *iface)
		// Setup packet capture
		handle, err = pcap.OpenLive(*iface, 1600, true, pcap.BlockForever)
	}

	if err != nil {
		panic(err)
	} else if err := handle.SetBPFFilter("tcp and port 80"); err != nil {
		panic(err)
	}
	log.Println("reading in packets. Press CTRL-C to end and report.")
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packets = packetSource.Packets()

	return
}

func (pl *PacketLayer) HandlePacket(packet gopacket.Packet) {
	if pl.storage.PacketInScope(packet) {
		pl.streamFactory.logPacketSize(packet)
		netFlow := packet.NetworkLayer().NetworkFlow()
		tcp := packet.TransportLayer().(*layers.TCP)

		pl.assembler.AssembleWithTimestamp(netFlow, tcp,
			packet.Metadata().Timestamp)
	}
}

func (pl *PacketLayer) Close() {

	// Cleanup the go routines
	// Ignore any http request/response parsing errors when closing the streams.
	pl.streamFactory.closed = true
	for _, v := range pl.streamFactory.bidiStreams {
		if v.in != nil {
			v.in.closed = true
		}
		if v.out != nil {
			v.out.closed = true
		}
	}

	pl.assembler.FlushAll()
}

type bidiStream struct {
	key      uint64
	in, out  *tcpStream
	requests chan *http.Request
}

// tcpStream will handle the actual decoding of http requests and responses.
type tcpStream struct {
	netFlow, tcpFlow gopacket.Flow
	readStream       tcpreader.ReaderStream
	storage          *Storage
	bidikey          uint64
	closed           bool
	reqInProgress    *http.Request
}

// runOut is a blocking function that reads HTTP requests from a stream.
func (h *tcpStream) runOut(bds *bidiStream) {
	buf := bufio.NewReader(&h.readStream)
	var reqID int64

	for {
		/*      _, err := buf.Peek(1)
		        if err == io.EOF {
		            return
		        }*/
		req, err := http.ReadRequest(buf)
		if err == io.EOF {
			//          log.Println("EOF while reading stream", h.netFlow, h.tcpFlow, ":", err)
			// We must read until we see an EOF... very important!
			err = h.storage.CloseTCPConnection(h.bidikey, time.Now())
			if err != nil {
				log.Println("Error storing connection close timestamp", err)
			}
			return
		} else if err != nil {
			tcpreader.DiscardBytesToFirstError(buf)

			if h.closed == true {
				// error occurred after stream was closed, ignore.
			} else {
				log.Println("Error reading stream", h.netFlow, h.tcpFlow, ":", err)
			}
		} else {
			// bodyBytes := tcpreader.DiscardBytesToEOF(req.Body)
			req.Body.Close()
			bds.requests <- req
			err = h.storage.SentRequest(h.bidikey, reqID, time.Now(), req)

			if err != nil {
				log.Println("Error storing request", err)
			}

			reqID++
			//          fmt.Print(".")
			// log.Println("Received request from stream", h.netFlow, h.tcpFlow,
			//  ":", req, "with", bodyBytes, "bytes in request body")
		}
	}
}

// runIn is a blocking function that reads HTTP responses from a stream.
func (h *tcpStream) runIn(bds *bidiStream) {
	buf := bufio.NewReader(&h.readStream)
	var reqID int64

	for {
		// Don't start reading a response if no data is available
		_, err := buf.Peek(1)
		if err == io.EOF {
			return
		}

		// Data available, read response.

		// Find the request to which this is the response.
		req := h.reqInProgress
		if req == nil {
			req = <-bds.requests
			h.reqInProgress = req
		}

		resp, err := http.ReadResponse(buf, req)
		if err == io.EOF {
			// We must read until we see an EOF... very important!
			//          log.Println("EOF while reading stream", h.netFlow, h.tcpFlow, ":", err)
			return
		} else if err != nil {
			tcpreader.DiscardBytesToFirstError(buf)
			if h.closed == true {
				// error occurred after stream was closed, ignore.
			} else {
				log.Println("Error reading stream", h.netFlow, h.tcpFlow, ":", err)
			}
		} else {
			_, err := tcpreader.DiscardBytesToFirstError(resp.Body)
			if err != nil && err != io.EOF {
				log.Println("Error discarding bytes ", err)
			}
			resp.Body.Close()
			err = h.storage.ReceivedResponse(h.bidikey, reqID, time.Now(), resp)
			if err != nil {
				log.Println("Error storing response", err)
			}

			reqID++
			h.reqInProgress = nil

			// fmt.Print(".")
			//log.Println("Received response from stream", h.netFlow, h.tcpFlow,
			//  ":", resp, "with", bodyBytes, "bytes in response body")
		}
	}

}

// httpStreamFactory implements tcpassembly.StreamFactory
type httpStreamFactory struct {
	bidiStreams map[uint64]*bidiStream
	storage     *Storage
	closed      bool
}

func newStreamFactory(s *Storage) *httpStreamFactory {
	return &httpStreamFactory{bidiStreams: make(map[uint64]*bidiStream),
		storage: s}
}

func (h *httpStreamFactory) New(netFlow, tcpFlow gopacket.Flow) tcpassembly.Stream {

	// Watch out: this function can still get called even after all
	// streams were flushed (via FlushAll) and closed.
	/*  if h.closed == true {
	        return tcpreader.NewReaderStream()
	    }
	*/
	// First the outgoing stream, then the incoming stream
	key := netFlow.FastHash() ^ tcpFlow.FastHash()

	hstream := &tcpStream{
		netFlow:    netFlow,
		tcpFlow:    tcpFlow,
		readStream: tcpreader.NewReaderStream(),
		storage:    h.storage,
		bidikey:    key,
	}

	bds := h.bidiStreams[key]
	if bds == nil {
		//      log.Println("reading stream", netFlow, tcpFlow)
		bds = &bidiStream{out: hstream, key: key,
			requests: make(chan *http.Request, 100)}
		h.bidiStreams[key] = bds
		// Start a coroutine per stream, to ensure that all data is read from
		// the reader stream
		go hstream.runOut(bds)
	} else {
		//      log.Println("opening TCP conn", netFlow, tcpFlow)
		bds.in = hstream
		err := h.storage.OpenTCPConnection(key, time.Now())
		if err != nil {
			log.Println("Error storing connection", err)
		}
		// Start a coroutine per stream, to ensure that all data is read from
		// the reader stream
		go hstream.runIn(bds)
	}

	// ReaderStream implements tcpassembly.Stream, so we can return a pointer to it.
	return &hstream.readStream
}

// logPacketSize calculates the payload length of a TCP packet and stores it
// in the storage layer.
func (h *httpStreamFactory) logPacketSize(packet gopacket.Packet) {
	netFlow := packet.NetworkLayer().NetworkFlow()
	tcpFlow := packet.TransportLayer().TransportFlow()
	key := netFlow.FastHash() ^ tcpFlow.FastHash()

	ipv4Layer := packet.Layer(layers.LayerTypeIPv4)
	ipv4, _ := ipv4Layer.(*layers.IPv4)

	tcpLayer := packet.Layer(layers.LayerTypeTCP)
	tcp, _ := tcpLayer.(*layers.TCP)

	bds := h.bidiStreams[key]
	if bds == nil || bds.in == nil || bds.out == nil {
		return
	}

	payloadLength := uint32(ipv4.Length - uint16(ipv4.IHL)*4 - uint16(tcp.DataOffset)*4)

	if bds.in.netFlow == netFlow {
		// This is an incoming packet
		err := h.storage.IncomingTCPPacket(key, payloadLength)
		if err != nil {
			panic(err)
		}
	} else {
		// This is an outgoing packet
		err := h.storage.OutgoingTCPPacket(key, payloadLength)
		if err != nil {
			panic(err)
		}
	}
}
<commit_msg>Avoid crash when sniffing non tcp+ipv4 packets.<commit_after>// Copyright 2014 Lieven Govaerts. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bufio"
	"code.google.com/p/gopacket"
	"code.google.com/p/gopacket/layers"
	"code.google.com/p/gopacket/pcap"
	"code.google.com/p/gopacket/tcpassembly"
	"code.google.com/p/gopacket/tcpassembly/tcpreader"
	"io"
	"log"
	"net/http"
	"time"
)

type PacketLayer struct {
	assembler     *tcpassembly.Assembler
	streamFactory *httpStreamFactory
	storage       *Storage
}

func NewPacketLayer(s *Storage) *PacketLayer {
	pl := &PacketLayer{storage: s}

	// Set up assembly
	pl.streamFactory = newStreamFactory(pl.storage)
	streamPool := tcpassembly.NewStreamPool(pl.streamFactory)
	pl.assembler = tcpassembly.NewAssembler(streamPool)

	return pl
}

func (pl *PacketLayer) CreatePacketsChannel() (packets chan gopacket.Packet) {
	var handle *pcap.Handle
	var err error

	if *inputfile != "" {
		handle, err = pcap.OpenOffline(*inputfile)
	} else {
		log.Printf("starting capture on interface %q", *iface)
		// Setup packet capture
		handle, err = pcap.OpenLive(*iface, 1600, true, pcap.BlockForever)
	}

	if err != nil {
		panic(err)
	} else if err := handle.SetBPFFilter("tcp and port 80"); err != nil {
		panic(err)
	}
	log.Println("reading in packets. Press CTRL-C to end and report.")
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packets = packetSource.Packets()

	return
}

func (pl *PacketLayer) HandlePacket(packet gopacket.Packet) {
	if pl.storage.PacketInScope(packet) {
		pl.streamFactory.logPacketSize(packet)
		netFlow := packet.NetworkLayer().NetworkFlow()
		tcp := packet.TransportLayer().(*layers.TCP)

		pl.assembler.AssembleWithTimestamp(netFlow, tcp,
			packet.Metadata().Timestamp)
	}
}

func (pl *PacketLayer) Close() {

	// Cleanup the go routines
	// Ignore any http request/response parsing errors when closing the streams.
	pl.streamFactory.closed = true
	for _, v := range pl.streamFactory.bidiStreams {
		if v.in != nil {
			v.in.closed = true
		}
		if v.out != nil {
			v.out.closed = true
		}
	}

	pl.assembler.FlushAll()
}

type bidiStream struct {
	key      uint64
	in, out  *tcpStream
	requests chan *http.Request
}

// tcpStream will handle the actual decoding of http requests and responses.
type tcpStream struct {
	netFlow, tcpFlow gopacket.Flow
	readStream       tcpreader.ReaderStream
	storage          *Storage
	bidikey          uint64
	closed           bool
	reqInProgress    *http.Request
}

// runOut is a blocking function that reads HTTP requests from a stream.
func (h *tcpStream) runOut(bds *bidiStream) {
	buf := bufio.NewReader(&h.readStream)
	var reqID int64

	for {
		/*      _, err := buf.Peek(1)
		        if err == io.EOF {
		            return
		        }*/
		req, err := http.ReadRequest(buf)
		if err == io.EOF {
			//          log.Println("EOF while reading stream", h.netFlow, h.tcpFlow, ":", err)
			// We must read until we see an EOF... very important!
			err = h.storage.CloseTCPConnection(h.bidikey, time.Now())
			if err != nil {
				log.Println("Error storing connection close timestamp", err)
			}
			return
		} else if err != nil {
			tcpreader.DiscardBytesToFirstError(buf)

			if h.closed == true {
				// error occurred after stream was closed, ignore.
			} else {
				log.Println("Error reading stream", h.netFlow, h.tcpFlow, ":", err)
			}
		} else {
			// bodyBytes := tcpreader.DiscardBytesToEOF(req.Body)
			req.Body.Close()
			bds.requests <- req
			err = h.storage.SentRequest(h.bidikey, reqID, time.Now(), req)

			if err != nil {
				log.Println("Error storing request", err)
			}

			reqID++
			//          fmt.Print(".")
			// log.Println("Received request from stream", h.netFlow, h.tcpFlow,
			//  ":", req, "with", bodyBytes, "bytes in request body")
		}
	}
}

// runIn is a blocking function that reads HTTP responses from a stream.
func (h *tcpStream) runIn(bds *bidiStream) {
	buf := bufio.NewReader(&h.readStream)
	var reqID int64

	for {
		// Don't start reading a response if no data is available
		_, err := buf.Peek(1)
		if err == io.EOF {
			return
		}

		// Data available, read response.

		// Find the request to which this is the response.
		req := h.reqInProgress
		if req == nil {
			req = <-bds.requests
			h.reqInProgress = req
		}

		resp, err := http.ReadResponse(buf, req)
		if err == io.EOF {
			// We must read until we see an EOF... very important!
			//          log.Println("EOF while reading stream", h.netFlow, h.tcpFlow, ":", err)
			return
		} else if err != nil {
			tcpreader.DiscardBytesToFirstError(buf)
			if h.closed == true {
				// error occurred after stream was closed, ignore.
			} else {
				log.Println("Error reading stream", h.netFlow, h.tcpFlow, ":", err)
			}
		} else {
			_, err := tcpreader.DiscardBytesToFirstError(resp.Body)
			if err != nil && err != io.EOF {
				log.Println("Error discarding bytes ", err)
			}
			resp.Body.Close()
			err = h.storage.ReceivedResponse(h.bidikey, reqID, time.Now(), resp)
			if err != nil {
				log.Println("Error storing response", err)
			}

			reqID++
			h.reqInProgress = nil

			// fmt.Print(".")
			//log.Println("Received response from stream", h.netFlow, h.tcpFlow,
			//  ":", resp, "with", bodyBytes, "bytes in response body")
		}
	}

}

// httpStreamFactory implements tcpassembly.StreamFactory
type httpStreamFactory struct {
	bidiStreams map[uint64]*bidiStream
	storage     *Storage
	closed      bool
}

func newStreamFactory(s *Storage) *httpStreamFactory {
	return &httpStreamFactory{bidiStreams: make(map[uint64]*bidiStream),
		storage: s}
}

func (h *httpStreamFactory) New(netFlow, tcpFlow gopacket.Flow) tcpassembly.Stream {

	// Watch out: this function can still get called even after all
	// streams were flushed (via FlushAll) and closed.
	/*  if h.closed == true {
	        return tcpreader.NewReaderStream()
	    }
	*/
	// First the outgoing stream, then the incoming stream
	key := netFlow.FastHash() ^ tcpFlow.FastHash()

	hstream := &tcpStream{
		netFlow:    netFlow,
		tcpFlow:    tcpFlow,
		readStream: tcpreader.NewReaderStream(),
		storage:    h.storage,
		bidikey:    key,
	}

	bds := h.bidiStreams[key]
	if bds == nil {
		//      log.Println("reading stream", netFlow, tcpFlow)
		bds = &bidiStream{out: hstream, key: key,
			requests: make(chan *http.Request, 100)}
		h.bidiStreams[key] = bds
		// Start a coroutine per stream, to ensure that all data is read from
		// the reader stream
		go hstream.runOut(bds)
	} else {
		//      log.Println("opening TCP conn", netFlow, tcpFlow)
		bds.in = hstream
		err := h.storage.OpenTCPConnection(key, time.Now())
		if err != nil {
			log.Println("Error storing connection", err)
		}
		// Start a coroutine per stream, to ensure that all data is read from
		// the reader stream
		go hstream.runIn(bds)
	}

	// ReaderStream implements tcpassembly.Stream, so we can return a pointer to it.
	return &hstream.readStream
}

// logPacketSize calculates the payload length of a TCP packet and stores it
// in the storage layer.
func (h *httpStreamFactory) logPacketSize(packet gopacket.Packet) {
	netFlow := packet.NetworkLayer().NetworkFlow()
	tcpFlow := packet.TransportLayer().TransportFlow()
	key := netFlow.FastHash() ^ tcpFlow.FastHash()

	ipv4Layer := packet.Layer(layers.LayerTypeIPv4)
	ipv4, _ := ipv4Layer.(*layers.IPv4)

	tcpLayer := packet.Layer(layers.LayerTypeTCP)
	tcp, _ := tcpLayer.(*layers.TCP)

	bds := h.bidiStreams[key]
	if bds == nil || bds.in == nil || bds.out == nil {
		return
	}

	payloadLength := uint32(0)

	if ipv4 != nil {
		payloadLength += uint32(ipv4.Length - uint16(ipv4.IHL)*4)
	}
	if tcp != nil {
		payloadLength -= uint32(uint16(tcp.DataOffset) * 4)
	}

	if bds.in.netFlow == netFlow {
		// This is an incoming packet
		err := h.storage.IncomingTCPPacket(key, payloadLength)
		if err != nil {
			panic(err)
		}
	} else {
		// This is an outgoing packet
		err := h.storage.OutgoingTCPPacket(key, payloadLength)
		if err != nil {
			panic(err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"math/rand"
)

// Router to nodes
type NodeRouter struct {
}

// Pick node
func (this *NodeRouter) PickNode() (string, error) {
	// @todo smarter node selection, e.g. https://labs.spotify.com/2015/12/08/els-part-1/
	nodes := gossip.GetNodeStates()
	nodeCount := len(nodes)
	selected := rand.Intn(nodeCount)
	i := 0
	for _, node := range nodes {
		if i == selected {
			return node.Node, nil
		}
		i++
	}
	return "", errors.New("Unable to pick node from router")
}

// New router
func newNodeRouter() *NodeRouter {
	return &NodeRouter{}
}
<commit_msg>No nodes route to local<commit_after>package main

import (
	"errors"
	"math/rand"
)

// Router to nodes
type NodeRouter struct {
}

// Pick node
func (this *NodeRouter) PickNode() (string, error) {
	// @todo smarter node selection, e.g. https://labs.spotify.com/2015/12/08/els-part-1/
	nodes := gossip.GetNodeStates()
	nodeCount := len(nodes)

	// No nodes? Route to localhost
	if nodeCount == 0 {
		return "127.0.0.1", nil
	}

	// Select
	selected := rand.Intn(nodeCount)
	i := 0
	for _, node := range nodes {
		if i == selected {
			return node.Node, nil
		}
		i++
	}
	return "", errors.New("Unable to pick node from router")
}

// New router
func newNodeRouter() *NodeRouter {
	return &NodeRouter{}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/bradfitz/gomemcache/memcache"
	"github.com/nranchev/go-libGeoIP"
	"io"
	"koding/kontrol/kontrolhelper"
	"koding/kontrol/kontrolproxy/proxyconfig"
	"koding/tools/config"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
)

func init() {
	log.SetPrefix("kontrol-proxy ")
}

type RabbitChannel struct {
	ReplyTo string
	Receive chan []byte
}

var proxyDB *proxyconfig.ProxyConfiguration
var amqpStream *AmqpStream
var connections map[string]RabbitChannel
var geoIP *libgeo.GeoIP
var memCache *memcache.Client
var hostname = kontrolhelper.CustomHostname()

func main() {
	log.Printf("kontrol proxy started ")
	connections = make(map[string]RabbitChannel)

	// open and read from DB
	var err error
	proxyDB, err = proxyconfig.Connect()
	if err != nil {
		log.Fatalf("proxyconfig mongodb connect: %s", err)
	}

	err = proxyDB.AddProxy(hostname)
	if err != nil {
		log.Println(err)
	}

	memCache = memcache.New("127.0.0.1:11211") // used for vm lookup

	// load GeoIP db into memory
	dbFile := "GeoIP.dat"
	geoIP, err = libgeo.Load(dbFile)
	if err != nil {
		log.Printf("load GeoIP.dat: %s\n", err.Error())
	}

	// create amqpStream for rabbitmq proxyieng
	amqpStream = setupAmqp()

	reverseProxy := &ReverseProxy{}
	// http.HandleFunc("/", reverseProxy.ServeHTTP) this works for 1.1
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		reverseProxy.ServeHTTP(w, r)
	})

	port := strconv.Itoa(config.Current.Kontrold.Proxy.Port)
	portssl := strconv.Itoa(config.Current.Kontrold.Proxy.PortSSL)
	sslips := strings.Split(config.Current.Kontrold.Proxy.SSLIPS, ",")

	for _, sslip := range sslips {
		go func(sslip string) {
			err = http.ListenAndServeTLS(sslip+":"+portssl, sslip+"_cert.pem", sslip+"_key.pem", nil)
			if err != nil {
				log.Printf("https mode is disabled. please add cert.pem and key.pem files. %s %s", err, sslip)
			} else {
				log.Printf("https mode is enabled. serving at :%s ...", portssl)
			}
		}(sslip)
	}

	log.Printf("normal mode is enabled. serving at :%s ...", port)
	err = http.ListenAndServe(":"+port, nil)
	if err != nil {
		log.Println(err)
	}
}

/*************************************************
*
*  util functions
*
*  - arslan
*************************************************/
// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func hasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Given a string of the form "host", "port", returns "host:port"
func addPort(host, port string) string {
	if ok := hasPort(host); ok {
		return host
	}

	return host + ":" + port
}

// Check if a server is alive or not
func checkServer(host string) error {
	c, err := net.Dial("tcp", host)
	if err != nil {
		return err
	}
	c.Close()
	return nil
}

/*************************************************
*
*  modified version of go's reverseProxy source code
*  has support for dynamic target url, websockets and amqp
*
*  - arslan
*************************************************/

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// redirect http to https
	if req.TLS == nil && req.Host == "new.koding.com" {
		http.Redirect(rw, req, "https://new.koding.com"+req.RequestURI, http.StatusMovedPermanently)
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	// if connection is of type websocket, hijacking is used instead of http proxy
	websocket := checkWebsocket(outreq)

	user, err := populateUser(outreq)
	if err != nil {
		log.Printf("error populating user %s: %s", outreq.Host, err.Error())
		io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
		return
	}

	if user.Redirect {
		http.Redirect(rw, req, user.Target.String(), http.StatusTemporaryRedirect)
		return
	}

	_, err = validate(user)
	if err != nil {
		log.Printf("error validating user %s: %s", user.IP, err.Error())
		io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
		return
	}

	target := user.Target
	fmt.Printf("proxy to %s\n", target.Host)

	// Smart handling incoming request path/query, example:
	// incoming : foo.com/dir
	// target	: bar.com/base
	// proxy to : bar.com/base/dir
	outreq.URL.Scheme = target.Scheme
	outreq.URL.Host = target.Host
	outreq.URL.Path = singleJoiningSlash(target.Path, outreq.URL.Path)

	// incoming : foo.com/name=arslan
	// target	: bar.com/q=example
	// proxy to : bar.com/q=example&name=arslan
	if target.RawQuery == "" || outreq.URL.RawQuery == "" {
		outreq.URL.RawQuery = target.RawQuery + outreq.URL.RawQuery
	} else {
		outreq.URL.RawQuery = target.RawQuery + "&" + outreq.URL.RawQuery
	}

	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// https://groups.google.com/d/msg/golang-nuts/KBx9pDlvFOc/edt4iad96nwJ
	if websocket {
		fmt.Println("connection via websocket")
		rConn, err := net.Dial("tcp", outreq.URL.Host)
		if err != nil {
			http.Error(rw, "Error contacting backend server.", http.StatusInternalServerError)
			log.Printf("Error dialing websocket backend %s: %v", outreq.URL.Host, err)
			return
		}

		hj, ok := rw.(http.Hijacker)
		if !ok {
			http.Error(rw, "Not a hijacker?", http.StatusInternalServerError)
			return
		}

		conn, _, err := hj.Hijack()
		if err != nil {
			log.Printf("Hijack error: %v", err)
			return
		}
		defer conn.Close()
		defer rConn.Close()

		err = req.Write(rConn)
		if err != nil {
			log.Printf("Error copying request to target: %v", err)
			return
		}

		go p.copyResponse(rConn, conn)
		p.copyResponse(conn, rConn)

	} else {
		go logDomainStat(outreq.Host)
		go logProxyStat(hostname, user.Country)

		transport := p.Transport
		if transport == nil {
			transport = http.DefaultTransport
		}

		// Display error when someone hits the main page
		if hostname == outreq.URL.Host {
			io.WriteString(rw, "{\"err\":\"no such host\"}\n")
			return
		}

		// Remove hop-by-hop headers to the backend.  Especially
		// important is "Connection" because we want a persistent
		// connection, regardless of what the client sent to us.  This
		// is modifying the same underlying map from req (shallow
		// copied above) so we only copy it if necessary.
		copiedHeaders := false
		for _, h := range hopHeaders {
			if outreq.Header.Get(h) != "" {
				if !copiedHeaders {
					outreq.Header = make(http.Header)
					copyHeader(outreq.Header, req.Header)
					copiedHeaders = true
				}
				outreq.Header.Del(h)
			}
		}

		if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
			// If we aren't the first proxy retain prior
			// X-Forwarded-For information as a comma+space
			// separated list and fold multiple headers into one.
			if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
				clientIP = strings.Join(prior, ", ") + ", " + clientIP
			}
			outreq.Header.Set("X-Forwarded-For", clientIP)
		}

		res := new(http.Response)

		rabbitKey, err := lookupRabbitKey(user.Username, user.Servicename, user.Key)
		if err != nil {
			// add :80 if not available
			ok := hasPort(outreq.URL.Host)
			if !ok {
				outreq.URL.Host = addPort(outreq.URL.Host, "80")
			}

			res, err = transport.RoundTrip(outreq)
			if err != nil {
				io.WriteString(rw, fmt.Sprint(err))
				return
			}
		} else {
			fmt.Println("connection via rabbitmq")
			res, err = rabbitTransport(outreq, user, rabbitKey)
			if err != nil {
				log.Printf("rabbit proxy %s", err.Error())
				io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
				return
			}

		}
		defer res.Body.Close()

		copyHeader(rw.Header(), res.Header)
		rw.WriteHeader(res.StatusCode)
		p.copyResponse(rw, res.Body)
	}

}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	io.Copy(dst, src)
}
<commit_msg>kontrolproxy: small change<commit_after>package main

import (
	"fmt"
	"github.com/bradfitz/gomemcache/memcache"
	"github.com/nranchev/go-libGeoIP"
	"io"
	"koding/kontrol/kontrolhelper"
	"koding/kontrol/kontrolproxy/proxyconfig"
	"koding/tools/config"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
)

func init() {
	log.SetPrefix("kontrol-proxy ")
}

type RabbitChannel struct {
	ReplyTo string
	Receive chan []byte
}

var proxyDB *proxyconfig.ProxyConfiguration
var amqpStream *AmqpStream
var connections map[string]RabbitChannel
var geoIP *libgeo.GeoIP
var memCache *memcache.Client
var hostname = kontrolhelper.CustomHostname()

func main() {
	log.Printf("kontrol proxy started ")
	connections = make(map[string]RabbitChannel)

	// open and read from DB
	var err error
	proxyDB, err = proxyconfig.Connect()
	if err != nil {
		log.Fatalf("proxyconfig mongodb connect: %s", err)
	}

	err = proxyDB.AddProxy(hostname)
	if err != nil {
		log.Println(err)
	}

	memCache = memcache.New("127.0.0.1:11211") // used for vm lookup

	// load GeoIP db into memory
	dbFile := "GeoIP.dat"
	geoIP, err = libgeo.Load(dbFile)
	if err != nil {
		log.Printf("load GeoIP.dat: %s\n", err.Error())
	}

	// create amqpStream for rabbitmq proxyieng
	amqpStream = setupAmqp()

	reverseProxy := &ReverseProxy{}
	// http.HandleFunc("/", reverseProxy.ServeHTTP) this works for 1.1
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		reverseProxy.ServeHTTP(w, r)
	})

	port := strconv.Itoa(config.Current.Kontrold.Proxy.Port)
	portssl := strconv.Itoa(config.Current.Kontrold.Proxy.PortSSL)
	sslips := strings.Split(config.Current.Kontrold.Proxy.SSLIPS, ",")

	for _, sslip := range sslips {
		go func(sslip string) {
			err = http.ListenAndServeTLS(sslip+":"+portssl, sslip+"_cert.pem", sslip+"_key.pem", nil)
			if err != nil {
				log.Printf("https mode is disabled. please add cert.pem and key.pem files. %s %s", err, sslip)
			} else {
				log.Printf("https mode is enabled. serving at :%s ...", portssl)
			}
		}(sslip)
	}

	log.Printf("normal mode is enabled. serving at :%s ...", port)
	err = http.ListenAndServe(":"+port, nil)
	if err != nil {
		log.Println(err)
	}
}

/*************************************************
*
*  util functions
*
*  - arslan
*************************************************/
// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func hasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Given a string of the form "host", "port", returns "host:port"
func addPort(host, port string) string {
	if ok := hasPort(host); ok {
		return host
	}

	return host + ":" + port
}

// Check if a server is alive or not
func checkServer(host string) error {
	c, err := net.Dial("tcp", host)
	if err != nil {
		return err
	}
	c.Close()
	return nil
}

/*************************************************
*
*  modified version of go's reverseProxy source code
*  has support for dynamic target url, websockets and amqp
*
*  - arslan
*************************************************/

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// redirect http to https
	if req.TLS == nil && req.Host == "new.koding.com" {
		http.Redirect(rw, req, "https://new.koding.com"+req.RequestURI, http.StatusMovedPermanently)
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	// if connection is of type websocket, hijacking is used instead of http proxy
	websocket := checkWebsocket(outreq)

	user, err := populateUser(outreq)
	if err != nil {
		log.Printf("error populating user %s: %s", outreq.Host, err.Error())
		io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
		return
	}

	if user.Redirect {
		http.Redirect(rw, req, user.Target.String(), http.StatusTemporaryRedirect)
		return
	}

	_, err = validate(user)
	if err != nil {
		log.Printf("error validating user %s: %s", user.IP, err.Error())
		io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
		return
	}

	target := user.Target
	fmt.Printf("proxy to %s\n", target.Host)

	// Smart handling incoming request path/query, example:
	// incoming : foo.com/dir
	// target	: bar.com/base
	// proxy to : bar.com/base/dir
	outreq.URL.Scheme = target.Scheme
	outreq.URL.Host = target.Host
	outreq.URL.Path = singleJoiningSlash(target.Path, outreq.URL.Path)

	// incoming : foo.com/name=arslan
	// target	: bar.com/q=example
	// proxy to : bar.com/q=example&name=arslan
	if target.RawQuery == "" || outreq.URL.RawQuery == "" {
		outreq.URL.RawQuery = target.RawQuery + outreq.URL.RawQuery
	} else {
		outreq.URL.RawQuery = target.RawQuery + "&" + outreq.URL.RawQuery
	}

	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// https://groups.google.com/d/msg/golang-nuts/KBx9pDlvFOc/edt4iad96nwJ
	if websocket {
		fmt.Println("connection via websocket")
		rConn, err := net.Dial("tcp", outreq.URL.Host)
		if err != nil {
			http.Error(rw, "Error contacting backend server.", http.StatusInternalServerError)
			log.Printf("Error dialing websocket backend %s: %v", outreq.URL.Host, err)
			return
		}

		hj, ok := rw.(http.Hijacker)
		if !ok {
			http.Error(rw, "Not a hijacker?", http.StatusInternalServerError)
			return
		}

		conn, _, err := hj.Hijack()
		if err != nil {
			log.Printf("Hijack error: %v", err)
			return
		}
		defer conn.Close()
		defer rConn.Close()

		err = req.Write(rConn)
		if err != nil {
			log.Printf("Error copying request to target: %v", err)
			return
		}

		go p.copyResponse(rConn, conn)
		p.copyResponse(conn, rConn)

	} else {
		go logDomainStat(outreq.Host)
		go logProxyStat(hostname, user.Country)

		transport := p.Transport
		if transport == nil {
			transport = http.DefaultTransport
		}

		// Display error when someone hits the main page
		if hostname == outreq.URL.Host {
			io.WriteString(rw, "{\"err\":\"no such host\"}\n")
			return
		}

		// Remove hop-by-hop headers to the backend.  Especially
		// important is "Connection" because we want a persistent
		// connection, regardless of what the client sent to us.  This
		// is modifying the same underlying map from req (shallow
		// copied above) so we only copy it if necessary.
		copiedHeaders := false
		for _, h := range hopHeaders {
			if outreq.Header.Get(h) != "" {
				if !copiedHeaders {
					outreq.Header = make(http.Header)
					copyHeader(outreq.Header, req.Header)
					copiedHeaders = true
				}
				outreq.Header.Del(h)
			}
		}

		if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
			// If we aren't the first proxy retain prior
			// X-Forwarded-For information as a comma+space
			// separated list and fold multiple headers into one.
			if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
				clientIP = strings.Join(prior, ", ") + ", " + clientIP
			}
			outreq.Header.Set("X-Forwarded-For", clientIP)
		}

		res := new(http.Response)

		rabbitKey, err := lookupRabbitKey(user.Username, user.Servicename, user.Key)
		if err != nil {
			// add :80 if not available
			if !hasPort(outreq.URL.Host) {
				outreq.URL.Host = addPort(outreq.URL.Host, "80")
			}

			res, err = transport.RoundTrip(outreq)
			if err != nil {
				io.WriteString(rw, fmt.Sprint(err))
				return
			}
		} else {
			fmt.Println("connection via rabbitmq")
			res, err = rabbitTransport(outreq, user, rabbitKey)
			if err != nil {
				log.Printf("rabbit proxy %s", err.Error())
				io.WriteString(rw, fmt.Sprintf("{\"err\":\"%s\"}\n", err.Error()))
				return
			}

		}
		defer res.Body.Close()

		copyHeader(rw.Header(), res.Header)
		rw.WriteHeader(res.StatusCode)
		p.copyResponse(rw, res.Body)
	}

}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	io.Copy(dst, src)
}
<|endoftext|>


<commit_before>package models

import (
	"socialapi/workers/common/runner"
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

func TestChannelparticipantNewChannelParticipant(t *testing.T) {
	Convey("while testing New Channel Participant", t, func() {
		Convey("Function call should return ChannelParticipant", func() {
			So(NewChannelParticipant(), ShouldNotBeNil)
		})
	})
}

func TestChannelparticipantGetId(t *testing.T) {
	Convey("While testing get id", t, func() {
		Convey("Initialized struct", func() {
			Convey("should return given id", func() {
				c := ChannelParticipant{Id: 45}
				So(c.GetId(), ShouldEqual, 45)
			})
		})
		Convey("Uninitialized struct", func() {
			Convey("should return 0", func() {
				So(NewChannelParticipant().GetId(), ShouldEqual, 0)
			})
			So(NewChannelParticipant, ShouldNotBeNil)
		})
	})
}

func TestChannelParticipantTableName(t *testing.T) {
	Convey("While getting table name", t, func() {
		Convey("table names should match", func() {
			c := ChannelParticipant{}
			So(c.TableName(), ShouldEqual, "api.channel_participant")
		})
	})
}

func TestChannelParticipantBeforeUpdate(t *testing.T) {
	Convey("While testing before update", t, func() {
		Convey("LastSeenAt should be updated", func() {
			lastSeenAt := time.Now().UTC()

			account, err := createAccount()
			So(err, ShouldBeNil)
			So(account, ShouldNotBeNil)
			So(account.Id, ShouldNotEqual, 0)

			c := NewChannelParticipant()
			c.LastSeenAt = lastSeenAt
			c.AccountId = account.Id

			// call before update
			err = c.BeforeUpdate()

			// make sure err is nil
			So(err, ShouldBeNil)
			// check preset last seen at is not same after calling
			// before update function
			So(c.LastSeenAt, ShouldNotEqual, lastSeenAt)
		})
	})
}

func TestChannelParticipantIsParticipant(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("While testing is participant", t, func() {
		Convey("it should have channel id", func() {
			cp := NewChannelParticipant()

			ip, err := cp.IsParticipant(1123)
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrChannelIdIsNotSet)
			So(ip, ShouldEqual, false)
		})

		Convey("it should return false if account is not exist", func() {
			cp := NewChannelParticipant()
			cp.ChannelId = 120

			ip, err := cp.IsParticipant(112345)
			So(err, ShouldBeNil)
			So(ip, ShouldEqual, false)
		})

		Convey("it should not have error if account is exist", func() {
			c := createNewChannelWithTest()
			So(c.Create(), ShouldBeNil)
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)
			cp := NewChannelParticipant()
			cp.ChannelId = c.Id
			cp.AccountId = acc.Id
			So(cp.Create(), ShouldBeNil)

			_, err := c.AddParticipant(acc.Id)

			ip, err := cp.IsParticipant(acc.Id)
			So(err, ShouldBeNil)
			So(ip, ShouldEqual, true)
			So(cp.StatusConstant, ShouldEqual, ChannelParticipant_STATUS_ACTIVE)
		})
	})
}

func TestChannelParticipantFetchParticipantCount(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("While fetching participant count", t, func() {
		Convey("it should have channel id", func() {
			cp := NewChannelParticipant()

			_, err := cp.FetchParticipantCount()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrChannelIdIsNotSet)
		})

		Convey("it should have zero value if account is not exist in the channel", func() {
			cp := NewChannelParticipant()
			cp.ChannelId = 2442

			fp, err := cp.FetchParticipantCount()
			So(err, ShouldBeNil)
			So(fp, ShouldEqual, 0)
		})

		Convey("it should return participant count successfully", func() {
			c := createNewChannelWithTest()
			So(c.Create(), ShouldBeNil)
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)
			acc1 := createAccountWithTest()
			So(acc1.Create(), ShouldBeNil)
			cp := NewChannelParticipant()
			cp.ChannelId = c.Id
			cp.AccountId = acc.Id
			So(cp.Create(), ShouldBeNil)

			c.AddParticipant(acc.Id)
			c.AddParticipant(acc1.Id)

			ip, err := cp.FetchParticipantCount()
			So(err, ShouldBeNil)
			So(ip, ShouldNotEqual, 2)

		})
	})

}
<commit_msg>social/tests: getAccountId tests are added, in channelparticipant<commit_after>package models

import (
	"socialapi/workers/common/runner"
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

func TestChannelparticipantNewChannelParticipant(t *testing.T) {
	Convey("while testing New Channel Participant", t, func() {
		Convey("Function call should return ChannelParticipant", func() {
			So(NewChannelParticipant(), ShouldNotBeNil)
		})
	})
}

func TestChannelparticipantGetId(t *testing.T) {
	Convey("While testing get id", t, func() {
		Convey("Initialized struct", func() {
			Convey("should return given id", func() {
				c := ChannelParticipant{Id: 45}
				So(c.GetId(), ShouldEqual, 45)
			})
		})
		Convey("Uninitialized struct", func() {
			Convey("should return 0", func() {
				So(NewChannelParticipant().GetId(), ShouldEqual, 0)
			})
			So(NewChannelParticipant, ShouldNotBeNil)
		})
	})
}

func TestChannelParticipantTableName(t *testing.T) {
	Convey("While getting table name", t, func() {
		Convey("table names should match", func() {
			c := ChannelParticipant{}
			So(c.TableName(), ShouldEqual, "api.channel_participant")
		})
	})
}

func TestChannelParticipantBeforeUpdate(t *testing.T) {
	Convey("While testing before update", t, func() {
		Convey("LastSeenAt should be updated", func() {
			lastSeenAt := time.Now().UTC()

			account, err := createAccount()
			So(err, ShouldBeNil)
			So(account, ShouldNotBeNil)
			So(account.Id, ShouldNotEqual, 0)

			c := NewChannelParticipant()
			c.LastSeenAt = lastSeenAt
			c.AccountId = account.Id

			// call before update
			err = c.BeforeUpdate()

			// make sure err is nil
			So(err, ShouldBeNil)
			// check preset last seen at is not same after calling
			// before update function
			So(c.LastSeenAt, ShouldNotEqual, lastSeenAt)
		})
	})
}

func TestChannelParticipantIsParticipant(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("While testing is participant", t, func() {
		Convey("it should have channel id", func() {
			cp := NewChannelParticipant()

			ip, err := cp.IsParticipant(1123)
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrChannelIdIsNotSet)
			So(ip, ShouldEqual, false)
		})

		Convey("it should return false if account is not exist", func() {
			cp := NewChannelParticipant()
			cp.ChannelId = 120

			ip, err := cp.IsParticipant(112345)
			So(err, ShouldBeNil)
			So(ip, ShouldEqual, false)
		})

		Convey("it should not have error if account is exist", func() {
			c := createNewChannelWithTest()
			So(c.Create(), ShouldBeNil)
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)
			cp := NewChannelParticipant()
			cp.ChannelId = c.Id
			cp.AccountId = acc.Id
			So(cp.Create(), ShouldBeNil)

			_, err := c.AddParticipant(acc.Id)

			ip, err := cp.IsParticipant(acc.Id)
			So(err, ShouldBeNil)
			So(ip, ShouldEqual, true)
			So(cp.StatusConstant, ShouldEqual, ChannelParticipant_STATUS_ACTIVE)
		})
	})
}

func TestChannelParticipantFetchParticipantCount(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("While fetching participant count", t, func() {
		Convey("it should have channel id", func() {
			cp := NewChannelParticipant()

			_, err := cp.FetchParticipantCount()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrChannelIdIsNotSet)
		})

		Convey("it should have zero value if account is not exist in the channel", func() {
			cp := NewChannelParticipant()
			cp.ChannelId = 2442

			fp, err := cp.FetchParticipantCount()
			So(err, ShouldBeNil)
			So(fp, ShouldEqual, 0)
		})

		Convey("it should return participant count successfully", func() {
			c := createNewChannelWithTest()
			So(c.Create(), ShouldBeNil)
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)
			acc1 := createAccountWithTest()
			So(acc1.Create(), ShouldBeNil)
			cp := NewChannelParticipant()
			cp.ChannelId = c.Id
			cp.AccountId = acc.Id
			So(cp.Create(), ShouldBeNil)

			c.AddParticipant(acc.Id)
			c.AddParticipant(acc1.Id)

			ip, err := cp.FetchParticipantCount()
			So(err, ShouldBeNil)
			So(ip, ShouldNotEqual, 2)

		})
	})

}

func TestChannelParticipantgetAccountId(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("While getting account id", t, func() {
		Convey("it should have channel id", func() {
			cp := NewChannelParticipant()

			gai, err := cp.getAccountId()
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldContainSubstring, "couldnt find accountId from content")
			So(gai, ShouldEqual, 0)
		})

		Convey("it should return account id if account is exist", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			cp := NewChannelParticipant()
			cp.Id = 1201
			cp.AccountId = acc.Id

			gai, err := cp.getAccountId()
			So(err, ShouldBeNil)
			So(gai, ShouldEqual, acc.Id)
		})

	})

}
<|endoftext|>


<commit_before>/*
Copyright 2015 Workiva, LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package ctrie

import (
	"hash"
	"hash/fnv"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestCtrie(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(nil)

	_, ok := ctrie.Lookup([]byte("foo"))
	assert.False(ok)

	ctrie.Insert([]byte("foo"), "bar")
	val, ok := ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("bar", val)

	ctrie.Insert([]byte("fooooo"), "baz")
	val, ok = ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("bar", val)
	val, ok = ctrie.Lookup([]byte("fooooo"))
	assert.True(ok)
	assert.Equal("baz", val)

	for i := 0; i < 100; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), "blah")
	}
	for i := 0; i < 100; i++ {
		val, ok = ctrie.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal("blah", val)
	}

	val, ok = ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("bar", val)
	ctrie.Insert([]byte("foo"), "qux")
	val, ok = ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("qux", val)

	val, ok = ctrie.Remove([]byte("foo"))
	assert.True(ok)
	assert.Equal("qux", val)

	_, ok = ctrie.Remove([]byte("foo"))
	assert.False(ok)

	val, ok = ctrie.Remove([]byte("fooooo"))
	assert.True(ok)
	assert.Equal("baz", val)

	for i := 0; i < 100; i++ {
		ctrie.Remove([]byte(strconv.Itoa(i)))
	}
}

type mockHash32 struct {
	hash.Hash32
}

func (m *mockHash32) Sum32() uint32 {
	return 0
}

func mockHashFactory() hash.Hash32 {
	return &mockHash32{fnv.New32a()}
}

func TestInsertLNode(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(mockHashFactory)

	for i := 0; i < 10; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}

	for i := 0; i < 10; i++ {
		val, ok := ctrie.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}
	_, ok := ctrie.Lookup([]byte("11"))
	assert.False(ok)

	for i := 0; i < 10; i++ {
		val, ok := ctrie.Remove([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}
}

func TestInsertTNode(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(nil)

	for i := 0; i < 100000; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}

	for i := 0; i < 50000; i++ {
		ctrie.Remove([]byte(strconv.Itoa(i)))
	}

	for i := 0; i < 100000; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}

	for i := 0; i < 100000; i++ {
		val, ok := ctrie.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}
}

func TestConcurrency(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(nil)
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		for i := 0; i < 100000; i++ {
			ctrie.Insert([]byte(strconv.Itoa(i)), i)
		}
		wg.Done()
	}()

	go func() {
		for i := 0; i < 100000; i++ {
			val, ok := ctrie.Lookup([]byte(strconv.Itoa(i)))
			if ok {
				assert.Equal(i, val)
			}
		}
		wg.Done()
	}()

	for i := 0; i < 100000; i++ {
		time.Sleep(5)
		ctrie.Remove([]byte(strconv.Itoa(i)))
	}

	wg.Wait()
}

func TestSnapshot(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(nil)
	for i := 0; i < 100; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}

	snapshot := ctrie.Snapshot()

	// Ensure snapshot contains expected keys.
	for i := 0; i < 100; i++ {
		val, ok := snapshot.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}

	for i := 0; i < 100; i++ {
		ctrie.Remove([]byte(strconv.Itoa(i)))
	}

	// Ensure snapshot was unaffected by removals.
	for i := 0; i < 100; i++ {
		val, ok := snapshot.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}

	ctrie = New(nil)
	for i := 0; i < 100; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	snapshot = ctrie.Snapshot()

	// Ensure snapshot is mutable.
	for i := 0; i < 100; i++ {
		snapshot.Remove([]byte(strconv.Itoa(i)))
	}
	snapshot.Insert([]byte("bat"), "man")

	for i := 0; i < 100; i++ {
		_, ok := snapshot.Lookup([]byte(strconv.Itoa(i)))
		assert.False(ok)
	}
	val, ok := snapshot.Lookup([]byte("bat"))
	assert.True(ok)
	assert.Equal("man", val)

	// Ensure original Ctrie was unaffected.
	for i := 0; i < 100; i++ {
		val, ok := ctrie.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}
	_, ok = ctrie.Lookup([]byte("bat"))
	assert.False(ok)

	snapshot = ctrie.ReadOnlySnapshot()
	for i := 0; i < 100; i++ {
		val, ok := snapshot.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}

	// Ensure read-only snapshots panic on writes.
	defer func() {
		assert.NotNil(recover())
	}()
	snapshot.Remove([]byte("blah"))

	// Ensure snapshots-of-snapshots work as expected.
	snapshot2 := snapshot.Snapshot()
	for i := 0; i < 100; i++ {
		val, ok := snapshot2.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}
	snapshot2.Remove([]byte("0"))
	_, ok = snapshot2.Lookup([]byte("0"))
	assert.False(ok)
	val, ok = snapshot.Lookup([]byte("0"))
	assert.True(ok)
	assert.Equal(0, val)
}

func TestIterator(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(nil)
	for i := 0; i < 10; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	expected := map[string]int{
		"0": 0,
		"1": 1,
		"2": 2,
		"3": 3,
		"4": 4,
		"5": 5,
		"6": 6,
		"7": 7,
		"8": 8,
		"9": 9,
	}

	count := 0
	for entry := range ctrie.Iterator(nil) {
		exp, ok := expected[string(entry.Key)]
		if assert.True(ok) {
			assert.Equal(exp, entry.Value)
		}
		count++
	}
	assert.Equal(len(expected), count)

	// Closing cancel channel should close iterator channel.
	cancel := make(chan struct{})
	iter := ctrie.Iterator(cancel)
	entry := <-iter
	exp, ok := expected[string(entry.Key)]
	if assert.True(ok) {
		assert.Equal(exp, entry.Value)
	}
	close(cancel)
	<-iter // Drain anything already put on the channel
	_, ok = <-iter
	assert.False(ok)
}

func TestSize(t *testing.T) {
	ctrie := New(nil)
	for i := 0; i < 10; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	assert.Equal(t, uint(10), ctrie.Size())
}

func TestClear(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(nil)
	for i := 0; i < 10; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	assert.Equal(uint(10), ctrie.Size())
	snapshot := ctrie.Snapshot()

	ctrie.Clear()

	assert.Equal(uint(0), ctrie.Size())
	assert.Equal(uint(10), snapshot.Size())
}

func BenchmarkInsert(b *testing.B) {
	ctrie := New(nil)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ctrie.Insert([]byte("foo"), 0)
	}
}

func BenchmarkLookup(b *testing.B) {
	numItems := 1000
	ctrie := New(nil)
	for i := 0; i < numItems; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	key := []byte(strconv.Itoa(numItems / 2))
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		ctrie.Lookup(key)
	}
}

func BenchmarkRemove(b *testing.B) {
	numItems := 1000
	ctrie := New(nil)
	for i := 0; i < numItems; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	key := []byte(strconv.Itoa(numItems / 2))
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		ctrie.Remove(key)
	}
}

func BenchmarkSnapshot(b *testing.B) {
	numItems := 1000
	ctrie := New(nil)
	for i := 0; i < numItems; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		ctrie.Snapshot()
	}
}

func BenchmarkReadOnlySnapshot(b *testing.B) {
	numItems := 1000
	ctrie := New(nil)
	for i := 0; i < numItems; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		ctrie.ReadOnlySnapshot()
	}
}
<commit_msg>Make ctrie tests faster<commit_after>/*
Copyright 2015 Workiva, LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package ctrie

import (
	"hash"
	"hash/fnv"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestCtrie(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(nil)

	_, ok := ctrie.Lookup([]byte("foo"))
	assert.False(ok)

	ctrie.Insert([]byte("foo"), "bar")
	val, ok := ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("bar", val)

	ctrie.Insert([]byte("fooooo"), "baz")
	val, ok = ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("bar", val)
	val, ok = ctrie.Lookup([]byte("fooooo"))
	assert.True(ok)
	assert.Equal("baz", val)

	for i := 0; i < 100; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), "blah")
	}
	for i := 0; i < 100; i++ {
		val, ok = ctrie.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal("blah", val)
	}

	val, ok = ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("bar", val)
	ctrie.Insert([]byte("foo"), "qux")
	val, ok = ctrie.Lookup([]byte("foo"))
	assert.True(ok)
	assert.Equal("qux", val)

	val, ok = ctrie.Remove([]byte("foo"))
	assert.True(ok)
	assert.Equal("qux", val)

	_, ok = ctrie.Remove([]byte("foo"))
	assert.False(ok)

	val, ok = ctrie.Remove([]byte("fooooo"))
	assert.True(ok)
	assert.Equal("baz", val)

	for i := 0; i < 100; i++ {
		ctrie.Remove([]byte(strconv.Itoa(i)))
	}
}

type mockHash32 struct {
	hash.Hash32
}

func (m *mockHash32) Sum32() uint32 {
	return 0
}

func mockHashFactory() hash.Hash32 {
	return &mockHash32{fnv.New32a()}
}

func TestInsertLNode(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(mockHashFactory)

	for i := 0; i < 10; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}

	for i := 0; i < 10; i++ {
		val, ok := ctrie.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}
	_, ok := ctrie.Lookup([]byte("11"))
	assert.False(ok)

	for i := 0; i < 10; i++ {
		val, ok := ctrie.Remove([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}
}

func TestInsertTNode(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(nil)

	for i := 0; i < 10000; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}

	for i := 0; i < 5000; i++ {
		ctrie.Remove([]byte(strconv.Itoa(i)))
	}

	for i := 0; i < 10000; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}

	for i := 0; i < 10000; i++ {
		val, ok := ctrie.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}
}

func TestConcurrency(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(nil)
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		for i := 0; i < 10000; i++ {
			ctrie.Insert([]byte(strconv.Itoa(i)), i)
		}
		wg.Done()
	}()

	go func() {
		for i := 0; i < 10000; i++ {
			val, ok := ctrie.Lookup([]byte(strconv.Itoa(i)))
			if ok {
				assert.Equal(i, val)
			}
		}
		wg.Done()
	}()

	for i := 0; i < 10000; i++ {
		time.Sleep(5)
		ctrie.Remove([]byte(strconv.Itoa(i)))
	}

	wg.Wait()
}

func TestSnapshot(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(nil)
	for i := 0; i < 100; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}

	snapshot := ctrie.Snapshot()

	// Ensure snapshot contains expected keys.
	for i := 0; i < 100; i++ {
		val, ok := snapshot.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}

	for i := 0; i < 100; i++ {
		ctrie.Remove([]byte(strconv.Itoa(i)))
	}

	// Ensure snapshot was unaffected by removals.
	for i := 0; i < 100; i++ {
		val, ok := snapshot.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}

	ctrie = New(nil)
	for i := 0; i < 100; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	snapshot = ctrie.Snapshot()

	// Ensure snapshot is mutable.
	for i := 0; i < 100; i++ {
		snapshot.Remove([]byte(strconv.Itoa(i)))
	}
	snapshot.Insert([]byte("bat"), "man")

	for i := 0; i < 100; i++ {
		_, ok := snapshot.Lookup([]byte(strconv.Itoa(i)))
		assert.False(ok)
	}
	val, ok := snapshot.Lookup([]byte("bat"))
	assert.True(ok)
	assert.Equal("man", val)

	// Ensure original Ctrie was unaffected.
	for i := 0; i < 100; i++ {
		val, ok := ctrie.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}
	_, ok = ctrie.Lookup([]byte("bat"))
	assert.False(ok)

	snapshot = ctrie.ReadOnlySnapshot()
	for i := 0; i < 100; i++ {
		val, ok := snapshot.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}

	// Ensure read-only snapshots panic on writes.
	defer func() {
		assert.NotNil(recover())
	}()
	snapshot.Remove([]byte("blah"))

	// Ensure snapshots-of-snapshots work as expected.
	snapshot2 := snapshot.Snapshot()
	for i := 0; i < 100; i++ {
		val, ok := snapshot2.Lookup([]byte(strconv.Itoa(i)))
		assert.True(ok)
		assert.Equal(i, val)
	}
	snapshot2.Remove([]byte("0"))
	_, ok = snapshot2.Lookup([]byte("0"))
	assert.False(ok)
	val, ok = snapshot.Lookup([]byte("0"))
	assert.True(ok)
	assert.Equal(0, val)
}

func TestIterator(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(nil)
	for i := 0; i < 10; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	expected := map[string]int{
		"0": 0,
		"1": 1,
		"2": 2,
		"3": 3,
		"4": 4,
		"5": 5,
		"6": 6,
		"7": 7,
		"8": 8,
		"9": 9,
	}

	count := 0
	for entry := range ctrie.Iterator(nil) {
		exp, ok := expected[string(entry.Key)]
		if assert.True(ok) {
			assert.Equal(exp, entry.Value)
		}
		count++
	}
	assert.Equal(len(expected), count)

	// Closing cancel channel should close iterator channel.
	cancel := make(chan struct{})
	iter := ctrie.Iterator(cancel)
	entry := <-iter
	exp, ok := expected[string(entry.Key)]
	if assert.True(ok) {
		assert.Equal(exp, entry.Value)
	}
	close(cancel)
	<-iter // Drain anything already put on the channel
	_, ok = <-iter
	assert.False(ok)
}

func TestSize(t *testing.T) {
	ctrie := New(nil)
	for i := 0; i < 10; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	assert.Equal(t, uint(10), ctrie.Size())
}

func TestClear(t *testing.T) {
	assert := assert.New(t)
	ctrie := New(nil)
	for i := 0; i < 10; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	assert.Equal(uint(10), ctrie.Size())
	snapshot := ctrie.Snapshot()

	ctrie.Clear()

	assert.Equal(uint(0), ctrie.Size())
	assert.Equal(uint(10), snapshot.Size())
}

func BenchmarkInsert(b *testing.B) {
	ctrie := New(nil)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ctrie.Insert([]byte("foo"), 0)
	}
}

func BenchmarkLookup(b *testing.B) {
	numItems := 1000
	ctrie := New(nil)
	for i := 0; i < numItems; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	key := []byte(strconv.Itoa(numItems / 2))
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		ctrie.Lookup(key)
	}
}

func BenchmarkRemove(b *testing.B) {
	numItems := 1000
	ctrie := New(nil)
	for i := 0; i < numItems; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	key := []byte(strconv.Itoa(numItems / 2))
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		ctrie.Remove(key)
	}
}

func BenchmarkSnapshot(b *testing.B) {
	numItems := 1000
	ctrie := New(nil)
	for i := 0; i < numItems; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		ctrie.Snapshot()
	}
}

func BenchmarkReadOnlySnapshot(b *testing.B) {
	numItems := 1000
	ctrie := New(nil)
	for i := 0; i < numItems; i++ {
		ctrie.Insert([]byte(strconv.Itoa(i)), i)
	}
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		ctrie.ReadOnlySnapshot()
	}
}
<|endoftext|>


<commit_before>package nimble

type Pixel uint32

const (
	alphaShift = 24
	redShift   = 16
	greenShift = 8
	blueShift  = 0
)

func component(c float32) Pixel {
	return Pixel(c*255 + (0.5 - 1.0/256))
}

//  RGB constructs a Pixel from its red, green, and blue components.
func RGB(red float32, green float32, blue float32) Pixel {
	return component(red)<<redShift |
		component(green)<<greenShift |
		component(blue)<<blueShift |
		0xFF<<alphaShift
}

//  RGB constructs a Pixel with equal red, green, and blue components.
func Gray(frac float32) Pixel {
	g := component(frac)
	return g<<redShift | g<<greenShift | g<<blueShift | 0xFF<<alphaShift
}
<commit_msg>Add Black and White constants.<commit_after>package nimble

type Pixel uint32

const (
	alphaShift = 24
	redShift   = 16
	greenShift = 8
	blueShift  = 0
)

func component(c float32) Pixel {
	return Pixel(c*255 + (0.5 - 1.0/256))
}

//  RGB constructs a Pixel from its red, green, and blue components.
func RGB(red float32, green float32, blue float32) Pixel {
	return component(red)<<redShift |
		component(green)<<greenShift |
		component(blue)<<blueShift |
		0xFF<<alphaShift
}

//  RGB constructs a Pixel with equal red, green, and blue components.
func Gray(frac float32) Pixel {
	g := component(frac)
	return g<<redShift | g<<greenShift | g<<blueShift | 0xFF<<alphaShift
}

const Black = Pixel(0xFF000000)
const White = Pixel(0xFFFFFFFF)
<|endoftext|>


<commit_before>package nyb

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	kitty "github.com/ugjka/kittybot"
	"gopkg.in/ugjka/go-tz.v2/tz"
)

const helpMsg = "COMMANDS: '%shny <location>', '%stime <location>', '%snext', '%sprevious', '%sremaining', '%shelp', '%ssource'"

func (bot *Settings) addTriggers() {
	irc := bot.irc

	//Log Notices
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "NOTICE"
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("[NOTICE] " + m.Content)
		},
	})

	//Trigger for !source
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"source")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Reply(m, "https://github.com/ugjka/newyearsbot")
		},
	})

	//Trigger for !help
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"help") ||
				normalize(m.Content) == bot.Prefix+"hny"
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying help...")
			b.Reply(m, fmt.Sprintf(helpMsg, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix))
		},
	})

	//Trigger for !next
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"next")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying next...")
			dur := time.Minute * time.Duration(bot.next.Offset*60)
			if timeNow().UTC().Add(dur).After(target) {
				b.Reply(m, fmt.Sprintf("No more next, %d is here AoE", target.Year()))
				return
			}
			hdur := humanDur(target.Sub(timeNow().UTC().Add(dur)))
			b.Reply(m, fmt.Sprintf("Next New Year in %s in %s",
				hdur, bot.next))
		},
	})

	//Trigger for !previous
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"previous")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying previous...")
			dur := time.Minute * time.Duration(bot.previous.Offset*60)
			hdur := humanDur(timeNow().UTC().Add(dur).Sub(target))
			if bot.previous.Offset == -12 {
				hdur = humanDur(timeNow().UTC().Add(dur).Sub(target.AddDate(-1, 0, 0)))
			}
			b.Reply(m, fmt.Sprintf("Previous New Year %s ago in %s",
				hdur, bot.previous))
		},
	})

	//Trigger for !remaining
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"remaining")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying remaining...")
			plural := "s"
			if bot.remaining == 1 {
				plural = ""
			}
			b.Reply(m, fmt.Sprintf("%d timezone%s remaining. %d%% are in the new year", bot.remaining, (bot.zones - bot.remaining) * 100 / (bot.zones * 100), plural))
		},
	})

	//Trigger for time in location
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"time ")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying time...")
			result, err := bot.time(normalize(m.Content)[len(bot.Prefix)+len("time")+1:])
			if err == errNoZone || err == errNoPlace {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, err.Error())
				return
			}
			if err != nil {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, "Some error occurred!")
				return
			}
			b.Reply(m, result)
		},
	})

	//Trigger for UTC time
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				normalize(m.Content) == bot.Prefix+"time"
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying time...")
			result := "Time is " + time.Now().UTC().Format("Mon Jan 2 15:04:05 -0700 MST 2006")
			b.Reply(m, result)
		},
	})

	//Trigger for new year in location
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"hny ")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			result, err := bot.newYear(normalize(m.Content)[len(bot.Prefix)+len("hny")+1:])
			if err == errNoZone || err == errNoPlace {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, err.Error())
				return
			}
			if err != nil {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, "Some error occurred!")
				return
			}
			b.Reply(m, result)
		},
	})
}

var (
	errNoZone  = errors.New("couldn't get timezone for that location")
	errNoPlace = errors.New("Couldn't find that place")
)

func (bot *Settings) time(location string) (string, error) {
	bot.irc.Info("Querying location: " + location)
	data, err := NominatimFetcher(&bot.Email, &bot.Nominatim, &location)
	if err != nil {
		bot.irc.Warn("Nominatim error: " + err.Error())
		return "", err
	}
	var res NominatimResults
	if err = json.Unmarshal(data, &res); err != nil {
		bot.irc.Warn("Nominatim JSON error: " + err.Error())
		return "", err
	}
	if len(res) == 0 {
		return "", errNoPlace
	}
	p := tz.Point{
		Lat: res[0].Lat,
		Lon: res[0].Lon,
	}
	tzid, err := tz.GetZone(p)
	if err != nil {
		return "", errNoZone
	}
	zone, err := time.LoadLocation(tzid[0])
	if err != nil {
		return "", errNoZone
	}
	address := res[0].DisplayName
	msg := fmt.Sprintf("Time in %s is %s", address, time.Now().In(zone).Format("Mon Jan 2 15:04:05 -0700 MST 2006"))
	return msg, nil
}

func (bot *Settings) newYear(location string) (string, error) {
	bot.irc.Info("Querying location: " + location)
	data, err := NominatimFetcher(&bot.Email, &bot.Nominatim, &location)
	if err != nil {
		bot.irc.Warn("Nominatim error: " + err.Error())
		return "", err
	}
	var res NominatimResults
	if err = json.Unmarshal(data, &res); err != nil {
		bot.irc.Warn("Nominatim JSON error: " + err.Error())
		return "", err
	}
	if len(res) == 0 {
		return "", errNoPlace
	}
	p := tz.Point{
		Lat: res[0].Lat,
		Lon: res[0].Lon,
	}
	tzid, err := tz.GetZone(p)
	if err != nil {
		return "", errNoZone
	}
	zone, err := time.LoadLocation(tzid[0])
	if err != nil {
		return "", errNoZone
	}
	offset := zoneOffset(target, zone)
	address := res[0].DisplayName
	if timeNow().UTC().Add(offset).Before(target) {
		hdur := humanDur(target.Sub(timeNow().UTC().Add(offset)))
		const newYearFutureMsg = "New Year in %s will happen in %s"
		return fmt.Sprintf(newYearFutureMsg, address, hdur), nil
	}
	hdur := humanDur(timeNow().UTC().Add(offset).Sub(target))
	const newYearPastMsg = "New Year in %s happened %s ago"
	return fmt.Sprintf(newYearPastMsg, address, hdur), nil
}
<commit_msg>Percentage as int out of 100<commit_after>package nyb

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	kitty "github.com/ugjka/kittybot"
	"gopkg.in/ugjka/go-tz.v2/tz"
)

const helpMsg = "COMMANDS: '%shny <location>', '%stime <location>', '%snext', '%sprevious', '%sremaining', '%shelp', '%ssource'"

func (bot *Settings) addTriggers() {
	irc := bot.irc

	//Log Notices
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "NOTICE"
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("[NOTICE] " + m.Content)
		},
	})

	//Trigger for !source
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"source")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Reply(m, "https://github.com/ugjka/newyearsbot")
		},
	})

	//Trigger for !help
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"help") ||
				normalize(m.Content) == bot.Prefix+"hny"
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying help...")
			b.Reply(m, fmt.Sprintf(helpMsg, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix))
		},
	})

	//Trigger for !next
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"next")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying next...")
			dur := time.Minute * time.Duration(bot.next.Offset*60)
			if timeNow().UTC().Add(dur).After(target) {
				b.Reply(m, fmt.Sprintf("No more next, %d is here AoE", target.Year()))
				return
			}
			hdur := humanDur(target.Sub(timeNow().UTC().Add(dur)))
			b.Reply(m, fmt.Sprintf("Next New Year in %s in %s",
				hdur, bot.next))
		},
	})

	//Trigger for !previous
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"previous")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying previous...")
			dur := time.Minute * time.Duration(bot.previous.Offset*60)
			hdur := humanDur(timeNow().UTC().Add(dur).Sub(target))
			if bot.previous.Offset == -12 {
				hdur = humanDur(timeNow().UTC().Add(dur).Sub(target.AddDate(-1, 0, 0)))
			}
			b.Reply(m, fmt.Sprintf("Previous New Year %s ago in %s",
				hdur, bot.previous))
		},
	})

	//Trigger for !remaining
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"remaining")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying remaining...")
			plural := "s"
			if bot.remaining == 1 {
				plural = ""
			}
			pct := (len(bot.zones) - bot.remaining) * 100 / (len(bot.zones) * 100)
			b.Reply(m, fmt.Sprintf("%d timezone%s remaining. %d%% are in the new year", bot.remaining, plural, pct))
		},
	})

	//Trigger for time in location
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"time ")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying time...")
			result, err := bot.time(normalize(m.Content)[len(bot.Prefix)+len("time")+1:])
			if err == errNoZone || err == errNoPlace {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, err.Error())
				return
			}
			if err != nil {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, "Some error occurred!")
				return
			}
			b.Reply(m, result)
		},
	})

	//Trigger for UTC time
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				normalize(m.Content) == bot.Prefix+"time"
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying time...")
			result := "Time is " + time.Now().UTC().Format("Mon Jan 2 15:04:05 -0700 MST 2006")
			b.Reply(m, result)
		},
	})

	//Trigger for new year in location
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"hny ")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			result, err := bot.newYear(normalize(m.Content)[len(bot.Prefix)+len("hny")+1:])
			if err == errNoZone || err == errNoPlace {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, err.Error())
				return
			}
			if err != nil {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, "Some error occurred!")
				return
			}
			b.Reply(m, result)
		},
	})
}

var (
	errNoZone  = errors.New("couldn't get timezone for that location")
	errNoPlace = errors.New("Couldn't find that place")
)

func (bot *Settings) time(location string) (string, error) {
	bot.irc.Info("Querying location: " + location)
	data, err := NominatimFetcher(&bot.Email, &bot.Nominatim, &location)
	if err != nil {
		bot.irc.Warn("Nominatim error: " + err.Error())
		return "", err
	}
	var res NominatimResults
	if err = json.Unmarshal(data, &res); err != nil {
		bot.irc.Warn("Nominatim JSON error: " + err.Error())
		return "", err
	}
	if len(res) == 0 {
		return "", errNoPlace
	}
	p := tz.Point{
		Lat: res[0].Lat,
		Lon: res[0].Lon,
	}
	tzid, err := tz.GetZone(p)
	if err != nil {
		return "", errNoZone
	}
	zone, err := time.LoadLocation(tzid[0])
	if err != nil {
		return "", errNoZone
	}
	address := res[0].DisplayName
	msg := fmt.Sprintf("Time in %s is %s", address, time.Now().In(zone).Format("Mon Jan 2 15:04:05 -0700 MST 2006"))
	return msg, nil
}

func (bot *Settings) newYear(location string) (string, error) {
	bot.irc.Info("Querying location: " + location)
	data, err := NominatimFetcher(&bot.Email, &bot.Nominatim, &location)
	if err != nil {
		bot.irc.Warn("Nominatim error: " + err.Error())
		return "", err
	}
	var res NominatimResults
	if err = json.Unmarshal(data, &res); err != nil {
		bot.irc.Warn("Nominatim JSON error: " + err.Error())
		return "", err
	}
	if len(res) == 0 {
		return "", errNoPlace
	}
	p := tz.Point{
		Lat: res[0].Lat,
		Lon: res[0].Lon,
	}
	tzid, err := tz.GetZone(p)
	if err != nil {
		return "", errNoZone
	}
	zone, err := time.LoadLocation(tzid[0])
	if err != nil {
		return "", errNoZone
	}
	offset := zoneOffset(target, zone)
	address := res[0].DisplayName
	if timeNow().UTC().Add(offset).Before(target) {
		hdur := humanDur(target.Sub(timeNow().UTC().Add(offset)))
		const newYearFutureMsg = "New Year in %s will happen in %s"
		return fmt.Sprintf(newYearFutureMsg, address, hdur), nil
	}
	hdur := humanDur(timeNow().UTC().Add(offset).Sub(target))
	const newYearPastMsg = "New Year in %s happened %s ago"
	return fmt.Sprintf(newYearPastMsg, address, hdur), nil
}
<|endoftext|>


<commit_before>package apns

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"time"
)

const (
	PriorityImmediate     = 10
	PriorityPowerConserve = 5
)

const (
	commandID = 2

	// Items IDs
	deviceTokenItemID            = 1
	payloadItemID                = 2
	notificationIdentifierItemID = 3
	expirationDateItemID         = 4
	priorityItemID               = 5

	// Item lengths
	deviceTokenItemLength            = 32
	notificationIdentifierItemLength = 4
	expirationDateItemLength         = 4
	priorityItemLength               = 1
)

type NotificationResult struct {
	Notif Notification
	Err   Error
}

type Alert struct {
	// Do not add fields without updating the implementation of isZero.
	Body         string   `json:"body,omitempty"`
	Title        string   `json:"title,omitempty"`
	Action       string   `json:"action,omitempty"`
	LocKey       string   `json:"loc-key,omitempty"`
	LocArgs      []string `json:"loc-args,omitempty"`
	ActionLocKey string   `json:"action-loc-key,omitempty"`
	LaunchImage  string   `json:"launch-image,omitempty"`
}

func (a *Alert) isSimple() bool {
	return len(a.Title) == 0 && len(a.Action) == 0 && len(a.LocKey) == 0 && len(a.LocArgs) == 0 && len(a.ActionLocKey) == 0 && len(a.LaunchImage) == 0
}

func (a *Alert) isZero() bool {
	return a.isSimple() && len(a.Body) == 0
}

type APS struct {
	Alert            Alert
	Badge            *int // 0 to clear notifications, nil to leave as is.
	Sound            string
	ContentAvailable int
	URLArgs          []string
	Category         string // requires iOS 8+
}

func (aps APS) MarshalJSON() ([]byte, error) {
	data := make(map[string]interface{})

	if !aps.Alert.isZero() {
		if aps.Alert.isSimple() {
			data["alert"] = aps.Alert.Body
		} else {
			data["alert"] = aps.Alert
		}
	}
	if aps.Badge != nil {
		data["badge"] = aps.Badge
	}
	if aps.Sound != "" {
		data["sound"] = aps.Sound
	}
	if aps.ContentAvailable != 0 {
		data["content-available"] = aps.ContentAvailable
	}
	if aps.Category != "" {
		data["category"] = aps.Category
	}
	if aps.URLArgs != nil && len(aps.URLArgs) != 0 {
		data["url-args"] = aps.URLArgs
	}

	return json.Marshal(data)
}

type Payload struct {
	APS APS
	// MDM for mobile device management
	MDM          string
	customValues map[string]interface{}
}

func (p *Payload) MarshalJSON() ([]byte, error) {
	if len(p.MDM) != 0 {
		p.customValues["mdm"] = p.MDM
	} else {
		p.customValues["aps"] = p.APS
	}

	return json.Marshal(p.customValues)
}

func (p *Payload) SetCustomValue(key string, value interface{}) error {
	if key == "aps" {
		return errors.New("cannot assign a custom APS value in payload")
	}

	p.customValues[key] = value

	return nil
}

type Notification struct {
	ID          string
	DeviceToken string
	Identifier  uint32
	Expiration  *time.Time
	Priority    int
	Payload     *Payload
}

func NewNotification() Notification {
	return Notification{Payload: NewPayload()}
}

func NewPayload() *Payload {
	return &Payload{customValues: map[string]interface{}{}}
}

func (n Notification) ToBinary() ([]byte, error) {
	b := []byte{}

	binTok, err := hex.DecodeString(n.DeviceToken)
	if err != nil {
		return b, fmt.Errorf("convert token to hex error: %s", err)
	}

	j, _ := json.Marshal(n.Payload)

	buf := bytes.NewBuffer(b)

	// Token
	binary.Write(buf, binary.BigEndian, uint8(deviceTokenItemID))
	binary.Write(buf, binary.BigEndian, uint16(deviceTokenItemLength))
	binary.Write(buf, binary.BigEndian, binTok)

	// Payload
	binary.Write(buf, binary.BigEndian, uint8(payloadItemID))
	binary.Write(buf, binary.BigEndian, uint16(len(j)))
	binary.Write(buf, binary.BigEndian, j)

	// Identifier
	binary.Write(buf, binary.BigEndian, uint8(notificationIdentifierItemID))
	binary.Write(buf, binary.BigEndian, uint16(notificationIdentifierItemLength))
	binary.Write(buf, binary.BigEndian, uint32(n.Identifier))

	// Expiry
	binary.Write(buf, binary.BigEndian, uint8(expirationDateItemID))
	binary.Write(buf, binary.BigEndian, uint16(expirationDateItemLength))
	if n.Expiration == nil {
		binary.Write(buf, binary.BigEndian, uint32(0))
	} else {
		binary.Write(buf, binary.BigEndian, uint32(n.Expiration.Unix()))
	}

	// Priority
	binary.Write(buf, binary.BigEndian, uint8(priorityItemID))
	binary.Write(buf, binary.BigEndian, uint16(priorityItemLength))
	binary.Write(buf, binary.BigEndian, uint8(n.Priority))

	framebuf := bytes.NewBuffer([]byte{})
	binary.Write(framebuf, binary.BigEndian, uint8(commandID))
	binary.Write(framebuf, binary.BigEndian, uint32(buf.Len()))
	binary.Write(framebuf, binary.BigEndian, buf.Bytes())

	return framebuf.Bytes(), nil
}
<commit_msg>Add support for email push notifications<commit_after>package apns

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"time"
)

const (
	PriorityImmediate     = 10
	PriorityPowerConserve = 5
)

const (
	commandID = 2

	// Items IDs
	deviceTokenItemID            = 1
	payloadItemID                = 2
	notificationIdentifierItemID = 3
	expirationDateItemID         = 4
	priorityItemID               = 5

	// Item lengths
	deviceTokenItemLength            = 32
	notificationIdentifierItemLength = 4
	expirationDateItemLength         = 4
	priorityItemLength               = 1
)

type NotificationResult struct {
	Notif Notification
	Err   Error
}

type Alert struct {
	// Do not add fields without updating the implementation of isZero.
	Body         string   `json:"body,omitempty"`
	Title        string   `json:"title,omitempty"`
	Action       string   `json:"action,omitempty"`
	LocKey       string   `json:"loc-key,omitempty"`
	LocArgs      []string `json:"loc-args,omitempty"`
	ActionLocKey string   `json:"action-loc-key,omitempty"`
	LaunchImage  string   `json:"launch-image,omitempty"`
}

func (a *Alert) isSimple() bool {
	return len(a.Title) == 0 && len(a.Action) == 0 && len(a.LocKey) == 0 && len(a.LocArgs) == 0 && len(a.ActionLocKey) == 0 && len(a.LaunchImage) == 0
}

func (a *Alert) isZero() bool {
	return a.isSimple() && len(a.Body) == 0
}

type APS struct {
	Alert            Alert
	Badge            *int // 0 to clear notifications, nil to leave as is.
	Sound            string
	ContentAvailable int
	URLArgs          []string
	Category         string // requires iOS 8+
	AccountId        string // for email push notifications
}

func (aps APS) MarshalJSON() ([]byte, error) {
	data := make(map[string]interface{})

	if !aps.Alert.isZero() {
		if aps.Alert.isSimple() {
			data["alert"] = aps.Alert.Body
		} else {
			data["alert"] = aps.Alert
		}
	}
	if aps.Badge != nil {
		data["badge"] = aps.Badge
	}
	if aps.Sound != "" {
		data["sound"] = aps.Sound
	}
	if aps.ContentAvailable != 0 {
		data["content-available"] = aps.ContentAvailable
	}
	if aps.Category != "" {
		data["category"] = aps.Category
	}
	if aps.URLArgs != nil && len(aps.URLArgs) != 0 {
		data["url-args"] = aps.URLArgs
	}
	if aps.AccountId != "" {
		data["account-id"] = aps.AccountId
	}

	return json.Marshal(data)
}

type Payload struct {
	APS APS
	// MDM for mobile device management
	MDM          string
	customValues map[string]interface{}
}

func (p *Payload) MarshalJSON() ([]byte, error) {
	if len(p.MDM) != 0 {
		p.customValues["mdm"] = p.MDM
	} else {
		p.customValues["aps"] = p.APS
	}

	return json.Marshal(p.customValues)
}

func (p *Payload) SetCustomValue(key string, value interface{}) error {
	if key == "aps" {
		return errors.New("cannot assign a custom APS value in payload")
	}

	p.customValues[key] = value

	return nil
}

type Notification struct {
	ID          string
	DeviceToken string
	Identifier  uint32
	Expiration  *time.Time
	Priority    int
	Payload     *Payload
}

func NewNotification() Notification {
	return Notification{Payload: NewPayload()}
}

func NewPayload() *Payload {
	return &Payload{customValues: map[string]interface{}{}}
}

func (n Notification) ToBinary() ([]byte, error) {
	b := []byte{}

	binTok, err := hex.DecodeString(n.DeviceToken)
	if err != nil {
		return b, fmt.Errorf("convert token to hex error: %s", err)
	}

	j, _ := json.Marshal(n.Payload)

	buf := bytes.NewBuffer(b)

	// Token
	binary.Write(buf, binary.BigEndian, uint8(deviceTokenItemID))
	binary.Write(buf, binary.BigEndian, uint16(deviceTokenItemLength))
	binary.Write(buf, binary.BigEndian, binTok)

	// Payload
	binary.Write(buf, binary.BigEndian, uint8(payloadItemID))
	binary.Write(buf, binary.BigEndian, uint16(len(j)))
	binary.Write(buf, binary.BigEndian, j)

	// Identifier
	binary.Write(buf, binary.BigEndian, uint8(notificationIdentifierItemID))
	binary.Write(buf, binary.BigEndian, uint16(notificationIdentifierItemLength))
	binary.Write(buf, binary.BigEndian, uint32(n.Identifier))

	// Expiry
	binary.Write(buf, binary.BigEndian, uint8(expirationDateItemID))
	binary.Write(buf, binary.BigEndian, uint16(expirationDateItemLength))
	if n.Expiration == nil {
		binary.Write(buf, binary.BigEndian, uint32(0))
	} else {
		binary.Write(buf, binary.BigEndian, uint32(n.Expiration.Unix()))
	}

	// Priority
	binary.Write(buf, binary.BigEndian, uint8(priorityItemID))
	binary.Write(buf, binary.BigEndian, uint16(priorityItemLength))
	binary.Write(buf, binary.BigEndian, uint8(n.Priority))

	framebuf := bytes.NewBuffer([]byte{})
	binary.Write(framebuf, binary.BigEndian, uint8(commandID))
	binary.Write(framebuf, binary.BigEndian, uint32(buf.Len()))
	binary.Write(framebuf, binary.BigEndian, buf.Bytes())

	return framebuf.Bytes(), nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, Oracle and/or its affiliates. All rights reserved.

package provider

import (
	"log"
)

const Version = "3.7.0"

func PrintVersion() {
	log.Printf("[INFO] terraform-provider-oci %s\n", Version)
}
<commit_msg>Update version to 3.8.0<commit_after>// Copyright (c) 2017, Oracle and/or its affiliates. All rights reserved.

package provider

import (
	"log"
)

const Version = "3.8.0"

func PrintVersion() {
	log.Printf("[INFO] terraform-provider-oci %s\n", Version)
}
<|endoftext|>


<commit_before>package notificator

import (
	"fmt"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
)

type Options struct {
	DefaultIcon string
	AppName     string
}

const (
	UR_NORMAL   = "normal"
	UR_CRITICAL = "critical"
)

type notifier interface {
	push(title string, text string, iconPath string) *exec.Cmd
	pushCritical(title string, text string, iconPath string) *exec.Cmd
}

type Notificator struct {
	notifier    notifier
	defaultIcon string
}

func (n Notificator) Push(title string, text string, iconPath string, urgency string) error {
	icon := n.defaultIcon

	if iconPath != "" {
		icon = iconPath
	}

	if urgency == UR_CRITICAL {
		return n.notifier.pushCritical(title, text, icon).Run()
	}

	return n.notifier.push(title, text, icon).Run()

}

type osxNotificator struct {
	AppName string
}

func (o osxNotificator) push(title string, text string, iconPath string) *exec.Cmd {

	// Checks if terminal-notifier exists, and is accessible.

	term_notif := CheckTermNotif()
	os_version_check := CheckMacOSVersion()

	// if terminal-notifier exists, use it.
	// else, fall back to osascript. (Mavericks and later.)

	if term_notif == true {
		return exec.Command("terminal-notifier", "-title", o.AppName, "-message", text, "-subtitle", title)
	} else if os_version_check == true {
		notification := fmt.Sprintf("display notification \"%s\" with title \"%s\" subtitle \"%s\"", text, o.AppName, title)
		return exec.Command("osascript", "-e", notification)
	}

	// finally falls back to growlnotify.

	return exec.Command("growlnotify", "-n", o.AppName, "--image", iconPath, "-m", title)
}

// Causes the notification to stick around until clicked.
func (o osxNotificator) pushCritical(title string, text string, iconPath string) *exec.Cmd {

	// same function as above...

	term_notif := CheckTermNotif()
	os_version_check := CheckMacOSVersion()

	if term_notif == true {
		// timeout set to 30 seconds, to show the importance of the notification
		return exec.Command("terminal-notifier", "-title", o.AppName, "-message", title, "-subtitle", title, "-timeout", "30")
	} else if os_version_check == true {
		notification := fmt.Sprintf("display notification \"%s\" with title \"%s\" subtitle \"%s\"", text, o.AppName, title)
		return exec.Command("osascript", "-e", notification)
	}

	return exec.Command("growlnotify", "-n", o.AppName, "--image", iconPath, "-m", title)

}

type linuxNotificator struct{}

func (l linuxNotificator) push(title string, text string, iconPath string) *exec.Cmd {
	return exec.Command("notify-send", "-i", iconPath, title, text)
}

// Causes the notification to stick around until clicked.
func (l linuxNotificator) pushCritical(title string, text string, iconPath string) *exec.Cmd {
	return exec.Command("notify-send", "-i", iconPath, title, text, "-u", "critical")
}

type windowsNotificator struct{}

func (w windowsNotificator) push(title string, text string, iconPath string) *exec.Cmd {
	return exec.Command("growlnotify", "/i:", iconPath, "/t:", title, text)
}

// Causes the notification to stick around until clicked.
func (w windowsNotificator) pushCritical(title string, text string, iconPath string) *exec.Cmd {
	return exec.Command("notify-send", "-i", iconPath, title, text, "/s", "true", "/p", "2")
}

func New(o Options) *Notificator {

	var Notifier notifier

	switch runtime.GOOS {

	case "darwin":
		Notifier = osxNotificator{AppName: o.AppName}
	case "linux":
		Notifier = linuxNotificator{}
	case "windows":
		Notifier = windowsNotificator{}

	}

	return &Notificator{notifier: Notifier, defaultIcon: o.DefaultIcon}
}

// Helper function for macOS

func CheckTermNotif() bool {
	// Checks if terminal-notifier exists, and is accessible.

	check_term_notif := exec.Command("which", "terminal-notifier")
	err := check_term_notif.Start()

	if err != nil {
		return false
	} else {
		err = check_term_notif.Wait()
		if err != nil {
			return false
		}
	}
	// no error, so return true. (terminal-notifier exists)
	return true
}

func CheckMacOSVersion() bool {
	// Checks if the version of macOS is 10.9 or Higher (osascript support for notifications.)

	cmd := exec.Command("sw_vers", "-productVersion")
	check, _ := cmd.Output()

	version := strings.Split(string(check), ".")

	// semantic versioning of macOS

	major, _ := strconv.Atoi(version[0])
	minor, _ := strconv.Atoi(version[1])

	if major < 10 {
		return false
	} else if major == 10 && minor < 9 {
		return false
	} else {
		return true
	}
}
<commit_msg>Change title to text<commit_after>package notificator

import (
	"fmt"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
)

type Options struct {
	DefaultIcon string
	AppName     string
}

const (
	UR_NORMAL   = "normal"
	UR_CRITICAL = "critical"
)

type notifier interface {
	push(title string, text string, iconPath string) *exec.Cmd
	pushCritical(title string, text string, iconPath string) *exec.Cmd
}

type Notificator struct {
	notifier    notifier
	defaultIcon string
}

func (n Notificator) Push(title string, text string, iconPath string, urgency string) error {
	icon := n.defaultIcon

	if iconPath != "" {
		icon = iconPath
	}

	if urgency == UR_CRITICAL {
		return n.notifier.pushCritical(title, text, icon).Run()
	}

	return n.notifier.push(title, text, icon).Run()

}

type osxNotificator struct {
	AppName string
}

func (o osxNotificator) push(title string, text string, iconPath string) *exec.Cmd {

	// Checks if terminal-notifier exists, and is accessible.

	term_notif := CheckTermNotif()
	os_version_check := CheckMacOSVersion()

	// if terminal-notifier exists, use it.
	// else, fall back to osascript. (Mavericks and later.)

	if term_notif == true {
		return exec.Command("terminal-notifier", "-title", o.AppName, "-message", text, "-subtitle", title)
	} else if os_version_check == true {
		notification := fmt.Sprintf("display notification \"%s\" with title \"%s\" subtitle \"%s\"", text, o.AppName, title)
		return exec.Command("osascript", "-e", notification)
	}

	// finally falls back to growlnotify.

	return exec.Command("growlnotify", "-n", o.AppName, "--image", iconPath, "-m", title)
}

// Causes the notification to stick around until clicked.
func (o osxNotificator) pushCritical(title string, text string, iconPath string) *exec.Cmd {

	// same function as above...

	term_notif := CheckTermNotif()
	os_version_check := CheckMacOSVersion()

	if term_notif == true {
		// timeout set to 30 seconds, to show the importance of the notification
		return exec.Command("terminal-notifier", "-title", o.AppName, "-message", text, "-subtitle", title, "-timeout", "30")
	} else if os_version_check == true {
		notification := fmt.Sprintf("display notification \"%s\" with title \"%s\" subtitle \"%s\"", text, o.AppName, title)
		return exec.Command("osascript", "-e", notification)
	}

	return exec.Command("growlnotify", "-n", o.AppName, "--image", iconPath, "-m", title)

}

type linuxNotificator struct{}

func (l linuxNotificator) push(title string, text string, iconPath string) *exec.Cmd {
	return exec.Command("notify-send", "-i", iconPath, title, text)
}

// Causes the notification to stick around until clicked.
func (l linuxNotificator) pushCritical(title string, text string, iconPath string) *exec.Cmd {
	return exec.Command("notify-send", "-i", iconPath, title, text, "-u", "critical")
}

type windowsNotificator struct{}

func (w windowsNotificator) push(title string, text string, iconPath string) *exec.Cmd {
	return exec.Command("growlnotify", "/i:", iconPath, "/t:", title, text)
}

// Causes the notification to stick around until clicked.
func (w windowsNotificator) pushCritical(title string, text string, iconPath string) *exec.Cmd {
	return exec.Command("notify-send", "-i", iconPath, title, text, "/s", "true", "/p", "2")
}

func New(o Options) *Notificator {

	var Notifier notifier

	switch runtime.GOOS {

	case "darwin":
		Notifier = osxNotificator{AppName: o.AppName}
	case "linux":
		Notifier = linuxNotificator{}
	case "windows":
		Notifier = windowsNotificator{}

	}

	return &Notificator{notifier: Notifier, defaultIcon: o.DefaultIcon}
}

// Helper function for macOS

func CheckTermNotif() bool {
	// Checks if terminal-notifier exists, and is accessible.

	check_term_notif := exec.Command("which", "terminal-notifier")
	err := check_term_notif.Start()

	if err != nil {
		return false
	} else {
		err = check_term_notif.Wait()
		if err != nil {
			return false
		}
	}
	// no error, so return true. (terminal-notifier exists)
	return true
}

func CheckMacOSVersion() bool {
	// Checks if the version of macOS is 10.9 or Higher (osascript support for notifications.)

	cmd := exec.Command("sw_vers", "-productVersion")
	check, _ := cmd.Output()

	version := strings.Split(string(check), ".")

	// semantic versioning of macOS

	major, _ := strconv.Atoi(version[0])
	minor, _ := strconv.Atoi(version[1])

	if major < 10 {
		return false
	} else if major == 10 && minor < 9 {
		return false
	} else {
		return true
	}
}
<|endoftext|>