<commit_before>package perfm

import (
	"fmt"
	"math"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	hist "github.com/arthurkiller/perfm/histogram"
)

// Job give out a job for parallel call
// 1. start workers
// 		1. workers call job.Copy()
// 		2. for-loop do
// 			* job.Pre()
// 			* job.Do()
// 		3. after for-loop call job.After()
// 2. caculate the summary
type Job interface {
	// Copy will copy a job for parallel call
	Copy() (Job, error)
	// Pre will called before do
	Pre() error
	// Do contains the core job here
	Do() error
	// After contains the clean job after job done
	After()
}

//PerfMonitor define the atcion about perfmonitor
type PerfMonitor interface {
	Regist(Job) //regist the job to perfm
	Start()     //start the perf monitor
	Wait()      //wait for the benchmark done
}

type perfmonitor struct {
	Sum   float64 //Sum of the per request cost
	Stdev float64 //Standard Deviation
	Mean  float64 //Mean about distribution
	Total int64   //total request by count

	Config                            //configration for perfm
	done           chan int           //stop the perfm
	startTime      time.Time          //keep the start time
	timer          <-chan time.Time   //the frequency sampling timer
	collector      chan time.Duration //get the request cost from every done()
	errCount       int64              //error counter count error request
	localCount     int                //count for the number in the sampling times
	localTimeCount time.Duration      //count for the sampling time total costs
	buffer         chan int64         //buffer the test time for latter add to the historgam
	histogram      hist.Histogram     //used to print the histogram
	wg             sync.WaitGroup     //wait group to block the stop and sync the work thread

	//job implement benchmark job
	//error occoured in job.Do will be collected
	job Job
}

func New(options ...Options) PerfMonitor { return &perfmonitor{Config: newConfig(options...)} }

// Regist a job into perfmonitor fro benchmark
func (p *perfmonitor) Regist(job Job) {
	p.timer = time.Tick(time.Second * time.Duration(p.Frequency))
	p.collector = make(chan time.Duration, p.BufferSize)
	p.histogram = hist.NewHistogram(p.BinsNumber)
	p.done = make(chan int, 0)
	p.buffer = make(chan int64, 100000000)
	p.wg = sync.WaitGroup{}
	p.job = job

	p.Sum = 0
	p.Stdev = 0
	p.Mean = 0
	p.Total = 0
	p.errCount = 0
	p.localCount = 0
	p.localTimeCount = 0
}

// Start the benchmark with given arguments on regisit
func (p *perfmonitor) Start() {
	if p.job == nil {
		panic("error job does not registed yet")
	}
	var localwg sync.WaitGroup

	// If job implement descripetion as Stringer
	if _, ok := p.job.(fmt.Stringer); ok {
		fmt.Println(p.job)
	}
	fmt.Println("===============================================")

	p.wg.Add(1)
	go func() {
		p.startTime = time.Now()
		var cost time.Duration
		for {
			select {
			case cost = <-p.collector:
				p.localCount++
				p.localTimeCount += cost
				p.buffer <- int64(cost)
			case <-p.timer:
				if p.localCount == 0 {
					continue
				}
				if !p.NoPrint {
					fmt.Printf("%s \t  Qps: %d \t  Avg Latency: %.3fms\n", time.Now().Format("15:04:05.000"),
						p.localCount, float64(p.localTimeCount.Nanoseconds()/int64(p.localCount))/1000000)
				}
				p.localCount = 0
				p.localTimeCount = 0
			case <-p.done:
				localwg.Wait()
				close(p.collector)
				for cost := range p.collector {
					p.localCount++
					p.localTimeCount += cost
					p.buffer <- int64(cost)
				}
				if !p.NoPrint {
					fmt.Printf("%s \t  Qps: %d \t  Avg Latency: %.3fms\n", time.Now().Format("15:04:05.000"),
						p.localCount, float64(p.localTimeCount.Nanoseconds()/int64(p.localCount))/1000000)
				}
				close(p.buffer)
				p.wg.Done()
				return
			}
		}
	}()

	if p.Number > 0 {
		// in total request module
		sum := int64(p.Number)
		for i := 0; i < p.Parallel; i++ {
			localwg.Add(1)
			go func() {
				defer localwg.Done()
				var err error
				job, err := p.job.Copy()
				if err != nil {
					fmt.Println("error in do copy", err)
					return
				}
				defer job.After()
				var start time.Time
				var l int64
				for {
					if l = atomic.AddInt64(&p.Total, 1); l > sum {
						if l == sum+1 {
							close(p.done)
						}
						// check if the request reach the goal
						return
					}
					if err = job.Pre(); err != nil {
						fmt.Println("error in do pre job", err)
						return
					}
					start = time.Now()
					err = job.Do()
					p.collector <- time.Since(start)
					if err != nil {
						atomic.AddInt64(&p.errCount, 1)
					}
				}
			}()
		}
	} else {
		// in test duration module
		// start all the worker and do job till cancelled
		starter := make(chan struct{})
		for i := 0; i < p.Parallel; i++ {
			localwg.Add(1)
			go func() {
				defer localwg.Done()
				var err error
				job, err := p.job.Copy()
				if err != nil {
					fmt.Println("error in do copy", err)
					return
				}
				defer job.After()
				var start time.Time
				<-starter
				for {
					select {
					case <-p.done:
						return
					default:
						if err = job.Pre(); err != nil {
							fmt.Println("error in do pre job", err)
							return
						}
						start = time.Now()
						err = job.Do()
						p.collector <- time.Since(start)
						if err != nil {
							atomic.AddInt64(&p.errCount, 1)
						}
						atomic.AddInt64(&p.Total, 1)
					}
				}
			}()
		}

		p.wg.Add(1)
		go func() {
			// stoper to cancell all the workers
			p.wg.Done()
			close(starter)
			time.Sleep(time.Second * time.Duration(p.Duration))
			close(p.done)
			return
		}()
	}
}

// Wait for the benchmark task done and caculate the result
func (p *perfmonitor) Wait() {
	var sum2, max, min, p70, p80, p90, p95 float64
	min = 0x7fffffffffffffff
	p.wg.Wait()
	p.Total--
	sortSlice := make([]float64, p.Total)
	i := 0
	for d := range p.buffer {
		sortSlice[i] = float64(d)
		p.histogram.Add(float64(d))
		p.Sum += float64(d)
		sum2 += float64(d * d)
		i++
	}
	sort.Slice(sortSlice, func(i, j int) bool { return sortSlice[i] < sortSlice[j] })
	p70 = sortSlice[int(float64(p.Total)*0.7)] / 1000000
	p80 = sortSlice[int(float64(p.Total)*0.8)] / 1000000
	p90 = sortSlice[int(float64(p.Total)*0.9)] / 1000000
	p95 = sortSlice[int(float64(p.Total)*0.95)] / 1000000
	min = sortSlice[0]
	max = sortSlice[p.Total-1]

	p.Mean = p.histogram.(*hist.NumericHistogram).Mean()
	p.Stdev = math.Sqrt((float64(sum2) - 2*float64(p.Mean*p.Sum) + float64(float64(p.Total)*p.Mean*p.Mean)) / float64(p.Total))

	fmt.Println("\n===============================================")
	// here show the histogram
	if p.errCount != 0 {
		fmt.Printf("Total errors: %v\t Error percentage: %.3f%%\n", p.errCount, float64(p.errCount*100)/float64(p.Total))
	}
	fmt.Printf("MAX: %.3fms MIN: %.3fms MEAN: %.3fms STDEV: %.3f CV: %.3f%% ", max/1000000, min/1000000, p.Mean/1000000, p.Stdev/1000000, p.Stdev/float64(p.Mean)*100)
	fmt.Println(p.histogram)
	fmt.Println("===============================================")
	fmt.Printf("Summary:\n70%% in:\t%.3fms\n80%% in:\t%.3fms\n90%% in:\t%.3fms\n95%% in:\t%.3fms\n", p70, p80, p90, p95)
}
<commit_msg>fix panix bug<commit_after>package perfm

import (
	"fmt"
	"math"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	hist "github.com/arthurkiller/perfm/histogram"
)

// Job give out a job for parallel call
// 1. start workers
// 		1. workers call job.Copy()
// 		2. for-loop do
// 			* job.Pre()
// 			* job.Do()
// 		3. after for-loop call job.After()
// 2. caculate the summary
type Job interface {
	// Copy will copy a job for parallel call
	Copy() (Job, error)
	// Pre will called before do
	Pre() error
	// Do contains the core job here
	Do() error
	// After contains the clean job after job done
	After()
}

//PerfMonitor define the atcion about perfmonitor
type PerfMonitor interface {
	Regist(Job) //regist the job to perfm
	Start()     //start the perf monitor
	Wait()      //wait for the benchmark done
}

type perfmonitor struct {
	Sum   float64 //Sum of the per request cost
	Stdev float64 //Standard Deviation
	Mean  float64 //Mean about distribution
	Total int64   //total request by count

	Config                            //configration for perfm
	done           chan int           //stop the perfm
	startTime      time.Time          //keep the start time
	timer          <-chan time.Time   //the frequency sampling timer
	collector      chan time.Duration //get the request cost from every done()
	errCount       int64              //error counter count error request
	localCount     int                //count for the number in the sampling times
	localTimeCount time.Duration      //count for the sampling time total costs
	buffer         chan int64         //buffer the test time for latter add to the historgam
	histogram      hist.Histogram     //used to print the histogram
	wg             sync.WaitGroup     //wait group to block the stop and sync the work thread

	//job implement benchmark job
	//error occoured in job.Do will be collected
	job Job
}

func New(options ...Options) PerfMonitor { return &perfmonitor{Config: newConfig(options...)} }

// Regist a job into perfmonitor fro benchmark
func (p *perfmonitor) Regist(job Job) {
	p.timer = time.Tick(time.Second * time.Duration(p.Frequency))
	p.collector = make(chan time.Duration, p.BufferSize)
	p.histogram = hist.NewHistogram(p.BinsNumber)
	p.done = make(chan int, 0)
	p.buffer = make(chan int64, 100000000)
	p.wg = sync.WaitGroup{}
	p.job = job

	p.Sum = 0
	p.Stdev = 0
	p.Mean = 0
	p.Total = 0
	p.errCount = 0
	p.localCount = 0
	p.localTimeCount = 0
}

// Start the benchmark with given arguments on regisit
func (p *perfmonitor) Start() {
	if p.job == nil {
		panic("error job does not registed yet")
	}
	var localwg sync.WaitGroup

	// If job implement descripetion as Stringer
	if _, ok := p.job.(fmt.Stringer); ok {
		fmt.Println(p.job)
	}
	fmt.Println("===============================================")

	p.wg.Add(1)
	go func() {
		p.startTime = time.Now()
		var cost time.Duration
		for {
			select {
			case cost = <-p.collector:
				p.localCount++
				p.localTimeCount += cost
				p.buffer <- int64(cost)
			case <-p.timer:
				if p.localCount == 0 {
					continue
				}
				if !p.NoPrint {
					fmt.Printf("%s \t  Qps: %d \t  Avg Latency: %.3fms\n", time.Now().Format("15:04:05.000"),
						p.localCount, float64(p.localTimeCount.Nanoseconds()/int64(p.localCount))/1000000)
				}
				p.localCount = 0
				p.localTimeCount = 0
			case <-p.done:
				localwg.Wait()
				close(p.collector)
				for cost := range p.collector {
					p.localCount++
					p.localTimeCount += cost
					p.buffer <- int64(cost)
				}
				if !p.NoPrint {
					fmt.Printf("%s \t  Qps: %d \t  Avg Latency: %.3fms\n", time.Now().Format("15:04:05.000"),
						p.localCount, float64(p.localTimeCount.Nanoseconds()/int64(p.localCount))/1000000)
				}
				close(p.buffer)
				p.wg.Done()
				return
			}
		}
	}()

	if p.Number > 0 {
		// in total request module
		sum := int64(p.Number)
		for i := 0; i < p.Parallel; i++ {
			localwg.Add(1)
			go func() {
				defer localwg.Done()
				var err error
				job, err := p.job.Copy()
				if err != nil {
					fmt.Println("error in do copy", err)
					return
				}
				defer job.After()
				var start time.Time
				var l int64
				for {
					if l = atomic.AddInt64(&p.Total, 1); l > sum {
						if l == sum+1 {
							close(p.done)
						}
						// check if the request reach the goal
						return
					}
					if err = job.Pre(); err != nil {
						fmt.Println("error in do pre job", err)
						return
					}
					start = time.Now()
					err = job.Do()
					p.collector <- time.Since(start)
					if err != nil {
						atomic.AddInt64(&p.errCount, 1)
					}
				}
			}()
		}
	} else {
		// in test duration module
		// start all the worker and do job till cancelled
		starter := make(chan struct{})
		for i := 0; i < p.Parallel; i++ {
			localwg.Add(1)
			go func() {
				defer localwg.Done()
				var err error
				job, err := p.job.Copy()
				if err != nil {
					fmt.Println("error in do copy", err)
					return
				}
				defer job.After()
				var start time.Time
				<-starter
				for {
					select {
					case <-p.done:
						return
					default:
						if err = job.Pre(); err != nil {
							fmt.Println("error in do pre job", err)
							return
						}
						start = time.Now()
						err = job.Do()
						p.collector <- time.Since(start)
						if err != nil {
							atomic.AddInt64(&p.errCount, 1)
						}
						atomic.AddInt64(&p.Total, 1)
					}
				}
			}()
		}

		p.wg.Add(1)
		go func() {
			// stoper to cancell all the workers
			p.wg.Done()
			close(starter)
			time.Sleep(time.Second * time.Duration(p.Duration))
			close(p.done)
			return
		}()
	}
}

// Wait for the benchmark task done and caculate the result
func (p *perfmonitor) Wait() {
	var sum2, max, min, p70, p80, p90, p95 float64
	min = 0x7fffffffffffffff
	p.wg.Wait()
	p.Total--
	sortSlice := make([]float64, 0, len(p.buffer))
	for d := range p.buffer {
		sortSlice = append(sortSlice, float64(d))
		p.histogram.Add(float64(d))
		p.Sum += float64(d)
		sum2 += float64(d * d)
	}
	sort.Slice(sortSlice, func(i, j int) bool { return sortSlice[i] < sortSlice[j] })
	p70 = sortSlice[int(float64(p.Total)*0.7)] / 1000000
	p80 = sortSlice[int(float64(p.Total)*0.8)] / 1000000
	p90 = sortSlice[int(float64(p.Total)*0.9)] / 1000000
	p95 = sortSlice[int(float64(p.Total)*0.95)] / 1000000
	min = sortSlice[0]
	max = sortSlice[p.Total-1]

	p.Mean = p.histogram.(*hist.NumericHistogram).Mean()
	p.Stdev = math.Sqrt((float64(sum2) - 2*float64(p.Mean*p.Sum) + float64(float64(p.Total)*p.Mean*p.Mean)) / float64(p.Total))

	fmt.Println("\n===============================================")
	// here show the histogram
	if p.errCount != 0 {
		fmt.Printf("Total errors: %v\t Error percentage: %.3f%%\n", p.errCount, float64(p.errCount*100)/float64(p.Total))
	}
	fmt.Printf("MAX: %.3fms MIN: %.3fms MEAN: %.3fms STDEV: %.3f CV: %.3f%% ", max/1000000, min/1000000, p.Mean/1000000, p.Stdev/1000000, p.Stdev/float64(p.Mean)*100)
	fmt.Println(p.histogram)
	fmt.Println("===============================================")
	fmt.Printf("Summary:\n70%% in:\t%.3fms\n80%% in:\t%.3fms\n90%% in:\t%.3fms\n95%% in:\t%.3fms\n", p70, p80, p90, p95)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"log"
	"time"
)

var (
	device        = flag.String("device", "lo", "")
	snapshotLen   = flag.Int("snapshot_len", 2048, "")
	BPFFilter     = flag.String("bpf_filter", "tcp and port 5432", "")
	queryFilter   = flag.String("query_filter", "", "not case-sensitive")
	queries       = make(map[string]query)
	slowQueryTime = flag.Int64("slow_query_time", 0, "in milliseconds")
)

type query struct {
	query string
	start time.Time
}

func main() {

	flag.Parse()

	handle, err := pcap.OpenLive(*device, int32(*snapshotLen), true, time.Second)

	defer handle.Close()

	if err != nil {

		log.Fatal(err)
	}

	handle.SetBPFFilter(*BPFFilter)

	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())

	var (
		ipLayer  *layers.IPv4
		tcpLayer *layers.TCP
		ok       bool
	)

	for packet := range packetSource.Packets() {

		if applicationLayer := packet.ApplicationLayer(); applicationLayer != nil {

			if ipLayer, ok = packet.Layer(layers.LayerTypeIPv4).(*layers.IPv4); !ok {

				continue
			}

			if tcpLayer, ok = packet.Layer(layers.LayerTypeTCP).(*layers.TCP); !ok {

				continue
			}

			playload := applicationLayer.Payload()

			if len(playload) < 5 {

				continue
			}

			length := _len(playload[1:5])

			if length > len(playload)-1 {

				continue
			}

			switch playload[0] {

			case 'Q', 'P':

				from := fmt.Sprintf("%s%d:%s%d\n", ipLayer.SrcIP, tcpLayer.SrcPort, ipLayer.DstIP, tcpLayer.DstPort)

				if *queryFilter == "" || bytes.Contains(bytes.ToLower(playload[5:length]), bytes.ToLower([]byte(*queryFilter))) {

					queries[from] = query{
						query: string(playload[5:length]),
						start: packet.Metadata().Timestamp,
					}
				}

			default:

				from := fmt.Sprintf("%s%d:%s%d\n", ipLayer.DstIP, tcpLayer.DstPort, ipLayer.SrcIP, tcpLayer.SrcPort)

				if query, found := queries[from]; found {

					queryTime := packet.Metadata().Timestamp.Sub(query.start)

					if *slowQueryTime == 0 || queryTime.Nanoseconds()/1000 > *slowQueryTime {

						fmt.Printf("-[ QUERY %f s]-:\n%s\n\n\n", queryTime.Seconds(), query.query)
					}

					delete(queries, from)
				}
			}
		}
	}
}

func _len(b []byte) int {

	return int(b[0])<<24 | int(b[1])<<16 | int(b[2])<<8 | int(b[3])
}
<commit_msg>slow time<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"log"
	"time"
)

var (
	device        = flag.String("device", "lo", "")
	snapshotLen   = flag.Int("snapshot_len", 2048, "")
	BPFFilter     = flag.String("bpf_filter", "tcp and port 5432", "")
	queryFilter   = flag.String("query_filter", "", "not case-sensitive")
	queries       = make(map[string]query)
	slowQueryTime = flag.Int64("slow_query_time", 0, "in milliseconds")
)

type query struct {
	query string
	start time.Time
}

func main() {

	flag.Parse()

	handle, err := pcap.OpenLive(*device, int32(*snapshotLen), true, time.Second)

	defer handle.Close()

	if err != nil {

		log.Fatal(err)
	}

	handle.SetBPFFilter(*BPFFilter)

	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())

	var (
		ipLayer  *layers.IPv4
		tcpLayer *layers.TCP
		ok       bool
	)

	for packet := range packetSource.Packets() {

		if applicationLayer := packet.ApplicationLayer(); applicationLayer != nil {

			if ipLayer, ok = packet.Layer(layers.LayerTypeIPv4).(*layers.IPv4); !ok {

				continue
			}

			if tcpLayer, ok = packet.Layer(layers.LayerTypeTCP).(*layers.TCP); !ok {

				continue
			}

			playload := applicationLayer.Payload()

			if len(playload) < 5 {

				continue
			}

			length := _len(playload[1:5])

			if length > len(playload)-1 {

				continue
			}

			switch playload[0] {

			case 'Q', 'P':

				from := fmt.Sprintf("%s%d:%s%d\n", ipLayer.SrcIP, tcpLayer.SrcPort, ipLayer.DstIP, tcpLayer.DstPort)

				if *queryFilter == "" || bytes.Contains(bytes.ToLower(playload[5:length]), bytes.ToLower([]byte(*queryFilter))) {

					queries[from] = query{
						query: string(playload[5:length]),
						start: packet.Metadata().Timestamp,
					}
				}

			default:

				from := fmt.Sprintf("%s%d:%s%d\n", ipLayer.DstIP, tcpLayer.DstPort, ipLayer.SrcIP, tcpLayer.SrcPort)

				if query, found := queries[from]; found {

					queryTime := packet.Metadata().Timestamp.Sub(query.start)

					if *slowQueryTime == 0 || queryTime.Nanoseconds()/1000000 > *slowQueryTime {

						fmt.Printf("-[ QUERY %f s]-:\n%s\n\n\n", queryTime.Seconds(), query.query)
					}

					delete(queries, from)
				}
			}
		}
	}
}

func _len(b []byte) int {

	return int(b[0])<<24 | int(b[1])<<16 | int(b[2])<<8 | int(b[3])
}
<|endoftext|>


<commit_before>/*
 * Copyright 2013 Nan Deng
 *
 * 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 msgcenter

import (
	"crypto/rsa"
	"fmt"
	"github.com/uniqush/uniqush-conn/config"
	"github.com/uniqush/uniqush-conn/proto/server"
	"github.com/uniqush/uniqush-conn/rpc"
	"net"
	"sync"
)

type MessageCenter struct {
	ln      net.Listener
	privkey *rsa.PrivateKey
	config  *config.Config

	srvCentersLock   sync.Mutex
	serviceCenterMap map[string]*serviceCenter
	fwdChan          chan *rpc.ForwardRequest
	peers            *rpc.MultiPeer
}

func (self *MessageCenter) processForwardRequest() {
	for req := range self.fwdChan {
		if req == nil {
			return
		}
		if len(req.ReceiverService) == 0 {
			continue
		}
		center := self.getServiceCenter(req.ReceiverService)
		if center != nil {
			go center.Forward(req)
		}
	}
}

func NewMessageCenter(ln net.Listener, privkey *rsa.PrivateKey, conf *config.Config) *MessageCenter {
	ret := new(MessageCenter)
	ret.ln = ln
	ret.privkey = privkey
	ret.config = conf

	ret.peers = rpc.NewMultiPeer()
	ret.fwdChan = make(chan *rpc.ForwardRequest)
	ret.serviceCenterMap = make(map[string]*serviceCenter, 10)

	go ret.processForwardRequest()
	return ret
}

func (self *MessageCenter) ServiceNames() []string {
	self.srvCentersLock.Lock()
	defer self.srvCentersLock.Unlock()

	ret := make([]string, 0, len(self.serviceCenterMap))

	for srv, _ := range self.serviceCenterMap {
		ret = append(ret, srv)
	}
	return ret
}

func (self *MessageCenter) getServiceCenter(srv string) *serviceCenter {
	self.srvCentersLock.Lock()
	defer self.srvCentersLock.Unlock()

	center, ok := self.serviceCenterMap[srv]
	if !ok {
		conf := self.config.ReadConfig(srv)
		if conf == nil {
			return nil
		}
		center = newServiceCenter(conf, self.fwdChan, self.peers)
		if center != nil {
			self.serviceCenterMap[srv] = center
		}
	}
	return center
}

func (self *MessageCenter) serveConn(c net.Conn) {
	if tcpConn, ok := c.(*net.TCPConn); ok {
		tcpConn.SetKeepAlive(true)
	}
	conn, err := server.AuthConn(c, self.privkey, self.config, self.config.HandshakeTimeout)
	if err != nil {
		if err != server.ErrAuthFail {
			self.config.OnError(c.RemoteAddr(), err)
		}
		c.Close()
		return
	}
	srv := conn.Service()

	center := self.getServiceCenter(srv)
	if center == nil {
		self.config.OnError(c.RemoteAddr(), fmt.Errorf("unknown service: %v", srv))
		c.Close()
		return
	}
	center.NewConn(conn)
	return
}

func (self *MessageCenter) Start() {
	for {
		conn, err := self.ln.Accept()
		if err != nil {
			if ne, ok := err.(net.Error); ok {
				// It's a temporary error.
				if ne.Temporary() {
					continue
				}
			}
			self.config.OnError(self.ln.Addr(), err)
			return
		}
		go self.serveConn(conn)
	}
}

// NOTE: you cannot restart it!
func (self *MessageCenter) Stop() {
	self.srvCentersLock.Lock()

	for _, center := range self.serviceCenterMap {
		center.Stop()
	}
	self.ln.Close()
	close(self.fwdChan)
}

func (self *MessageCenter) AllServices() []string {
	self.srvCentersLock.Lock()
	defer self.srvCentersLock.Lock()
	ret := make([]string, 0, len(self.serviceCenterMap))
	for srv, _ := range self.serviceCenterMap {
		ret = append(ret, srv)
	}
	return ret
}

func (self *MessageCenter) AllUsernames(srv string) []string {
	center := self.getServiceCenter(srv)
	if center == nil {
		return nil
	}
	return center.AllUsernames()
}

func (self *MessageCenter) NrConns(srv string) int {
	center := self.getServiceCenter(srv)
	if center == nil {
		return 0
	}
	return center.NrConns()
}

func (self *MessageCenter) NrUsers(srv string) int {
	center := self.getServiceCenter(srv)
	if center == nil {
		return 0
	}
	return center.NrUsers()
}

func (self *MessageCenter) do(srv string, f func(center *serviceCenter) *rpc.Result) *rpc.Result {
	center := self.getServiceCenter(srv)
	if center == nil {
		ret := new(rpc.Result)
		//ret.SetError(fmt.Errorf("unknown service: %v", srv))
		return ret
	}
	return f(center)
}

func (self *MessageCenter) Send(req *rpc.SendRequest) *rpc.Result {
	return self.do(req.ReceiverService, func(center *serviceCenter) *rpc.Result {
		return center.Send(req)
	})
}

func (self *MessageCenter) Forward(req *rpc.ForwardRequest) *rpc.Result {
	return self.do(req.ReceiverService, func(center *serviceCenter) *rpc.Result {
		return center.Forward(req)
	})
}

func (self *MessageCenter) Redirect(req *rpc.RedirectRequest) *rpc.Result {
	return self.do(req.ReceiverService, func(center *serviceCenter) *rpc.Result {
		return center.Redirect(req)
	})
}

func (self *MessageCenter) CheckUserStatus(req *rpc.UserStatusQuery) *rpc.Result {
	return self.do(req.Service, func(center *serviceCenter) *rpc.Result {
		return center.CheckUserStatus(req)
	})
}

func (self *MessageCenter) AddPeer(peer rpc.UniqushConnPeer) {
	self.peers.AddPeer(peer)
}
<commit_msg>use nagle<commit_after>/*
 * Copyright 2013 Nan Deng
 *
 * 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 msgcenter

import (
	"crypto/rsa"
	"fmt"
	"github.com/uniqush/uniqush-conn/config"
	"github.com/uniqush/uniqush-conn/proto/server"
	"github.com/uniqush/uniqush-conn/rpc"
	"net"
	"sync"
)

type MessageCenter struct {
	ln      net.Listener
	privkey *rsa.PrivateKey
	config  *config.Config

	srvCentersLock   sync.Mutex
	serviceCenterMap map[string]*serviceCenter
	fwdChan          chan *rpc.ForwardRequest
	peers            *rpc.MultiPeer
}

func (self *MessageCenter) processForwardRequest() {
	for req := range self.fwdChan {
		if req == nil {
			return
		}
		if len(req.ReceiverService) == 0 {
			continue
		}
		center := self.getServiceCenter(req.ReceiverService)
		if center != nil {
			go center.Forward(req)
		}
	}
}

func NewMessageCenter(ln net.Listener, privkey *rsa.PrivateKey, conf *config.Config) *MessageCenter {
	ret := new(MessageCenter)
	ret.ln = ln
	ret.privkey = privkey
	ret.config = conf

	ret.peers = rpc.NewMultiPeer()
	ret.fwdChan = make(chan *rpc.ForwardRequest)
	ret.serviceCenterMap = make(map[string]*serviceCenter, 10)

	go ret.processForwardRequest()
	return ret
}

func (self *MessageCenter) ServiceNames() []string {
	self.srvCentersLock.Lock()
	defer self.srvCentersLock.Unlock()

	ret := make([]string, 0, len(self.serviceCenterMap))

	for srv, _ := range self.serviceCenterMap {
		ret = append(ret, srv)
	}
	return ret
}

func (self *MessageCenter) getServiceCenter(srv string) *serviceCenter {
	self.srvCentersLock.Lock()
	defer self.srvCentersLock.Unlock()

	center, ok := self.serviceCenterMap[srv]
	if !ok {
		conf := self.config.ReadConfig(srv)
		if conf == nil {
			return nil
		}
		center = newServiceCenter(conf, self.fwdChan, self.peers)
		if center != nil {
			self.serviceCenterMap[srv] = center
		}
	}
	return center
}

func (self *MessageCenter) serveConn(c net.Conn) {
	if tcpConn, ok := c.(*net.TCPConn); ok {
		// Rather than keeping an application leve heart beat,
		// we rely on TCP-level keep-alive.
		// XXX Is this a good idea?
		tcpConn.SetKeepAlive(true)

		// Use Nagle.
		tcpConn.SetNoDelay(false)
	}
	conn, err := server.AuthConn(c, self.privkey, self.config, self.config.HandshakeTimeout)
	if err != nil {
		if err != server.ErrAuthFail {
			self.config.OnError(c.RemoteAddr(), err)
		}
		c.Close()
		return
	}
	srv := conn.Service()

	center := self.getServiceCenter(srv)
	if center == nil {
		self.config.OnError(c.RemoteAddr(), fmt.Errorf("unknown service: %v", srv))
		c.Close()
		return
	}
	center.NewConn(conn)
	return
}

func (self *MessageCenter) Start() {
	for {
		conn, err := self.ln.Accept()
		if err != nil {
			if ne, ok := err.(net.Error); ok {
				// It's a temporary error.
				if ne.Temporary() {
					continue
				}
			}
			self.config.OnError(self.ln.Addr(), err)
			return
		}
		go self.serveConn(conn)
	}
}

// NOTE: you cannot restart it!
func (self *MessageCenter) Stop() {
	self.srvCentersLock.Lock()

	for _, center := range self.serviceCenterMap {
		center.Stop()
	}
	self.ln.Close()
	close(self.fwdChan)
}

func (self *MessageCenter) AllServices() []string {
	self.srvCentersLock.Lock()
	defer self.srvCentersLock.Lock()
	ret := make([]string, 0, len(self.serviceCenterMap))
	for srv, _ := range self.serviceCenterMap {
		ret = append(ret, srv)
	}
	return ret
}

func (self *MessageCenter) AllUsernames(srv string) []string {
	center := self.getServiceCenter(srv)
	if center == nil {
		return nil
	}
	return center.AllUsernames()
}

func (self *MessageCenter) NrConns(srv string) int {
	center := self.getServiceCenter(srv)
	if center == nil {
		return 0
	}
	return center.NrConns()
}

func (self *MessageCenter) NrUsers(srv string) int {
	center := self.getServiceCenter(srv)
	if center == nil {
		return 0
	}
	return center.NrUsers()
}

func (self *MessageCenter) do(srv string, f func(center *serviceCenter) *rpc.Result) *rpc.Result {
	center := self.getServiceCenter(srv)
	if center == nil {
		ret := new(rpc.Result)
		//ret.SetError(fmt.Errorf("unknown service: %v", srv))
		return ret
	}
	return f(center)
}

func (self *MessageCenter) Send(req *rpc.SendRequest) *rpc.Result {
	return self.do(req.ReceiverService, func(center *serviceCenter) *rpc.Result {
		return center.Send(req)
	})
}

func (self *MessageCenter) Forward(req *rpc.ForwardRequest) *rpc.Result {
	return self.do(req.ReceiverService, func(center *serviceCenter) *rpc.Result {
		return center.Forward(req)
	})
}

func (self *MessageCenter) Redirect(req *rpc.RedirectRequest) *rpc.Result {
	return self.do(req.ReceiverService, func(center *serviceCenter) *rpc.Result {
		return center.Redirect(req)
	})
}

func (self *MessageCenter) CheckUserStatus(req *rpc.UserStatusQuery) *rpc.Result {
	return self.do(req.Service, func(center *serviceCenter) *rpc.Result {
		return center.CheckUserStatus(req)
	})
}

func (self *MessageCenter) AddPeer(peer rpc.UniqushConnPeer) {
	self.peers.AddPeer(peer)
}
<|endoftext|>


<commit_before>/*
 * Copyright 2013 Nan Deng
 *
 * 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 msgcenter

import (
	"errors"
	"fmt"
	"github.com/uniqush/uniqush-conn/evthandler"
	"github.com/uniqush/uniqush-conn/msgcache"
	"github.com/uniqush/uniqush-conn/proto"
	"github.com/uniqush/uniqush-conn/proto/server"
	"github.com/uniqush/uniqush-conn/push"
	"strings"
	"time"
)

type eventConnIn struct {
	errChan chan error
	conn    server.Conn
}

type eventConnLeave struct {
	conn server.Conn
	err  error
}

type EventConnError struct {
	Err error
	C   server.Conn
}

func (self *EventConnError) Service() string {
	return self.C.Service()
}

func (self *EventConnError) Username() string {
	return self.C.Username()
}

func (self *EventConnError) Error() string {
	return fmt.Sprintf("[Service=%v][User=%v] %v", self.C.Service(), self.C.Username(), self.Err)
}

type ServiceConfig struct {
	MaxNrConns        int
	MaxNrUsers        int
	MaxNrConnsPerUser int

	MsgCache msgcache.Cache

	LoginHandler          evthandler.LoginHandler
	LogoutHandler         evthandler.LogoutHandler
	MessageHandler        evthandler.MessageHandler
	ForwardRequestHandler evthandler.ForwardRequestHandler
	ErrorHandler          evthandler.ErrorHandler

	// Push related web hooks
	SubscribeHandler   evthandler.SubscribeHandler
	UnsubscribeHandler evthandler.UnsubscribeHandler
	PushHandler        evthandler.PushHandler

	PushService push.Push
}

type writeMessageResponse struct {
	err []error
	n   int
}

type writeMessageRequest struct {
	user      string
	msg       *proto.Message
	posterKey string
	ttl       time.Duration
	extra     map[string]string
	resChan   chan<- *writeMessageResponse
}

type serviceCenter struct {
	serviceName string
	config      *ServiceConfig
	fwdChan     chan<- *server.ForwardRequest

	writeReqChan chan *writeMessageRequest
	connIn       chan *eventConnIn
	connLeave    chan *eventConnLeave
}

var ErrTooManyConns = errors.New("too many connections")
var ErrInvalidConnType = errors.New("invalid connection type")

func (self *serviceCenter) ReceiveForward(fwdreq *server.ForwardRequest) {
	shouldFwd := false
	if self.config != nil {
		if self.config.ForwardRequestHandler != nil {
			shouldFwd = self.config.ForwardRequestHandler.ShouldForward(fwdreq)
			maxttl := self.config.ForwardRequestHandler.MaxTTL()
			if fwdreq.TTL < 1 * time.Second || fwdreq.TTL > maxttl {
				fwdreq.TTL = maxttl
			}
		}
	}
	if !shouldFwd {
		return
	}
	receiver := fwdreq.Receiver
	extra := getPushInfo(fwdreq.Message, nil, true)
	self.SendMail(receiver, fwdreq.Message, extra, fwdreq.TTL)
}

func getPushInfo(msg *proto.Message, extra map[string]string, fwd bool) map[string]string {
	if extra == nil {
		extra = make(map[string]string, len(msg.Header)+3)
	}
	if fwd {
		extra["sender"] = msg.Sender
		extra["sender-service"] = msg.SenderService
		for k, v := range msg.Header {
			if len(k) > 7 {
				extra[k] = v
			}
		}
	}
	if msg.Header != nil {
		if title, ok := msg.Header["title"]; ok {
			extra["notif.msg"] = title
		}
	}
	return extra
}

func (self *serviceCenter) shouldPush(service, username string, msg *proto.Message, extra map[string]string, fwd bool) bool {
	if self.config != nil {
		if self.config.PushHandler != nil {
			info := getPushInfo(msg, extra, fwd)
			return self.config.PushHandler.ShouldPush(service, username, info)
		} else {
			fmt.Printf("Should not push it, no push handler\n")
		}
	}
	return false
}

func (self *serviceCenter) nrDeliveryPoints(service, username string) int {
	n := 0
	if self.config != nil {
		if self.config.PushService != nil {
			n = self.config.PushService.NrDeliveryPoints(service, username)
		}
	}
	return n
}

func (self *serviceCenter) pushNotif(service, username string, msg *proto.Message, extra map[string]string, msgIds[]string, fwd bool) {
	if self.config != nil {
		if self.config.PushService != nil {
			info := getPushInfo(msg, extra, fwd)
			err := self.config.PushService.Push(service, username, info, msgIds)
			if err != nil {
				self.reportError(service, username, "", err)
			}
		}
	}
}

func (self *serviceCenter) reportError(service, username, connId string, err error) {
	if self.config != nil {
		if self.config.ErrorHandler != nil {
			self.config.ErrorHandler.OnError(service, username, connId, err)
		}
	}
}

func (self *serviceCenter) reportLogin(service, username, connId string) {
	if self.config != nil {
		if self.config.LoginHandler != nil {
			self.config.LoginHandler.OnLogin(service, username, connId)
		}
	}
}

func (self *serviceCenter) reportMessage(connId string, msg *proto.Message) {
	if self.config != nil {
		if self.config.MessageHandler != nil {
			self.config.MessageHandler.OnMessage(connId, msg)
		}
	}
}

func (self *serviceCenter) reportLogout(service, username, connId string, err error) {
	if self.config != nil {
		if self.config.LogoutHandler != nil {
			self.config.LogoutHandler.OnLogout(service, username, connId, err)
		}
	}
}

func (self *serviceCenter) setPoster(service, username, key string, msg *proto.Message, ttl time.Duration) (id string, err error) {
	if self.config != nil {
		if self.config.MsgCache != nil {
			id, err = self.config.MsgCache.SetPoster(service, username, key, msg, ttl)
		}
	}
	return
}

func (self *serviceCenter) setMail(service, username string, msg *proto.Message, ttl time.Duration) (id string, err error) {
	if self.config != nil {
		if self.config.MsgCache != nil {
			id, err = self.config.MsgCache.SetMail(service, username, msg, ttl)
		}
	}
	return
}

func (self *serviceCenter) process(maxNrConns, maxNrConnsPerUser, maxNrUsers int) {
	connMap := newTreeBasedConnMap()
	nrConns := 0
	for {
		select {
		case connInEvt := <-self.connIn:
			if maxNrConns > 0 && nrConns >= maxNrConns {
				if connInEvt.errChan != nil {
					connInEvt.errChan <- ErrTooManyConns
				}
				continue
			}
			err := connMap.AddConn(connInEvt.conn, maxNrConnsPerUser, maxNrUsers)
			if err != nil {
				if connInEvt.errChan != nil {
					connInEvt.errChan <- err
				}
				continue
			}
			nrConns++
			if connInEvt.errChan != nil {
				connInEvt.errChan <- nil
			}
		case leaveEvt := <-self.connLeave:
			connMap.DelConn(leaveEvt.conn)
			leaveEvt.conn.Close()
			nrConns--
			conn := leaveEvt.conn
			self.reportLogout(conn.Service(), conn.Username(), conn.UniqId(), leaveEvt.err)
		case wreq := <-self.writeReqChan:
			wres := new(writeMessageResponse)
			wres.n = 0
			conns := connMap.GetConn(wreq.user)
			if len(wreq.posterKey) != 0 && len(conns) > 0 {
				self.setPoster(self.serviceName, wreq.user, wreq.posterKey, wreq.msg, wreq.ttl)
			}
			for _, conn := range conns {
				if conn == nil {
					continue
				}
				var err error
				sconn, ok := conn.(server.Conn)
				if !ok {
					wres.err = append(wres.err, ErrInvalidConnType)
					break
				}
				if len(wreq.posterKey) == 0 {
					_, err = sconn.SendMail(wreq.msg, wreq.extra, wreq.ttl)
				} else {
					_, err = sconn.SendPoster(wreq.msg, wreq.extra, wreq.posterKey, wreq.ttl, false)
				}
				if err != nil {
					wres.err = append(wres.err, err)
					self.reportError(sconn.Service(), sconn.Username(), sconn.UniqId(), err)
					continue
				}
				if sconn.Visible() {
					wres.n++
				}
			}

			if wres.n == 0 {
				fmt.Printf("Should push this message\n")
				msg := wreq.msg
				extra := wreq.extra
				username := wreq.user
				service := self.serviceName
				fwd := false
				if len(msg.Sender) > 0 && len(msg.SenderService) > 0 {
					if msg.Sender != username || msg.SenderService != service {
						fwd = true
					}
				}
				go func() {
					if !self.shouldPush(service, username, msg, extra, fwd) {
						return
					}
					fmt.Printf("OK! Let's push it\n")
					n := self.nrDeliveryPoints(service, username)
					if n <= 0 {
						return
					}
					var msgIds []string
					if len(wreq.posterKey) == 0 {
						msgIds = make([]string, n)
						var e error
						for i := 0; i < n; i++ {
							msgIds[i], e = self.setMail(service, username, msg, wreq.ttl)
							if e != nil {
								// FIXME: Dark side of the force
								return
							}
						}
					} else {
						id, e := self.setPoster(service, wreq.user, wreq.posterKey, wreq.msg, wreq.ttl)
						if e != nil {
							// FIXME: Dark side of the force
							return
						}
						msgIds = []string{id}
					}
						self.pushNotif(service, username, msg, extra, msgIds, fwd)
				}()
			}
			if wreq.resChan != nil {
				wreq.resChan <- wres
			}
		}
	}
}

func (self *serviceCenter) SendMail(username string, msg *proto.Message, extra map[string]string, ttl time.Duration) (n int, err []error) {
	req := new(writeMessageRequest)
	ch := make(chan *writeMessageResponse)
	req.msg = msg
	req.posterKey = ""
	req.user = username
	req.ttl = ttl
	req.resChan = ch
	req.extra = extra
	self.writeReqChan <- req
	res := <-ch
	n = res.n
	err = res.err
	return
}

func (self *serviceCenter) SendPoster(username string, msg *proto.Message, extra map[string]string, key string, ttl time.Duration) (n int, err []error) {
	req := new(writeMessageRequest)
	ch := make(chan *writeMessageResponse)
	req.msg = msg
	req.posterKey = key
	req.ttl = ttl
	req.extra = extra
	req.user = username
	req.resChan = ch
	self.writeReqChan <- req
	res := <-ch
	n = res.n
	err = res.err
	return
}

func (self *serviceCenter) serveConn(conn server.Conn) {
	conn.SetForwardRequestChannel(self.fwdChan)
	var err error
	defer func() {
		self.connLeave <- &eventConnLeave{conn: conn, err: err}
	}()
	for {
		var msg *proto.Message
		msg, err = conn.ReadMessage()
		if err != nil {
			return
		}
		self.reportMessage(conn.UniqId(), msg)
	}
}

func (self *serviceCenter) NewConn(conn server.Conn) error {
	usr := conn.Username()
	if len(usr) == 0 || strings.Contains(usr, ":") || strings.Contains(usr, "\n") {
		return fmt.Errorf("[Username=%v] Invalid Username")
	}
	evt := new(eventConnIn)
	ch := make(chan error)

	conn.SetMessageCache(self.config.MsgCache)
	evt.conn = conn
	evt.errChan = ch
	self.connIn <- evt
	err := <-ch
	if err == nil {
		go self.serveConn(conn)
		self.reportLogin(conn.Service(), usr, conn.UniqId())
	}
	return err
}

func newServiceCenter(serviceName string, conf *ServiceConfig, fwdChan chan<- *server.ForwardRequest) *serviceCenter {
	ret := new(serviceCenter)
	ret.config = conf
	if ret.config == nil {
		ret.config = new(ServiceConfig)
	}
	ret.serviceName = serviceName
	ret.fwdChan = fwdChan

	ret.connIn = make(chan *eventConnIn)
	ret.connLeave = make(chan *eventConnLeave)
	ret.writeReqChan = make(chan *writeMessageRequest)
	go ret.process(conf.MaxNrConns, conf.MaxNrConnsPerUser, conf.MaxNrUsers)
	return ret
}
<commit_msg>Passed basic test with uniqush-push<commit_after>/*
 * Copyright 2013 Nan Deng
 *
 * 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 msgcenter

import (
	"errors"
	"fmt"
	"github.com/uniqush/uniqush-conn/evthandler"
	"github.com/uniqush/uniqush-conn/msgcache"
	"github.com/uniqush/uniqush-conn/proto"
	"github.com/uniqush/uniqush-conn/proto/server"
	"github.com/uniqush/uniqush-conn/push"
	"strings"
	"time"
)

type eventConnIn struct {
	errChan chan error
	conn    server.Conn
}

type eventConnLeave struct {
	conn server.Conn
	err  error
}

type EventConnError struct {
	Err error
	C   server.Conn
}

func (self *EventConnError) Service() string {
	return self.C.Service()
}

func (self *EventConnError) Username() string {
	return self.C.Username()
}

func (self *EventConnError) Error() string {
	return fmt.Sprintf("[Service=%v][User=%v] %v", self.C.Service(), self.C.Username(), self.Err)
}

type ServiceConfig struct {
	MaxNrConns        int
	MaxNrUsers        int
	MaxNrConnsPerUser int

	MsgCache msgcache.Cache

	LoginHandler          evthandler.LoginHandler
	LogoutHandler         evthandler.LogoutHandler
	MessageHandler        evthandler.MessageHandler
	ForwardRequestHandler evthandler.ForwardRequestHandler
	ErrorHandler          evthandler.ErrorHandler

	// Push related web hooks
	SubscribeHandler   evthandler.SubscribeHandler
	UnsubscribeHandler evthandler.UnsubscribeHandler
	PushHandler        evthandler.PushHandler

	PushService push.Push
}

type writeMessageResponse struct {
	err []error
	n   int
}

type writeMessageRequest struct {
	user      string
	msg       *proto.Message
	posterKey string
	ttl       time.Duration
	extra     map[string]string
	resChan   chan<- *writeMessageResponse
}

type serviceCenter struct {
	serviceName string
	config      *ServiceConfig
	fwdChan     chan<- *server.ForwardRequest

	writeReqChan chan *writeMessageRequest
	connIn       chan *eventConnIn
	connLeave    chan *eventConnLeave
}

var ErrTooManyConns = errors.New("too many connections")
var ErrInvalidConnType = errors.New("invalid connection type")

func (self *serviceCenter) ReceiveForward(fwdreq *server.ForwardRequest) {
	shouldFwd := false
	if self.config != nil {
		if self.config.ForwardRequestHandler != nil {
			shouldFwd = self.config.ForwardRequestHandler.ShouldForward(fwdreq)
			maxttl := self.config.ForwardRequestHandler.MaxTTL()
			if fwdreq.TTL < 1*time.Second || fwdreq.TTL > maxttl {
				fwdreq.TTL = maxttl
			}
		}
	}
	if !shouldFwd {
		return
	}
	receiver := fwdreq.Receiver
	extra := getPushInfo(fwdreq.Message, nil, true)
	self.SendMail(receiver, fwdreq.Message, extra, fwdreq.TTL)
}

func getPushInfo(msg *proto.Message, extra map[string]string, fwd bool) map[string]string {
	if extra == nil {
		extra = make(map[string]string, len(msg.Header)+3)
	}
	if fwd {
		extra["sender"] = msg.Sender
		extra["sender-service"] = msg.SenderService
		for k, v := range msg.Header {
			if len(k) > 7 {
				extra[k] = v
			}
		}
	}
	if msg.Header != nil {
		if title, ok := msg.Header["title"]; ok {
			extra["notif.msg"] = title
		}
	}
	return extra
}

func (self *serviceCenter) shouldPush(service, username string, msg *proto.Message, extra map[string]string, fwd bool) bool {
	if self.config != nil {
		if self.config.PushHandler != nil {
			info := getPushInfo(msg, extra, fwd)
			return self.config.PushHandler.ShouldPush(service, username, info)
		}
	}
	return false
}

func (self *serviceCenter) nrDeliveryPoints(service, username string) int {
	n := 0
	if self.config != nil {
		if self.config.PushService != nil {
			n = self.config.PushService.NrDeliveryPoints(service, username)
		}
	}
	return n
}

func (self *serviceCenter) pushNotif(service, username string, msg *proto.Message, extra map[string]string, msgIds []string, fwd bool) {
	if self.config != nil {
		if self.config.PushService != nil {
			info := getPushInfo(msg, extra, fwd)
			err := self.config.PushService.Push(service, username, info, msgIds)
			if err != nil {
				self.reportError(service, username, "", err)
			}
		}
	}
}

func (self *serviceCenter) reportError(service, username, connId string, err error) {
	if self.config != nil {
		if self.config.ErrorHandler != nil {
			self.config.ErrorHandler.OnError(service, username, connId, err)
		}
	}
}

func (self *serviceCenter) reportLogin(service, username, connId string) {
	if self.config != nil {
		if self.config.LoginHandler != nil {
			self.config.LoginHandler.OnLogin(service, username, connId)
		}
	}
}

func (self *serviceCenter) reportMessage(connId string, msg *proto.Message) {
	if self.config != nil {
		if self.config.MessageHandler != nil {
			self.config.MessageHandler.OnMessage(connId, msg)
		}
	}
}

func (self *serviceCenter) reportLogout(service, username, connId string, err error) {
	if self.config != nil {
		if self.config.LogoutHandler != nil {
			self.config.LogoutHandler.OnLogout(service, username, connId, err)
		}
	}
}

func (self *serviceCenter) setPoster(service, username, key string, msg *proto.Message, ttl time.Duration) (id string, err error) {
	if self.config != nil {
		if self.config.MsgCache != nil {
			id, err = self.config.MsgCache.SetPoster(service, username, key, msg, ttl)
		}
	}
	return
}

func (self *serviceCenter) setMail(service, username string, msg *proto.Message, ttl time.Duration) (id string, err error) {
	if self.config != nil {
		if self.config.MsgCache != nil {
			id, err = self.config.MsgCache.SetMail(service, username, msg, ttl)
		}
	}
	return
}

func (self *serviceCenter) process(maxNrConns, maxNrConnsPerUser, maxNrUsers int) {
	connMap := newTreeBasedConnMap()
	nrConns := 0
	for {
		select {
		case connInEvt := <-self.connIn:
			if maxNrConns > 0 && nrConns >= maxNrConns {
				if connInEvt.errChan != nil {
					connInEvt.errChan <- ErrTooManyConns
				}
				continue
			}
			err := connMap.AddConn(connInEvt.conn, maxNrConnsPerUser, maxNrUsers)
			if err != nil {
				if connInEvt.errChan != nil {
					connInEvt.errChan <- err
				}
				continue
			}
			nrConns++
			if connInEvt.errChan != nil {
				connInEvt.errChan <- nil
			}
		case leaveEvt := <-self.connLeave:
			connMap.DelConn(leaveEvt.conn)
			leaveEvt.conn.Close()
			nrConns--
			conn := leaveEvt.conn
			self.reportLogout(conn.Service(), conn.Username(), conn.UniqId(), leaveEvt.err)
		case wreq := <-self.writeReqChan:
			wres := new(writeMessageResponse)
			wres.n = 0
			conns := connMap.GetConn(wreq.user)
			if len(wreq.posterKey) != 0 && len(conns) > 0 {
				self.setPoster(self.serviceName, wreq.user, wreq.posterKey, wreq.msg, wreq.ttl)
			}
			for _, conn := range conns {
				if conn == nil {
					continue
				}
				var err error
				sconn, ok := conn.(server.Conn)
				if !ok {
					wres.err = append(wres.err, ErrInvalidConnType)
					break
				}
				if len(wreq.posterKey) == 0 {
					_, err = sconn.SendMail(wreq.msg, wreq.extra, wreq.ttl)
				} else {
					_, err = sconn.SendPoster(wreq.msg, wreq.extra, wreq.posterKey, wreq.ttl, false)
				}
				if err != nil {
					wres.err = append(wres.err, err)
					self.reportError(sconn.Service(), sconn.Username(), sconn.UniqId(), err)
					continue
				}
				if sconn.Visible() {
					wres.n++
				}
			}

			if wres.n == 0 {
				msg := wreq.msg
				extra := wreq.extra
				username := wreq.user
				service := self.serviceName
				fwd := false
				if len(msg.Sender) > 0 && len(msg.SenderService) > 0 {
					if msg.Sender != username || msg.SenderService != service {
						fwd = true
					}
				}
				go func() {
					if !self.shouldPush(service, username, msg, extra, fwd) {
						return
					}
					n := self.nrDeliveryPoints(service, username)
					if n <= 0 {
						return
					}
					var msgIds []string
					if len(wreq.posterKey) == 0 {
						msgIds = make([]string, n)
						var e error
						for i := 0; i < n; i++ {
							msgIds[i], e = self.setMail(service, username, msg, wreq.ttl)
							if e != nil {
								// FIXME: Dark side of the force
								return
							}
						}
					} else {
						id, e := self.setPoster(service, wreq.user, wreq.posterKey, wreq.msg, wreq.ttl)
						if e != nil {
							// FIXME: Dark side of the force
							return
						}
						msgIds = []string{id}
					}
					self.pushNotif(service, username, msg, extra, msgIds, fwd)
				}()
			}
			if wreq.resChan != nil {
				wreq.resChan <- wres
			}
		}
	}
}

func (self *serviceCenter) SendMail(username string, msg *proto.Message, extra map[string]string, ttl time.Duration) (n int, err []error) {
	req := new(writeMessageRequest)
	ch := make(chan *writeMessageResponse)
	req.msg = msg
	req.posterKey = ""
	req.user = username
	req.ttl = ttl
	req.resChan = ch
	req.extra = extra
	self.writeReqChan <- req
	res := <-ch
	n = res.n
	err = res.err
	return
}

func (self *serviceCenter) SendPoster(username string, msg *proto.Message, extra map[string]string, key string, ttl time.Duration) (n int, err []error) {
	req := new(writeMessageRequest)
	ch := make(chan *writeMessageResponse)
	req.msg = msg
	req.posterKey = key
	req.ttl = ttl
	req.extra = extra
	req.user = username
	req.resChan = ch
	self.writeReqChan <- req
	res := <-ch
	n = res.n
	err = res.err
	return
}

func (self *serviceCenter) serveConn(conn server.Conn) {
	conn.SetForwardRequestChannel(self.fwdChan)
	var err error
	defer func() {
		self.connLeave <- &eventConnLeave{conn: conn, err: err}
	}()
	for {
		var msg *proto.Message
		msg, err = conn.ReadMessage()
		if err != nil {
			return
		}
		self.reportMessage(conn.UniqId(), msg)
	}
}

func (self *serviceCenter) NewConn(conn server.Conn) error {
	usr := conn.Username()
	if len(usr) == 0 || strings.Contains(usr, ":") || strings.Contains(usr, "\n") {
		return fmt.Errorf("[Username=%v] Invalid Username")
	}
	evt := new(eventConnIn)
	ch := make(chan error)

	conn.SetMessageCache(self.config.MsgCache)
	evt.conn = conn
	evt.errChan = ch
	self.connIn <- evt
	err := <-ch
	if err == nil {
		go self.serveConn(conn)
		self.reportLogin(conn.Service(), usr, conn.UniqId())
	}
	return err
}

func newServiceCenter(serviceName string, conf *ServiceConfig, fwdChan chan<- *server.ForwardRequest) *serviceCenter {
	ret := new(serviceCenter)
	ret.config = conf
	if ret.config == nil {
		ret.config = new(ServiceConfig)
	}
	ret.serviceName = serviceName
	ret.fwdChan = fwdChan

	ret.connIn = make(chan *eventConnIn)
	ret.connLeave = make(chan *eventConnLeave)
	ret.writeReqChan = make(chan *writeMessageRequest)
	go ret.process(conf.MaxNrConns, conf.MaxNrConnsPerUser, conf.MaxNrUsers)
	return ret
}
<|endoftext|>


<commit_before>package estafette

import (
	"fmt"
	"strconv"
	"sync"

	"github.com/estafette/estafette-ci-api/cockroach"
	"github.com/rs/zerolog/log"
)

// EventWorker processes events pushed to channels
type EventWorker interface {
	ListenToCiBuilderEventChannels()
	RemoveJobForEstafetteBuild(CiBuilderEvent) error
	UpdateBuildStatus(CiBuilderEvent) error
}

type eventWorkerImpl struct {
	waitGroup              *sync.WaitGroup
	stopChannel            <-chan struct{}
	ciBuilderWorkerPool    chan chan CiBuilderEvent
	ciBuilderClient        CiBuilderClient
	cockroachDBClient      cockroach.DBClient
	ciBuilderEventsChannel chan CiBuilderEvent
}

// NewEstafetteEventWorker returns a new estafette.EventWorker
func NewEstafetteEventWorker(stopChannel <-chan struct{}, waitGroup *sync.WaitGroup, ciBuilderWorkerPool chan chan CiBuilderEvent, ciBuilderClient CiBuilderClient, cockroachDBClient cockroach.DBClient) EventWorker {
	return &eventWorkerImpl{
		waitGroup:              waitGroup,
		stopChannel:            stopChannel,
		ciBuilderWorkerPool:    ciBuilderWorkerPool,
		ciBuilderClient:        ciBuilderClient,
		cockroachDBClient:      cockroachDBClient,
		ciBuilderEventsChannel: make(chan CiBuilderEvent),
	}
}

func (w *eventWorkerImpl) ListenToCiBuilderEventChannels() {
	go func() {
		// handle estafette events via channels
		for {
			// register the current worker into the worker queue.
			w.ciBuilderWorkerPool <- w.ciBuilderEventsChannel

			select {
			case ciBuilderEvent := <-w.ciBuilderEventsChannel:
				go func(ciBuilderEvent CiBuilderEvent) {
					w.waitGroup.Add(1)
					err := w.UpdateBuildStatus(ciBuilderEvent)
					if err != nil {
						log.Error().Err(err).Msgf("Failed updating build status for job %v to %v, not removing the job", ciBuilderEvent.JobName, ciBuilderEvent.BuildStatus)
					} else {
						err = w.RemoveJobForEstafetteBuild(ciBuilderEvent)
						if err != nil {
							log.Error().Err(err).Msgf("Failed removing job %v", ciBuilderEvent.JobName)
						}
					}
					w.waitGroup.Done()
				}(ciBuilderEvent)
			case <-w.stopChannel:
				log.Debug().Msg("Stopping Estafette event worker...")
				return
			}
		}
	}()
}

func (w *eventWorkerImpl) RemoveJobForEstafetteBuild(ciBuilderEvent CiBuilderEvent) (err error) {

	// create ci builder job
	return w.ciBuilderClient.RemoveCiBuilderJob(ciBuilderEvent.JobName)
}

func (w *eventWorkerImpl) UpdateBuildStatus(ciBuilderEvent CiBuilderEvent) (err error) {

	log.Debug().Interface("ciBuilderEvent", ciBuilderEvent).Msgf("UpdateBuildStatus executing...")

	if ciBuilderEvent.BuildStatus != "" && ciBuilderEvent.ReleaseID != "" {

		releaseID, err := strconv.Atoi(ciBuilderEvent.ReleaseID)
		if err != nil {
			return err
		}

		log.Debug().Msgf("Converted release id %v", releaseID)

		err = w.cockroachDBClient.UpdateReleaseStatus(ciBuilderEvent.RepoSource, ciBuilderEvent.RepoOwner, ciBuilderEvent.RepoName, releaseID, ciBuilderEvent.BuildStatus)
		if err != nil {
			return err
		}

		log.Debug().Msgf("Updated release status for job %v to %v", ciBuilderEvent.JobName, ciBuilderEvent.BuildStatus)

		return err

	} else if ciBuilderEvent.BuildStatus != "" && ciBuilderEvent.BuildID != "" {

		buildID, err := strconv.Atoi(ciBuilderEvent.BuildID)
		if err != nil {
			return err
		}

		log.Debug().Msgf("Converted build id %v", buildID)

		err = w.cockroachDBClient.UpdateBuildStatusByID(ciBuilderEvent.RepoSource, ciBuilderEvent.RepoOwner, ciBuilderEvent.RepoName, buildID, ciBuilderEvent.BuildStatus)
		if err != nil {
			return err
		}

		log.Debug().Msgf("Updated build status for job %v to %v", ciBuilderEvent.JobName, ciBuilderEvent.BuildStatus)

		return err

		// check build status for backwards compatibility of builder
	} else if ciBuilderEvent.BuildStatus != "" {

		err := w.cockroachDBClient.UpdateBuildStatus(ciBuilderEvent.RepoSource, ciBuilderEvent.RepoOwner, ciBuilderEvent.RepoName, ciBuilderEvent.RepoBranch, ciBuilderEvent.RepoRevision, ciBuilderEvent.BuildStatus)
		if err != nil {
			return err
		}

		log.Debug().Msgf("Updated build status for job %v", ciBuilderEvent.JobName)

		return err
	}

	return fmt.Errorf("CiBuilderEvent has invalid state, not updating build status")
}
<commit_msg>log event in case of error<commit_after>package estafette

import (
	"fmt"
	"strconv"
	"sync"

	"github.com/estafette/estafette-ci-api/cockroach"
	"github.com/rs/zerolog/log"
)

// EventWorker processes events pushed to channels
type EventWorker interface {
	ListenToCiBuilderEventChannels()
	RemoveJobForEstafetteBuild(CiBuilderEvent) error
	UpdateBuildStatus(CiBuilderEvent) error
}

type eventWorkerImpl struct {
	waitGroup              *sync.WaitGroup
	stopChannel            <-chan struct{}
	ciBuilderWorkerPool    chan chan CiBuilderEvent
	ciBuilderClient        CiBuilderClient
	cockroachDBClient      cockroach.DBClient
	ciBuilderEventsChannel chan CiBuilderEvent
}

// NewEstafetteEventWorker returns a new estafette.EventWorker
func NewEstafetteEventWorker(stopChannel <-chan struct{}, waitGroup *sync.WaitGroup, ciBuilderWorkerPool chan chan CiBuilderEvent, ciBuilderClient CiBuilderClient, cockroachDBClient cockroach.DBClient) EventWorker {
	return &eventWorkerImpl{
		waitGroup:              waitGroup,
		stopChannel:            stopChannel,
		ciBuilderWorkerPool:    ciBuilderWorkerPool,
		ciBuilderClient:        ciBuilderClient,
		cockroachDBClient:      cockroachDBClient,
		ciBuilderEventsChannel: make(chan CiBuilderEvent),
	}
}

func (w *eventWorkerImpl) ListenToCiBuilderEventChannels() {
	go func() {
		// handle estafette events via channels
		for {
			// register the current worker into the worker queue.
			w.ciBuilderWorkerPool <- w.ciBuilderEventsChannel

			select {
			case ciBuilderEvent := <-w.ciBuilderEventsChannel:
				go func(ciBuilderEvent CiBuilderEvent) {
					w.waitGroup.Add(1)
					err := w.UpdateBuildStatus(ciBuilderEvent)
					if err != nil {
						log.Error().Err(err).Interface("ciBuilderEvent", ciBuilderEvent).Msgf("Failed updating build status for job %v to %v, not removing the job", ciBuilderEvent.JobName, ciBuilderEvent.BuildStatus)
					} else {
						err = w.RemoveJobForEstafetteBuild(ciBuilderEvent)
						if err != nil {
							log.Error().Err(err).Interface("ciBuilderEvent", ciBuilderEvent).Msgf("Failed removing job %v", ciBuilderEvent.JobName)
						}
					}
					w.waitGroup.Done()
				}(ciBuilderEvent)
			case <-w.stopChannel:
				log.Debug().Msg("Stopping Estafette event worker...")
				return
			}
		}
	}()
}

func (w *eventWorkerImpl) RemoveJobForEstafetteBuild(ciBuilderEvent CiBuilderEvent) (err error) {

	// create ci builder job
	return w.ciBuilderClient.RemoveCiBuilderJob(ciBuilderEvent.JobName)
}

func (w *eventWorkerImpl) UpdateBuildStatus(ciBuilderEvent CiBuilderEvent) (err error) {

	log.Debug().Interface("ciBuilderEvent", ciBuilderEvent).Msgf("UpdateBuildStatus executing...")

	if ciBuilderEvent.BuildStatus != "" && ciBuilderEvent.ReleaseID != "" {

		releaseID, err := strconv.Atoi(ciBuilderEvent.ReleaseID)
		if err != nil {
			return err
		}

		log.Debug().Msgf("Converted release id %v", releaseID)

		err = w.cockroachDBClient.UpdateReleaseStatus(ciBuilderEvent.RepoSource, ciBuilderEvent.RepoOwner, ciBuilderEvent.RepoName, releaseID, ciBuilderEvent.BuildStatus)
		if err != nil {
			return err
		}

		log.Debug().Msgf("Updated release status for job %v to %v", ciBuilderEvent.JobName, ciBuilderEvent.BuildStatus)

		return err

	} else if ciBuilderEvent.BuildStatus != "" && ciBuilderEvent.BuildID != "" {

		buildID, err := strconv.Atoi(ciBuilderEvent.BuildID)
		if err != nil {
			return err
		}

		log.Debug().Msgf("Converted build id %v", buildID)

		err = w.cockroachDBClient.UpdateBuildStatusByID(ciBuilderEvent.RepoSource, ciBuilderEvent.RepoOwner, ciBuilderEvent.RepoName, buildID, ciBuilderEvent.BuildStatus)
		if err != nil {
			return err
		}

		log.Debug().Msgf("Updated build status for job %v to %v", ciBuilderEvent.JobName, ciBuilderEvent.BuildStatus)

		return err

		// check build status for backwards compatibility of builder
	} else if ciBuilderEvent.BuildStatus != "" {

		err := w.cockroachDBClient.UpdateBuildStatus(ciBuilderEvent.RepoSource, ciBuilderEvent.RepoOwner, ciBuilderEvent.RepoName, ciBuilderEvent.RepoBranch, ciBuilderEvent.RepoRevision, ciBuilderEvent.BuildStatus)
		if err != nil {
			return err
		}

		log.Debug().Msgf("Updated build status for job %v", ciBuilderEvent.JobName)

		return err
	}

	return fmt.Errorf("CiBuilderEvent has invalid state, not updating build status")
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package goapp

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"appengine/datastore"
	"appengine/urlfetch"
	"appengine/user"
	mpg "github.com/MiniProfiler/go/miniprofiler_gae"
	"github.com/mjibson/goon"
)

type Plan struct {
	Id, Name, Desc string
	Amount         int
}

// parent: User, key: 1
type UserCharge struct {
	_kind  string         `goon:"kind,UC"`
	Id     int64          `datastore:"-" goon:"id"`
	Parent *datastore.Key `datastore:"-" goon:"parent"`

	Customer string    `datastore:"c,noindex" json:"-"`
	Created  time.Time `datastore:"r,noindex"`
	Last4    string    `datastore:"l,noindex" json:"-"`
	Next     time.Time `datastore:"n,noindex"`
	Amount   int       `datastore:"a,noindex"`
	Interval string    `datastore:"i,noindex"`
	Plan     string    `datastore:"p,noindex"`
}

type StripeCustomer struct {
	Id      string `json:"id"`
	Created int64  `json:"created"`
	Card    struct {
		Last4 string `json:"last4"`
	} `json:"active_card"`
	Subscription struct {
		Plan struct {
			Interval string `json:"interval"`
			Id       string `json:"id"`
			Amount   int    `json:"amount"`
		} `json:"plan"`
		End int64 `json:"current_period_end"`
	} `json:"subscription"`
}

type StripeError struct {
	Error struct {
		Message string `json:"message"`
	} `json:"error"`
}

func Charge(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	cu := user.Current(c)
	gn := goon.FromContext(c)
	u := User{Id: cu.ID}
	uc := &UserCharge{Id: 1, Parent: gn.Key(&u)}
	if err := gn.Get(&u); err != nil {
		serveError(w, err)
		return
	} else if u.Account != AFree {
		serveError(w, fmt.Errorf("You're already subscribed."))
		return
	}
	if err := gn.Get(uc); err == nil && len(uc.Customer) > 0 {
		serveError(w, fmt.Errorf("You're already subscribed."))
		return
	} else if err != datastore.ErrNoSuchEntity {
		serveError(w, err)
		return
	}
	resp, err := stripe(c, "POST", "customers", url.Values{
		"email":       {u.Email},
		"description": {u.Id},
		"card":        {r.FormValue("token")},
		"plan":        {r.FormValue("plan")},
	}.Encode())
	if err != nil {
		serveError(w, err)
		return
	} else if resp.StatusCode != http.StatusOK {
		var se StripeError
		defer resp.Body.Close()
		b, _ := ioutil.ReadAll(resp.Body)
		if err := json.Unmarshal(b, &se); err == nil {
			serveError(w, fmt.Errorf(se.Error.Message))
		} else {
			serveError(w, fmt.Errorf("Error"))
		}
		c.Errorf("status: %v, %s", resp.StatusCode, b)
		return
	}
	uc, err = setCharge(c, resp)
	if err != nil {
		serveError(w, err)
		return
	}
	b, _ := json.Marshal(&uc)
	w.Write(b)
}

func setCharge(c mpg.Context, r *http.Response) (*UserCharge, error) {
	var sc StripeCustomer
	defer r.Body.Close()
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return nil, err
	}
	if err := json.Unmarshal(b, &sc); err != nil {
		return nil, err
	}
	cu := user.Current(c)
	gn := goon.FromContext(c)
	u := User{Id: cu.ID}
	uc := UserCharge{Id: 1, Parent: gn.Key(&u)}
	if err := gn.RunInTransaction(func(gn *goon.Goon) error {
		if err := gn.Get(&u); err != nil && err != datastore.ErrNoSuchEntity {
			return err
		}
		if err := gn.Get(&uc); err != nil && err != datastore.ErrNoSuchEntity {
			return err
		}
		u.Account = APaid
		uc.Customer = sc.Id
		uc.Last4 = sc.Card.Last4
		uc.Created = time.Unix(sc.Created, 0)
		uc.Next = time.Unix(sc.Subscription.End, 0)
		uc.Amount = sc.Subscription.Plan.Amount
		uc.Interval = sc.Subscription.Plan.Interval
		uc.Plan = sc.Subscription.Plan.Id
		_, err := gn.PutMulti([]interface{}{&u, &uc})
		return err
	}, nil); err != nil {
		return nil, err
	}
	return &uc, nil
}

func Account(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	cu := user.Current(c)
	gn := goon.FromContext(c)
	u := User{Id: cu.ID}
	uc := &UserCharge{Id: 1, Parent: gn.Key(&u)}
	if err := gn.Get(uc); err == nil {
		if uc.Next.Before(time.Now()) {
			if resp, err := stripe(c, "GET", "customers/"+uc.Customer, ""); err == nil {
				if nuc, err := setCharge(c, resp); err == nil {
					uc = nuc
					c.Infof("updated user charge %v", cu.ID)
				}
			}
		}
		b, _ := json.Marshal(&uc)
		w.Write(b)
	} else if err != datastore.ErrNoSuchEntity {
		serveError(w, err)
		return
	}
}

func Uncheckout(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	uc, err := doUncheckout(c)
	if err != nil {
		serveError(w, err)
		return
	}
	b, _ := json.Marshal(uc)
	w.Write(b)
}

func doUncheckout(c mpg.Context) (*UserCharge, error) {
	cu := user.Current(c)
	gn := goon.FromContext(c)
	u := User{Id: cu.ID}
	uc := UserCharge{Id: 1, Parent: gn.Key(&u)}
	if err := gn.Get(&u); err != nil {
		return nil, err
	}
	if err := gn.Get(&uc); err != nil || len(uc.Customer) == 0 {
		return nil, err
	}
	resp, err := stripe(c, "DELETE", "customers/"+uc.Customer, "")
	if err != nil {
		return nil, err
	} else if resp.StatusCode != http.StatusOK {
		c.Errorf("%s", resp.Body)
		c.Errorf("stripe delete error, but proceeding")
	}
	if err := gn.RunInTransaction(func(gn *goon.Goon) error {
		if err := gn.Get(&u); err != nil && err != datastore.ErrNoSuchEntity {
			return err
		}
		u.Account = AFree
		u.Until = uc.Next
		if err := gn.Delete(gn.Key(&uc)); err != nil {
			return err
		}
		_, err := gn.Put(&u)
		return err
	}, nil); err != nil {
		return nil, err
	}
	return &uc, nil
}

func stripe(c mpg.Context, method, urlStr, body string) (*http.Response, error) {
	cl := &http.Client{
		Transport: &urlfetch.Transport{
			Context:  c,
			Deadline: time.Minute,
		},
	}
	req, err := http.NewRequest(method, fmt.Sprintf("https://api.stripe.com/v1/%s", urlStr), strings.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.SetBasicAuth(STRIPE_SECRET, "")
	return cl.Do(req)
}
<commit_msg>Never set Until to an earlier date<commit_after>/*
 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package goapp

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"appengine/datastore"
	"appengine/urlfetch"
	"appengine/user"
	mpg "github.com/MiniProfiler/go/miniprofiler_gae"
	"github.com/mjibson/goon"
)

type Plan struct {
	Id, Name, Desc string
	Amount         int
}

// parent: User, key: 1
type UserCharge struct {
	_kind  string         `goon:"kind,UC"`
	Id     int64          `datastore:"-" goon:"id"`
	Parent *datastore.Key `datastore:"-" goon:"parent"`

	Customer string    `datastore:"c,noindex" json:"-"`
	Created  time.Time `datastore:"r,noindex"`
	Last4    string    `datastore:"l,noindex" json:"-"`
	Next     time.Time `datastore:"n,noindex"`
	Amount   int       `datastore:"a,noindex"`
	Interval string    `datastore:"i,noindex"`
	Plan     string    `datastore:"p,noindex"`
}

type StripeCustomer struct {
	Id      string `json:"id"`
	Created int64  `json:"created"`
	Card    struct {
		Last4 string `json:"last4"`
	} `json:"active_card"`
	Subscription struct {
		Plan struct {
			Interval string `json:"interval"`
			Id       string `json:"id"`
			Amount   int    `json:"amount"`
		} `json:"plan"`
		End int64 `json:"current_period_end"`
	} `json:"subscription"`
}

type StripeError struct {
	Error struct {
		Message string `json:"message"`
	} `json:"error"`
}

func Charge(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	cu := user.Current(c)
	gn := goon.FromContext(c)
	u := User{Id: cu.ID}
	uc := &UserCharge{Id: 1, Parent: gn.Key(&u)}
	if err := gn.Get(&u); err != nil {
		serveError(w, err)
		return
	} else if u.Account != AFree {
		serveError(w, fmt.Errorf("You're already subscribed."))
		return
	}
	if err := gn.Get(uc); err == nil && len(uc.Customer) > 0 {
		serveError(w, fmt.Errorf("You're already subscribed."))
		return
	} else if err != datastore.ErrNoSuchEntity {
		serveError(w, err)
		return
	}
	resp, err := stripe(c, "POST", "customers", url.Values{
		"email":       {u.Email},
		"description": {u.Id},
		"card":        {r.FormValue("token")},
		"plan":        {r.FormValue("plan")},
	}.Encode())
	if err != nil {
		serveError(w, err)
		return
	} else if resp.StatusCode != http.StatusOK {
		var se StripeError
		defer resp.Body.Close()
		b, _ := ioutil.ReadAll(resp.Body)
		if err := json.Unmarshal(b, &se); err == nil {
			serveError(w, fmt.Errorf(se.Error.Message))
		} else {
			serveError(w, fmt.Errorf("Error"))
		}
		c.Errorf("status: %v, %s", resp.StatusCode, b)
		return
	}
	uc, err = setCharge(c, resp)
	if err != nil {
		serveError(w, err)
		return
	}
	b, _ := json.Marshal(&uc)
	w.Write(b)
}

func setCharge(c mpg.Context, r *http.Response) (*UserCharge, error) {
	var sc StripeCustomer
	defer r.Body.Close()
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return nil, err
	}
	if err := json.Unmarshal(b, &sc); err != nil {
		return nil, err
	}
	cu := user.Current(c)
	gn := goon.FromContext(c)
	u := User{Id: cu.ID}
	uc := UserCharge{Id: 1, Parent: gn.Key(&u)}
	if err := gn.RunInTransaction(func(gn *goon.Goon) error {
		if err := gn.Get(&u); err != nil && err != datastore.ErrNoSuchEntity {
			return err
		}
		if err := gn.Get(&uc); err != nil && err != datastore.ErrNoSuchEntity {
			return err
		}
		u.Account = APaid
		uc.Customer = sc.Id
		uc.Last4 = sc.Card.Last4
		uc.Created = time.Unix(sc.Created, 0)
		uc.Next = time.Unix(sc.Subscription.End, 0)
		uc.Amount = sc.Subscription.Plan.Amount
		uc.Interval = sc.Subscription.Plan.Interval
		uc.Plan = sc.Subscription.Plan.Id
		_, err := gn.PutMulti([]interface{}{&u, &uc})
		return err
	}, nil); err != nil {
		return nil, err
	}
	return &uc, nil
}

func Account(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	cu := user.Current(c)
	gn := goon.FromContext(c)
	u := User{Id: cu.ID}
	uc := &UserCharge{Id: 1, Parent: gn.Key(&u)}
	if err := gn.Get(uc); err == nil {
		if uc.Next.Before(time.Now()) {
			if resp, err := stripe(c, "GET", "customers/"+uc.Customer, ""); err == nil {
				if nuc, err := setCharge(c, resp); err == nil {
					uc = nuc
					c.Infof("updated user charge %v", cu.ID)
				}
			}
		}
		b, _ := json.Marshal(&uc)
		w.Write(b)
	} else if err != datastore.ErrNoSuchEntity {
		serveError(w, err)
		return
	}
}

func Uncheckout(c mpg.Context, w http.ResponseWriter, r *http.Request) {
	uc, err := doUncheckout(c)
	if err != nil {
		serveError(w, err)
		return
	}
	b, _ := json.Marshal(uc)
	w.Write(b)
}

func doUncheckout(c mpg.Context) (*UserCharge, error) {
	cu := user.Current(c)
	gn := goon.FromContext(c)
	u := User{Id: cu.ID}
	uc := UserCharge{Id: 1, Parent: gn.Key(&u)}
	if err := gn.Get(&u); err != nil {
		return nil, err
	}
	if err := gn.Get(&uc); err != nil || len(uc.Customer) == 0 {
		return nil, err
	}
	resp, err := stripe(c, "DELETE", "customers/"+uc.Customer, "")
	if err != nil {
		return nil, err
	} else if resp.StatusCode != http.StatusOK {
		c.Errorf("%s", resp.Body)
		c.Errorf("stripe delete error, but proceeding")
	}
	if err := gn.RunInTransaction(func(gn *goon.Goon) error {
		if err := gn.Get(&u); err != nil && err != datastore.ErrNoSuchEntity {
			return err
		}
		u.Account = AFree
		if uc.Next.After(u.Until) {
			u.Until = uc.Next
		}
		if err := gn.Delete(gn.Key(&uc)); err != nil {
			return err
		}
		_, err := gn.Put(&u)
		return err
	}, nil); err != nil {
		return nil, err
	}
	return &uc, nil
}

func stripe(c mpg.Context, method, urlStr, body string) (*http.Response, error) {
	cl := &http.Client{
		Transport: &urlfetch.Transport{
			Context:  c,
			Deadline: time.Minute,
		},
	}
	req, err := http.NewRequest(method, fmt.Sprintf("https://api.stripe.com/v1/%s", urlStr), strings.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.SetBasicAuth(STRIPE_SECRET, "")
	return cl.Do(req)
}
<|endoftext|>


<commit_before>package spvwallet

import (
	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/peer"
	"github.com/btcsuite/btcd/txscript"
	btc "github.com/btcsuite/btcutil"
	hd "github.com/btcsuite/btcutil/hdkeychain"
	"github.com/op/go-logging"
	b39 "github.com/tyler-smith/go-bip39"
	"net"
	"os"
	"path"
	"sync"
)

type SPVWallet struct {
	params *chaincfg.Params

	masterPrivateKey *hd.ExtendedKey
	masterPublicKey  *hd.ExtendedKey

	maxFee      uint64
	priorityFee uint64
	normalFee   uint64
	economicFee uint64
	feeAPI      string

	repoPath string

	blockchain  *Blockchain
	txstore     *TxStore
	PeerManager *PeerManager

	fPositives    chan *peer.Peer
	stopChan      chan int
	fpAccumulator map[int32]int32
	blockQueue    chan chainhash.Hash
	toDownload    map[chainhash.Hash]int32
	mutex         *sync.RWMutex

	running bool

	config *Config
}

var log = logging.MustGetLogger("bitcoin")

const WALLET_VERSION = "0.1.0"

func NewSPVWallet(mnemonic string, params *chaincfg.Params, maxFee uint64, lowFee uint64, mediumFee uint64, highFee uint64, feeApi,
	repoPath string, db Datastore, userAgent string, trustedPeer string, logger logging.LeveledBackend) (*SPVWallet, error) {

	log.SetBackend(logger)

	seed := b39.NewSeed(mnemonic, "")

	mPrivKey, err := hd.NewMaster(seed, params)
	if err != nil {
		return nil, err
	}
	mPubKey, err := mPrivKey.Neuter()
	if err != nil {
		return nil, err
	}

	w := &SPVWallet{
		repoPath:         repoPath,
		masterPrivateKey: mPrivKey,
		masterPublicKey:  mPubKey,
		params:           params,
		maxFee:           maxFee,
		priorityFee:      highFee,
		normalFee:        mediumFee,
		economicFee:      lowFee,
		feeAPI:           feeApi,
		fPositives:       make(chan *peer.Peer),
		stopChan:         make(chan int),
		fpAccumulator:    make(map[int32]int32),
		blockQueue:       make(chan chainhash.Hash, 32),
		toDownload:       make(map[chainhash.Hash]int32),
		mutex:            new(sync.RWMutex),
	}

	w.txstore, err = NewTxStore(w.params, db, w.masterPrivateKey)
	if err != nil {
		return nil, err
	}
	w.blockchain, err = NewBlockchain(w.repoPath, w.params)
	if err != nil {
		return nil, err
	}

	listeners := &peer.MessageListeners{
		OnMerkleBlock: w.onMerkleBlock,
		OnInv:         w.onInv,
		OnTx:          w.onTx,
		OnGetData:     w.onGetData,
	}

	getNewestBlock := func() (*chainhash.Hash, int32, error) {
		storedHeader, err := w.blockchain.db.GetBestHeader()
		if err != nil {
			return nil, 0, err
		}
		height, err := w.blockchain.db.Height()
		if err != nil {
			return nil, 0, err
		}
		hash := storedHeader.header.BlockHash()
		return &hash, int32(height), nil
	}

	w.config = &Config{
		UserAgentName:      userAgent,
		UserAgentVersion:   WALLET_VERSION,
		Params:             w.params,
		AddressCacheDir:    repoPath,
		GetFilter:          w.txstore.GimmeFilter,
		StartChainDownload: w.startChainDownload,
		GetNewestBlock:     getNewestBlock,
		Listeners:          listeners,
	}

	if trustedPeer != "" {
		addr, err := net.ResolveTCPAddr("tcp", trustedPeer)
		if err != nil {
			return nil, err
		}
		w.config.TrustedPeer = addr
	}

	w.PeerManager, err = NewPeerManager(w.config)
	if err != nil {
		return nil, err
	}

	return w, nil
}

func (w *SPVWallet) Start() {
	go w.PeerManager.Start()
	go w.fPositiveHandler(w.stopChan)
	w.running = true
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//
// API
//
//////////////

func (w *SPVWallet) CurrencyCode() string {
	if w.params.Name == chaincfg.MainNetParams.Name {
		return "btc"
	} else {
		return "tbtc"
	}
}

func (w *SPVWallet) MasterPrivateKey() *hd.ExtendedKey {
	return w.masterPrivateKey
}

func (w *SPVWallet) MasterPublicKey() *hd.ExtendedKey {
	return w.masterPublicKey
}

func (w *SPVWallet) CurrentAddress(purpose KeyPurpose) btc.Address {
	key := w.txstore.GetCurrentKey(purpose)
	addr, _ := key.Address(w.params)
	return btc.Address(addr)
}

func (w *SPVWallet) HasKey(addr btc.Address) bool {
	script, err := txscript.PayToAddrScript(addr)
	if err != nil {
		return false
	}
	_, err = w.txstore.GetKeyForScript(script)
	if err != nil {
		return false
	}
	return true
}

func (w *SPVWallet) Balance() (confirmed, unconfirmed int64) {
	utxos, _ := w.txstore.Utxos().GetAll()
	stxos, _ := w.txstore.Stxos().GetAll()
	for _, utxo := range utxos {
		if !utxo.Freeze {
			if utxo.AtHeight > 0 {
				confirmed += utxo.Value
			} else {
				if w.checkIfStxoIsConfirmed(utxo, stxos) {
					confirmed += utxo.Value
				} else {
					unconfirmed += utxo.Value
				}
			}
		}
	}
	return confirmed, unconfirmed
}

func (w *SPVWallet) Transactions() ([]Txn, error) {
	return w.txstore.Txns().GetAll()
}

func (w *SPVWallet) checkIfStxoIsConfirmed(utxo Utxo, stxos []Stxo) bool {
	for _, stxo := range stxos {
		if stxo.SpendTxid.IsEqual(&utxo.Op.Hash) {
			if stxo.Utxo.AtHeight > 0 {
				return true
			} else {
				return w.checkIfStxoIsConfirmed(stxo.Utxo, stxos)
			}
		}
	}
	return false
}

func (w *SPVWallet) Params() *chaincfg.Params {
	return w.params
}

func (w *SPVWallet) AcceptStealth() bool {
	return true
}

func (w *SPVWallet) AddTransactionListener(callback func(TransactionCallback)) {
	w.txstore.listeners = append(w.txstore.listeners, callback)
}

func (w *SPVWallet) ChainTip() uint32 {
	height, _ := w.blockchain.db.Height()
	return uint32(height)
}

func (w *SPVWallet) AddWatchedScript(script []byte) error {
	err := w.txstore.WatchedScripts().Put(script)
	w.txstore.PopulateAdrs()

	for _, peer := range w.PeerManager.ConnectedPeers() {
		w.updateFilterAndSend(peer)
	}
	return err
}

func (w *SPVWallet) GenerateMultisigScript(keys []hd.ExtendedKey, threshold int) (addr btc.Address, redeemScript []byte, err error) {
	var addrPubKeys []*btc.AddressPubKey
	for _, key := range keys {
		ecKey, err := key.ECPubKey()
		if err != nil {
			return nil, nil, err
		}
		k, err := btc.NewAddressPubKey(ecKey.SerializeCompressed(), w.params)
		if err != nil {
			return nil, nil, err
		}
		addrPubKeys = append(addrPubKeys, k)
	}
	redeemScript, err = txscript.MultiSigScript(addrPubKeys, threshold)
	if err != nil {
		return nil, nil, err
	}
	addr, err = btc.NewAddressScriptHash(redeemScript, w.params)
	if err != nil {
		return nil, nil, err
	}
	return addr, redeemScript, nil
}

func (w *SPVWallet) Close() {
	if w.running {
		log.Info("Disconnecting from peers and shutting down")
		w.PeerManager.Stop()
		w.blockchain.Close()
		w.stopChan <- 1
		w.running = false
	}
}

func (w *SPVWallet) ReSyncBlockchain(fromHeight int32) {
	w.Close()
	os.Remove(path.Join(w.repoPath, "headers.bin"))
	blockchain, err := NewBlockchain(w.repoPath, w.params)
	if err != nil {
		return
	}
	w.blockchain = blockchain
	go w.Start()
}
<commit_msg>Update spvwallet<commit_after>package spvwallet

import (
	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/peer"
	"github.com/btcsuite/btcd/txscript"
	btc "github.com/btcsuite/btcutil"
	hd "github.com/btcsuite/btcutil/hdkeychain"
	"github.com/op/go-logging"
	b39 "github.com/tyler-smith/go-bip39"
	"net"
	"os"
	"path"
	"sync"
)

type SPVWallet struct {
	params *chaincfg.Params

	masterPrivateKey *hd.ExtendedKey
	masterPublicKey  *hd.ExtendedKey

	maxFee      uint64
	priorityFee uint64
	normalFee   uint64
	economicFee uint64
	feeAPI      string

	repoPath string

	blockchain  *Blockchain
	txstore     *TxStore
	PeerManager *PeerManager

	fPositives    chan *peer.Peer
	stopChan      chan int
	fpAccumulator map[int32]int32
	blockQueue    chan chainhash.Hash
	toDownload    map[chainhash.Hash]int32
	mutex         *sync.RWMutex

	running bool

	config *Config
}

var log = logging.MustGetLogger("bitcoin")

const WALLET_VERSION = "0.1.0"

func NewSPVWallet(mnemonic string, params *chaincfg.Params, maxFee uint64, lowFee uint64, mediumFee uint64, highFee uint64, feeApi,
	repoPath string, db Datastore, userAgent string, trustedPeer string, logger logging.LeveledBackend) (*SPVWallet, error) {

	log.SetBackend(logger)

	seed := b39.NewSeed(mnemonic, "")

	mPrivKey, err := hd.NewMaster(seed, params)
	if err != nil {
		return nil, err
	}
	mPubKey, err := mPrivKey.Neuter()
	if err != nil {
		return nil, err
	}

	w := &SPVWallet{
		repoPath:         repoPath,
		masterPrivateKey: mPrivKey,
		masterPublicKey:  mPubKey,
		params:           params,
		maxFee:           maxFee,
		priorityFee:      highFee,
		normalFee:        mediumFee,
		economicFee:      lowFee,
		feeAPI:           feeApi,
		fPositives:       make(chan *peer.Peer),
		stopChan:         make(chan int),
		fpAccumulator:    make(map[int32]int32),
		blockQueue:       make(chan chainhash.Hash, 32),
		toDownload:       make(map[chainhash.Hash]int32),
		mutex:            new(sync.RWMutex),
	}

	w.txstore, err = NewTxStore(w.params, db, w.masterPrivateKey)
	if err != nil {
		return nil, err
	}
	w.blockchain, err = NewBlockchain(w.repoPath, w.params)
	if err != nil {
		return nil, err
	}

	listeners := &peer.MessageListeners{
		OnMerkleBlock: w.onMerkleBlock,
		OnInv:         w.onInv,
		OnTx:          w.onTx,
		OnGetData:     w.onGetData,
	}

	getNewestBlock := func() (*chainhash.Hash, int32, error) {
		storedHeader, err := w.blockchain.db.GetBestHeader()
		if err != nil {
			return nil, 0, err
		}
		height, err := w.blockchain.db.Height()
		if err != nil {
			return nil, 0, err
		}
		hash := storedHeader.header.BlockHash()
		return &hash, int32(height), nil
	}

	w.config = &Config{
		UserAgentName:      userAgent,
		UserAgentVersion:   WALLET_VERSION,
		Params:             w.params,
		AddressCacheDir:    repoPath,
		GetFilter:          w.txstore.GimmeFilter,
		StartChainDownload: w.startChainDownload,
		GetNewestBlock:     getNewestBlock,
		Listeners:          listeners,
	}

	if trustedPeer != "" {
		addr, err := net.ResolveTCPAddr("tcp", trustedPeer)
		if err != nil {
			return nil, err
		}
		w.config.TrustedPeer = addr
	}

	w.PeerManager, err = NewPeerManager(w.config)
	if err != nil {
		return nil, err
	}

	return w, nil
}

func (w *SPVWallet) Start() {
	go w.PeerManager.Start()
	go w.fPositiveHandler(w.stopChan)
	w.running = true
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//
// API
//
//////////////

func (w *SPVWallet) CurrencyCode() string {
	if w.params.Name == chaincfg.MainNetParams.Name {
		return "btc"
	} else {
		return "tbtc"
	}
}

func (w *SPVWallet) MasterPrivateKey() *hd.ExtendedKey {
	return w.masterPrivateKey
}

func (w *SPVWallet) MasterPublicKey() *hd.ExtendedKey {
	return w.masterPublicKey
}

func (w *SPVWallet) CurrentAddress(purpose KeyPurpose) btc.Address {
	key := w.txstore.GetCurrentKey(purpose)
	addr, _ := key.Address(w.params)
	return btc.Address(addr)
}

func (w *SPVWallet) HasKey(addr btc.Address) bool {
	script, err := txscript.PayToAddrScript(addr)
	if err != nil {
		return false
	}
	_, err = w.txstore.GetKeyForScript(script)
	if err != nil {
		return false
	}
	return true
}

func (w *SPVWallet) Balance() (confirmed, unconfirmed int64) {
	utxos, _ := w.txstore.Utxos().GetAll()
	stxos, _ := w.txstore.Stxos().GetAll()
	for _, utxo := range utxos {
		if !utxo.Freeze {
			if utxo.AtHeight > 0 {
				confirmed += utxo.Value
			} else {
				if w.checkIfStxoIsConfirmed(utxo, stxos) {
					confirmed += utxo.Value
				} else {
					unconfirmed += utxo.Value
				}
			}
		}
	}
	return confirmed, unconfirmed
}

func (w *SPVWallet) Transactions() ([]Txn, error) {
	return w.txstore.Txns().GetAll()
}

func (w *SPVWallet) checkIfStxoIsConfirmed(utxo Utxo, stxos []Stxo) bool {
	for _, stxo := range stxos {
		if stxo.SpendTxid.IsEqual(&utxo.Op.Hash) {
			if stxo.Utxo.AtHeight > 0 {
				return true
			} else {
				return w.checkIfStxoIsConfirmed(stxo.Utxo, stxos)
			}
		}
	}
	return false
}

func (w *SPVWallet) Params() *chaincfg.Params {
	return w.params
}

func (w *SPVWallet) AcceptStealth() bool {
	return true
}

func (w *SPVWallet) AddTransactionListener(callback func(TransactionCallback)) {
	w.txstore.listeners = append(w.txstore.listeners, callback)
}

func (w *SPVWallet) ChainTip() uint32 {
	height, _ := w.blockchain.db.Height()
	return uint32(height)
}

func (w *SPVWallet) AddWatchedScript(script []byte) error {
	err := w.txstore.WatchedScripts().Put(script)
	w.txstore.PopulateAdrs()

	for _, peer := range w.PeerManager.ConnectedPeers() {
		w.updateFilterAndSend(peer)
	}
	return err
}

func (w *SPVWallet) GenerateMultisigScript(keys []hd.ExtendedKey, threshold int) (addr btc.Address, redeemScript []byte, err error) {
	var addrPubKeys []*btc.AddressPubKey
	for _, key := range keys {
		ecKey, err := key.ECPubKey()
		if err != nil {
			return nil, nil, err
		}
		k, err := btc.NewAddressPubKey(ecKey.SerializeCompressed(), w.params)
		if err != nil {
			return nil, nil, err
		}
		addrPubKeys = append(addrPubKeys, k)
	}
	redeemScript, err = txscript.MultiSigScript(addrPubKeys, threshold)
	if err != nil {
		return nil, nil, err
	}
	addr, err = btc.NewAddressScriptHash(redeemScript, w.params)
	if err != nil {
		return nil, nil, err
	}
	return addr, redeemScript, nil
}

func (w *SPVWallet) Close() {
	if w.running {
		log.Info("Disconnecting from peers and shutting down")
		w.PeerManager.Stop()
		w.blockchain.Close()
		w.stopChan <- 1
		w.running = false
	}
}

func (w *SPVWallet) ReSyncBlockchain(fromHeight int32) {
	w.Close()
	os.Remove(path.Join(w.repoPath, "headers.bin"))
	blockchain, err := NewBlockchain(w.repoPath, w.params)
	if err != nil {
		return
	}
	w.blockchain = blockchain
	w.PeerManager, err = NewPeerManager(w.config)
	if err != nil {
		return
	}
	w.blockQueue = make(chan chainhash.Hash, 32)
	go w.Start()
}
<|endoftext|>


<commit_before>//
// Copyright © 2012-2017 Guy M. Allard
//
// 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 stompngo

import (
	"testing"
)

/*
	Data Test: Headers Basic.
*/
func TestHeadersBasic(t *testing.T) {
	k := "keya"
	v := "valuea"
	h := Headers{k, v}
	if nil != h.Validate() {
		t.Fatalf("Header validate error: [%v]\n", h.Validate())
	}
	if len(h) != 2 {
		t.Fatalf("Header Unexpected length error 1, length: [%v]\n", len(h))
	}
	h = h.Add("keyb", "valueb").Add("keya", "valuea2")
	if len(h) != 6 {
		t.Fatalf("Header Unexpected length error 2, length after add: [%v]\n", len(h))
	}
	if _, ok := h.Contains(k); !ok {
		t.Fatalf("Header Unexpected false for key: [%v]\n", k)
	}
	k = "xyz"
	if _, ok := h.Contains(k); ok {
		t.Fatalf("Header Unexpected true for key: [%v]\n", k)
	}
	//
	h = Headers{k}
	if e = h.Validate(); e != EHDRLEN {
		t.Fatalf("Header Validate, got [%v], expected [%v]\n", e, EHDRLEN)
	}
}

/*
	Data Test: Headers UTF8.
*/
func TestHeadersUTF8(t *testing.T) {
	k := "keya"
	v := "valuea"
	wh := Headers{k, v}
	var e error   // An error
	var rs string // Result string
	if rs, e = wh.ValidateUTF8(); e != nil {
		t.Fatalf("Unexpected UTF8 error 1: [%v]\n", e)
	}
	if rs != "" {
		t.Fatalf("Unexpected UTF8 error 1B, got [%v], expected [%v]\n", rs, "")
	}
	//
	wh = Headers{k, v, `“Iñtërnâtiônàlizætiøn”`, "valueb", "keyc", `“Iñtërnâtiônàlizætiøn”`}
	if _, e = wh.ValidateUTF8(); e != nil {
		t.Fatalf("Unexpected UTF8 error 2: [%v]\n", e)
	}
	//
	wh = Headers{k, v, `“Iñtërnâtiônàlizætiøn”`, "\x80", "keyc", `“Iñtërnâtiônàlizætiøn”`}
	if rs, e = wh.ValidateUTF8(); e == nil {
		t.Fatalf("Unexpected UTF8 error  3, got nil, expected an error")
	}
	if e != EHDRUTF8 {
		t.Fatalf("Unexpected UTF8 error 4, got [%v], expected [%v]\n", e, EHDRUTF8)
	}
	if rs != "\x80" {
		t.Fatalf("Unexpected UTF8 error 5, got [%v], expected [%v]\n", rs, "\x80")
	}
}

/*.
Data Test: Headers Clone
*/
func TestHeadersClone(t *testing.T) {
	wh := Headers{"ka", "va"}.Add("kb", "vb").Add("kc", "vc")
	hc := wh.Clone()
	if !wh.Compare(hc) {
		t.Fatalf("Unexpected false for clone: [%v], [%v]\n", wh, hc)
	}
}

/*
	Data Test: Headers Add / Delete.
*/
func TestHeadersAddDelete(t *testing.T) {
	ha := Headers{"ka", "va", "kb", "vb", "kc", "vc"}
	hb := Headers{"kaa", "va", "kbb", "vb", "kcc", "vc"}
	hn := ha.AddHeaders(hb)
	if len(ha)+len(hb) != len(hn) {
		t.Fatalf("Unexpected length AddHeaders, got: [%v], expected: [%v]\n", len(hn), len(ha)+len(hb))
	}
	ol := len(hn)
	hn = hn.Delete("ka")
	if len(hn) != ol-2 {
		t.Fatalf("Unexpected length Delete 1, got: [%v], expected: [%v]\n", len(hn), ol-2)
	}
	hn = hn.Delete("kcc")
	if len(hn) != ol-4 {
		t.Fatalf("Unexpected length Delete 2, got: [%v], expected: [%v]\n", len(hn), ol-4)
	}
}

/*
	Data Test: Headers ContainsKV
*/
func TestHeadersContainsKV(t *testing.T) {
	ha := Headers{"ka", "va", "kb", "vb", "kc", "vc"}
	b := ha.ContainsKV("kb", "vb")
	if !b {
		t.Fatalf("KV01 got false, expected true")
	}
	b = ha.ContainsKV("kb", "zz")
	if b {
		t.Fatalf("KV02 got true, expected false")
	}
}

/*
	Data Test: Headers Compare
*/
func TestHeadersCompare(t *testing.T) {
	ha := Headers{"ka", "va", "kb", "vb", "kc", "vc"}
	hb := Headers{"ka", "va", "kb", "vb", "kc", "vc"}
	hc := Headers{"ka", "va"}
	hd := Headers{"k1", "v1", "k2", "v2", "k3", "v3"}
	b := ha.Compare(hb)
	if !b {
		t.Fatalf("CMP01 Expected true, got false")
	}
	b = ha.Compare(hc)
	if b {
		t.Fatalf("CMP02 Expected false, got true")
	}
	b = ha.Compare(hd)
	if b {
		t.Fatalf("CMP03 Expected false, got true")
	}
	b = hd.Compare(ha)
	if b {
		t.Fatalf("CMP04 Expected false, got true")
	}
}

/*
	Data Test: Headers Size
*/
func TestHeadersSize(t *testing.T) {
	ha := Headers{"k", "v"}
	s := ha.Size(false)
	var w int64 = 4
	if s != w {
		t.Fatalf("SIZ01 size, got [%d], expected [%v]\n", s, w)
	}
	//
	ha = Headers{"kaa", "vaa2", "kba", "vba2", "kca", "vca2"}
	s = ha.Size(true)
	w = 3 + 1 + 4 + 1 + 3 + 1 + 4 + 1 + 3 + 1 + 4 + 1
	if s != w {
		t.Fatalf("SIZ02 size, got [%d] expected [%v]\n", s, w)
	}
}

/*
	Data Test: Empty Header Key / Value
*/
func TestHeadersEmtKV(t *testing.T) {
	wh := Headers{"a", "b", "c", "d"} // work headers
	ek := Headers{"a", "b", "", "d"}  // empty key
	ev := Headers{"a", "", "c", "d"}  // empty value
	//
	e = checkHeaders(wh, SPL_10)
	if e != nil {
		t.Fatalf("CHD01 Expected [nil], got [%v]\n", e)
	}
	e = checkHeaders(wh, SPL_11)
	if e != nil {
		t.Fatalf("CHD02 Expected [nil], got [%v]\n", e)
	}
	//
	e = checkHeaders(ek, SPL_10)
	if e != EHDRMTK {
		t.Fatalf("CHD03 Expected [%v], got [%v]\n", EHDRMTK, e)
	}
	e = checkHeaders(ek, SPL_11)
	if e != EHDRMTK {
		t.Fatalf("CHD04 Expected [%v], got [%v]\n", EHDRMTK, e)
	}
	//
	e = checkHeaders(ev, SPL_10)
	if e != EHDRMTV {
		t.Fatalf("CHD05 Expected [%v], got [%v]\n", EHDRMTV, e)
	}
	e = checkHeaders(ev, SPL_11)
	if e != nil {
		t.Fatalf("CHD06 Expected [nil], got [%v]\n", e)
	}
}
<commit_msg>Rework headers tests.<commit_after>//
// Copyright © 2012-2017 Guy M. Allard
//
// 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 stompngo

import (
	"testing"
)

/*
	Data Test: Headers Basic.
*/
func TestHeadersBasic(t *testing.T) {
	for _, _ = range Protocols() {
		k := "keya"
		v := "valuea"
		h := Headers{k, v}
		if nil != h.Validate() {
			t.Fatalf("Header validate error: [%v]\n", h.Validate())
		}
		if len(h) != 2 {
			t.Fatalf("Header Unexpected length error 1, length: [%v]\n", len(h))
		}
		h = h.Add("keyb", "valueb").Add("keya", "valuea2")
		if len(h) != 6 {
			t.Fatalf("Header Unexpected length error 2, length after add: [%v]\n", len(h))
		}
		if _, ok := h.Contains(k); !ok {
			t.Fatalf("Header Unexpected false for key: [%v]\n", k)
		}
		k = "xyz"
		if _, ok := h.Contains(k); ok {
			t.Fatalf("Header Unexpected true for key: [%v]\n", k)
		}
		//
		h = Headers{k}
		if e = h.Validate(); e != EHDRLEN {
			t.Fatalf("Header Validate, got [%v], expected [%v]\n", e, EHDRLEN)
		}
	}
}

/*
	Data Test: Headers UTF8.
*/
func TestHeadersUTF8(t *testing.T) {
	for _, _ = range Protocols() {
		k := "keya"
		v := "valuea"
		wh := Headers{k, v}
		var e error   // An error
		var rs string // Result string
		if rs, e = wh.ValidateUTF8(); e != nil {
			t.Fatalf("Unexpected UTF8 error 1: [%v]\n", e)
		}
		if rs != "" {
			t.Fatalf("Unexpected UTF8 error 1B, got [%v], expected [%v]\n", rs, "")
		}
		//
		wh = Headers{k, v, `“Iñtërnâtiônàlizætiøn”`, "valueb", "keyc", `“Iñtërnâtiônàlizætiøn”`}
		if _, e = wh.ValidateUTF8(); e != nil {
			t.Fatalf("Unexpected UTF8 error 2: [%v]\n", e)
		}
		//
		wh = Headers{k, v, `“Iñtërnâtiônàlizætiøn”`, "\x80", "keyc", `“Iñtërnâtiônàlizætiøn”`}
		if rs, e = wh.ValidateUTF8(); e == nil {
			t.Fatalf("Unexpected UTF8 error  3, got nil, expected an error")
		}
		if e != EHDRUTF8 {
			t.Fatalf("Unexpected UTF8 error 4, got [%v], expected [%v]\n", e, EHDRUTF8)
		}
		if rs != "\x80" {
			t.Fatalf("Unexpected UTF8 error 5, got [%v], expected [%v]\n", rs, "\x80")
		}
	}
}

/*.
Data Test: Headers Clone
*/
func TestHeadersClone(t *testing.T) {
	for _, _ = range Protocols() {
		wh := Headers{"ka", "va"}.Add("kb", "vb").Add("kc", "vc")
		hc := wh.Clone()
		if !wh.Compare(hc) {
			t.Fatalf("Unexpected false for clone: [%v], [%v]\n", wh, hc)
		}
	}
}

/*
	Data Test: Headers Add / Delete.
*/
func TestHeadersAddDelete(t *testing.T) {
	for _, _ = range Protocols() {
		ha := Headers{"ka", "va", "kb", "vb", "kc", "vc"}
		hb := Headers{"kaa", "va", "kbb", "vb", "kcc", "vc"}
		hn := ha.AddHeaders(hb)
		if len(ha)+len(hb) != len(hn) {
			t.Fatalf("Unexpected length AddHeaders, got: [%v], expected: [%v]\n", len(hn), len(ha)+len(hb))
		}
		ol := len(hn)
		hn = hn.Delete("ka")
		if len(hn) != ol-2 {
			t.Fatalf("Unexpected length Delete 1, got: [%v], expected: [%v]\n", len(hn), ol-2)
		}
		hn = hn.Delete("kcc")
		if len(hn) != ol-4 {
			t.Fatalf("Unexpected length Delete 2, got: [%v], expected: [%v]\n", len(hn), ol-4)
		}
	}
}

/*
	Data Test: Headers ContainsKV
*/
func TestHeadersContainsKV(t *testing.T) {
	for _, _ = range Protocols() {
		ha := Headers{"ka", "va", "kb", "vb", "kc", "vc"}
		b := ha.ContainsKV("kb", "vb")
		if !b {
			t.Fatalf("KV01 got false, expected true")
		}
		b = ha.ContainsKV("kb", "zz")
		if b {
			t.Fatalf("KV02 got true, expected false")
		}
	}
}

/*
	Data Test: Headers Compare
*/
func TestHeadersCompare(t *testing.T) {
	for _, _ = range Protocols() {
		ha := Headers{"ka", "va", "kb", "vb", "kc", "vc"}
		hb := Headers{"ka", "va", "kb", "vb", "kc", "vc"}
		hc := Headers{"ka", "va"}
		hd := Headers{"k1", "v1", "k2", "v2", "k3", "v3"}
		b := ha.Compare(hb)
		if !b {
			t.Fatalf("CMP01 Expected true, got false")
		}
		b = ha.Compare(hc)
		if b {
			t.Fatalf("CMP02 Expected false, got true")
		}
		b = ha.Compare(hd)
		if b {
			t.Fatalf("CMP03 Expected false, got true")
		}
		b = hd.Compare(ha)
		if b {
			t.Fatalf("CMP04 Expected false, got true")
		}
	}
}

/*
	Data Test: Headers Size
*/
func TestHeadersSize(t *testing.T) {
	for _, _ = range Protocols() {
		ha := Headers{"k", "v"}
		s := ha.Size(false)
		var w int64 = 4
		if s != w {
			t.Fatalf("SIZ01 size, got [%d], expected [%v]\n", s, w)
		}
		//
		ha = Headers{"kaa", "vaa2", "kba", "vba2", "kca", "vca2"}
		s = ha.Size(true)
		w = 3 + 1 + 4 + 1 + 3 + 1 + 4 + 1 + 3 + 1 + 4 + 1
		if s != w {
			t.Fatalf("SIZ02 size, got [%d] expected [%v]\n", s, w)
		}
	}
}

/*
	Data Test: Empty Header Key / Value
*/
func TestHeadersEmtKV(t *testing.T) {
	for _, _ = range Protocols() {
		wh := Headers{"a", "b", "c", "d"} // work headers
		ek := Headers{"a", "b", "", "d"}  // empty key
		ev := Headers{"a", "", "c", "d"}  // empty value
		//
		e = checkHeaders(wh, SPL_10)
		if e != nil {
			t.Fatalf("CHD01 Expected [nil], got [%v]\n", e)
		}
		e = checkHeaders(wh, SPL_11)
		if e != nil {
			t.Fatalf("CHD02 Expected [nil], got [%v]\n", e)
		}
		//
		e = checkHeaders(ek, SPL_10)
		if e != EHDRMTK {
			t.Fatalf("CHD03 Expected [%v], got [%v]\n", EHDRMTK, e)
		}
		e = checkHeaders(ek, SPL_11)
		if e != EHDRMTK {
			t.Fatalf("CHD04 Expected [%v], got [%v]\n", EHDRMTK, e)
		}
		//
		e = checkHeaders(ev, SPL_10)
		if e != EHDRMTV {
			t.Fatalf("CHD05 Expected [%v], got [%v]\n", EHDRMTV, e)
		}
		e = checkHeaders(ev, SPL_11)
		if e != nil {
			t.Fatalf("CHD06 Expected [nil], got [%v]\n", e)
		}
	}
}
<|endoftext|>


<commit_before>package datadog

type Check struct {
	Check     string   `json:"check"`
	HostName  string   `json:"host_name"`
	Status    status   `json:"status"`
	Timestamp string   `json:"timestamp,omitempty"`
	Message   string   `json:"message,omitempty"`
	Tags      []string `json:"tags,omitempty"`
}

type status int

const (
	OK status = iota
	WARNING
	CRITICAL
	UNKNOWN
)

// PostCheck posts the result of a check run to the server
func (client *Client) PostCheck(check Check) error {
	return client.doJsonRequest("POST", "/v1/check_run",
		check, nil)
}
<commit_msg>make checks.go/status public as its exposed in the API<commit_after>package datadog

type Check struct {
	Check     string   `json:"check"`
	HostName  string   `json:"host_name"`
	Status    Status   `json:"status"`
	Timestamp string   `json:"timestamp,omitempty"`
	Message   string   `json:"message,omitempty"`
	Tags      []string `json:"tags,omitempty"`
}

type Status int

const (
	OK Status = iota
	WARNING
	CRITICAL
	UNKNOWN
)

// PostCheck posts the result of a check run to the server
func (client *Client) PostCheck(check Check) error {
	return client.doJsonRequest("POST", "/v1/check_run",
		check, nil)
}
<|endoftext|>