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

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) receiver开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			/*if err != nil {
				Log.Infoc(func() string { return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err) })
				//         if err != io.EOF {
				//         }
				return
			}*/
			if err != nil {

				if err == io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
					})
					return
				} else {
					Log.Infoc(func() string {
						return fmt.Sprintf("(%s) info reading from connection: %v", this.cid(), err)
					})
				}

				continue
			} else {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s)向ringbuffer些数据成功！", this.cid())
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) sender开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting sender", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			_, err := this.out.WriteTo(conn)

			if err != nil {
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
}

// peekMessageSize() reads, but not commits, enough bytes to determine the size of
// the next message and returns the type and size.
func (this *service) peekMessageSize() (message.Message, int, error) {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize开始", this.cid())
	})
	var (
		b   *[]byte
		err error
		cnt int = 2
	)

	if this.in == nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in is nil", this.cid())
		})
		err = ErrBufferNotReady
		return nil, 0, err
	}

	// Let's read enough bytes to get the message header (msg type, remaining length)
	//for {
	// If we have read 5 bytes and still not done, then there's a problem.
	//if cnt > 5 {
	//	return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//}

	// Peek cnt bytes from the input buffer.
	b, err = this.in.ReadWait(cnt)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in.ReadWait falure", this.cid())
		})
		return nil, 0, err
	}
	//fmt.Println("b=", b)
	//// If not enough bytes are returned, then continue until there's enough.
	//if len(b) < cnt {
	//	continue
	//}

	// If we got enough bytes, then check the last byte to see if the continuation
	// bit is set. If so, increment cnt and continue peeking
	/*if b[cnt-1] >= 0x80 {
		cnt++
	} else {
		break
	}*/
	//for {
	//	if cnt > 5 {
	//		return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//	}
	//
	//	if b[cnt-1] >= 0x80 {
	//		cnt++
	//	} else {
	//		break
	//	}
	//}
	//}

	// Get the remaining length of the message
	//remlen, m := binary.Uvarint(b[1:cnt])

	// Total message length is remlen + 1 (msg type) + m (remlen bytes)
	//total := int(remlen) + 1 + m

	mtype := message.MessageType(*b[0] >> 4)

	//return mtype, total, err
	var msg message.Message
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize mtype.New() falure", this.cid())
		})
		return nil, 0, err
	}
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) 开始创建对象(%s)", this.cid(), msg.Name())
	})
	_, err = msg.Decode(*b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize msg.Decode falure", this.cid())
		})
		return nil, 0, err
	}
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize结束(%s)", this.cid(), msg.Name())
	})
	return msg, len(*b), err
}

// peekMessage() reads a message from the buffer, but the bytes are NOT committed.
// This means the buffer still thinks the bytes are not read yet.
/*func (this *service) peekMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b    []byte
		err  error
		i, n int
		msg  message.Message
	)

	if this.in == nil {
		return nil, 0, ErrBufferNotReady
	}

	// Peek until we get total bytes
	for i = 0; ; i++ {
		// Peek remlen bytes from the input buffer.
		b, err = this.in.ReadWait(total)
		if err != nil && err != ErrBufferInsufficientData {
			return nil, 0, err
		}

		// If not enough bytes are returned, then continue until there's enough.
		if len(b) >= total {
			break
		}
	}

	msg, err = mtype.New()
	if err != nil {
		return nil, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}*/

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b   []byte
		err error
		n   int
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, 0, err
	}

	if len(this.intmp) < total {
		this.intmp = make([]byte, total)
	}

	// Read until we get total bytes
	l := 0
	for l < total {
		n, err = this.in.Read(this.intmp[l:])
		l += n
		Log.Debugc(func() string {
			return fmt.Sprintf("read %d bytes, total %d", n, l)
		})
		if err != nil {
			return nil, 0, err
		}
	}

	b = this.intmp[:total]

	msg, err = mtype.New()
	if err != nil {
		return msg, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (int, error) {
	var (
		l    int = msg.Len()
		m, n int
		err  error
		buf  *[]byte
		wrap bool
	)

	if this.out == nil {
		return 0, ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	this.wmu.Lock()
	defer this.wmu.Unlock()

	buf, wrap, err = this.out.WriteWait(l)
	if err != nil {
		return 0, err
	}

	if wrap {
		if len(this.outtmp) < l {
			this.outtmp = make([]byte, l)
		}

		n, err = msg.Encode(this.outtmp[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.Write(this.outtmp[0:n])
		if err != nil {
			return m, err
		}
	} else {
		n, err = msg.Encode(*buf[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.WriteCommit(n)
		if err != nil {
			return 0, err
		}
	}

	this.outStat.increment(int64(m))

	return m, nil
}
<commit_msg>修改了buffer.go;sendrecv.go;process.go<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) receiver开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			/*if err != nil {
				Log.Infoc(func() string { return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err) })
				//         if err != io.EOF {
				//         }
				return
			}*/
			if err != nil {

				if err == io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
					})
					return
				} else {
					Log.Infoc(func() string {
						return fmt.Sprintf("(%s) info reading from connection: %v", this.cid(), err)
					})
				}

				continue
			} else {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s)向ringbuffer些数据成功！", this.cid())
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) sender开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting sender", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			_, err := this.out.WriteTo(conn)

			if err != nil {
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
}

// peekMessageSize() reads, but not commits, enough bytes to determine the size of
// the next message and returns the type and size.
func (this *service) peekMessageSize() (message.Message, int, error) {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize开始", this.cid())
	})
	var (
		b   *[]byte
		err error
		cnt int = 2
	)

	if this.in == nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in is nil", this.cid())
		})
		err = ErrBufferNotReady
		return nil, 0, err
	}

	// Let's read enough bytes to get the message header (msg type, remaining length)
	//for {
	// If we have read 5 bytes and still not done, then there's a problem.
	//if cnt > 5 {
	//	return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//}

	// Peek cnt bytes from the input buffer.
	b, err = this.in.ReadWait(cnt)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in.ReadWait falure", this.cid())
		})
		return nil, 0, err
	}
	//fmt.Println("b=", b)
	//// If not enough bytes are returned, then continue until there's enough.
	//if len(b) < cnt {
	//	continue
	//}

	// If we got enough bytes, then check the last byte to see if the continuation
	// bit is set. If so, increment cnt and continue peeking
	/*if b[cnt-1] >= 0x80 {
		cnt++
	} else {
		break
	}*/
	//for {
	//	if cnt > 5 {
	//		return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//	}
	//
	//	if b[cnt-1] >= 0x80 {
	//		cnt++
	//	} else {
	//		break
	//	}
	//}
	//}

	// Get the remaining length of the message
	//remlen, m := binary.Uvarint(b[1:cnt])

	// Total message length is remlen + 1 (msg type) + m (remlen bytes)
	//total := int(remlen) + 1 + m

	mtype := message.MessageType((*b)[0] >> 4)

	//return mtype, total, err
	var msg message.Message
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize mtype.New() falure", this.cid())
		})
		return nil, 0, err
	}
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) 开始创建对象(%s)", this.cid(), msg.Name())
	})
	_, err = msg.Decode(*b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize msg.Decode falure", this.cid())
		})
		return nil, 0, err
	}
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize结束(%s)", this.cid(), msg.Name())
	})
	return msg, len(*b), err
}

// peekMessage() reads a message from the buffer, but the bytes are NOT committed.
// This means the buffer still thinks the bytes are not read yet.
/*func (this *service) peekMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b    []byte
		err  error
		i, n int
		msg  message.Message
	)

	if this.in == nil {
		return nil, 0, ErrBufferNotReady
	}

	// Peek until we get total bytes
	for i = 0; ; i++ {
		// Peek remlen bytes from the input buffer.
		b, err = this.in.ReadWait(total)
		if err != nil && err != ErrBufferInsufficientData {
			return nil, 0, err
		}

		// If not enough bytes are returned, then continue until there's enough.
		if len(b) >= total {
			break
		}
	}

	msg, err = mtype.New()
	if err != nil {
		return nil, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}*/

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b   []byte
		err error
		n   int
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, 0, err
	}

	if len(this.intmp) < total {
		this.intmp = make([]byte, total)
	}

	// Read until we get total bytes
	l := 0
	for l < total {
		n, err = this.in.Read(this.intmp[l:])
		l += n
		Log.Debugc(func() string {
			return fmt.Sprintf("read %d bytes, total %d", n, l)
		})
		if err != nil {
			return nil, 0, err
		}
	}

	b = this.intmp[:total]

	msg, err = mtype.New()
	if err != nil {
		return msg, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (int, error) {
	var (
		l    int = msg.Len()
		m, n int
		err  error
		buf  *[]byte
		wrap bool
	)

	if this.out == nil {
		return 0, ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	this.wmu.Lock()
	defer this.wmu.Unlock()

	buf, wrap, err = this.out.WriteWait(l)
	if err != nil {
		return 0, err
	}

	if wrap {
		if len(this.outtmp) < l {
			this.outtmp = make([]byte, l)
		}

		n, err = msg.Encode(this.outtmp[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.Write(this.outtmp[0:n])
		if err != nil {
			return m, err
		}
	} else {
		n, err = msg.Encode((*buf)[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.WriteCommit(n)
		if err != nil {
			return 0, err
		}
	}

	this.outStat.increment(int64(m))

	return m, nil
}
<|endoftext|>


<commit_before>package realm

import "github.com/miekg/dns"

// RecordsEntry is used to hold a mapping of DNS request types to DNS records
type RecordsEntry map[uint16][]dns.RR

// GetRecords will fetch the appropriate DNS records to the requested type
func (entry RecordsEntry) GetRecords(rrType uint16) []dns.RR {
	var records []dns.RR
	records = make([]dns.RR, 0)

	if rrType == dns.TypeANY {
		for _, rrs := range entry {
			records = append(records, rrs...)
		}
	} else if rrs, ok := entry[rrType]; ok {
		records = append(records, rrs...)
	}

	return records
}

// DomainEntry is used to hold a mapping of DNS request classes to RecordEntrys
type DomainEntry map[uint16]RecordsEntry

// AddEntry is used to add a new DNS record to this mapping
func (entry DomainEntry) AddEntry(record dns.RR) {
	var header *dns.RR_Header
	header = record.Header()

	if _, ok := entry[header.Class]; !ok {
		entry[header.Class] = make(RecordsEntry)
	}
	if _, ok := entry[header.Class][header.Rrtype]; !ok {
		entry[header.Class][header.Rrtype] = make([]dns.RR, 0)
	}

	entry[header.Class][header.Rrtype] = append(entry[header.Class][header.Rrtype], record)
}

// GetEntries is used to find the appropriate RecordEntrys for the requested DNS class
func (entry DomainEntry) GetEntries(rrClass uint16) []RecordsEntry {
	var entries []RecordsEntry
	entries = make([]RecordsEntry, 0)

	if rrClass == dns.ClassANY {
		for _, entry := range entry {
			entries = append(entries, entry)
		}
	} else if entry, ok := entry[rrClass]; ok {
		entries = append(entries, entry)
	}

	return entries
}

// Registry is a container for looking up DNS records for any request
type Registry map[string]DomainEntry

// NewRegistry will allocate and return a new *Registry
func NewRegistry() *Registry {
	return &Registry{
		records: make(map[string]DomainEntry),
	}
}

// addRecord is used to add a new DNS record to this registry
func (r *Registry) addRecord(record dns.RR) {
	var header *dns.RR_Header
	header = record.Header()

	var name string
	name = dns.Fqdn(header.Name)

	if _, ok := r.records[name]; !ok {
		r.records[name] = make(DomainEntry)
	}
	r.records[name].AddEntry(record)

	// If this record is an SOA record then also store under the Mbox name
	if header.Rrtype == dns.TypeSOA {
		var soa *dns.SOA
		soa = record.(*dns.SOA)

		if _, ok := r.records[soa.Mbox]; !ok {
			r.records[soa.Mbox] = make(DomainEntry)
		}
		r.records[soa.Mbox].AddEntry(record)
	}
}

// AddZone is used to add the records from a *Zone into this *Registry
func (r *Registry) AddZone(z *Zone) {
	for _, record := range z.Records() {
		r.addRecord(record)
	}
}

// Lookup will find all records which we should respond with for the given name, request type, and request class.
func (r *Registry) Lookup(name string, reqType uint16, reqClass uint16) []dns.RR {
	name = dns.Fqdn(name)
	var records []dns.RR
	records = make([]dns.RR, 0)

	var domainEntry DomainEntry
	var ok bool
	domainEntry, ok = r.records[name]
	if !ok {
		return records
	}

	var recordEntries []RecordsEntry
	recordEntries = domainEntry.GetEntries(reqClass)

	for _, recordEntry := range recordEntries {
		var rrs []dns.RR
		rrs = recordEntry.GetRecords(reqType)
		records = append(records, rrs...)

		if len(rrs) == 0 && reqType == dns.TypeA {
			rrs = recordEntry.GetRecords(dns.TypeCNAME)
			for _, rr := range rrs {
				records = append(records, rr)
				var header *dns.RR_Header
				header = rr.Header()
				if header.Rrtype == dns.TypeCNAME && reqType != dns.TypeCNAME {
					// Attempt to resolve this CNAME record
					var cname *dns.CNAME
					cname = rr.(*dns.CNAME)
					var cnameRecords []dns.RR
					cnameRecords = r.Lookup(dns.Fqdn(cname.Target), reqType, reqClass)
					records = append(records, cnameRecords...)
				}

			}
		}

	}
	return records
}
<commit_msg>revert back to a struct<commit_after>package realm

import "github.com/miekg/dns"

// RecordsEntry is used to hold a mapping of DNS request types to DNS records
type RecordsEntry map[uint16][]dns.RR

// GetRecords will fetch the appropriate DNS records to the requested type
func (entry RecordsEntry) GetRecords(rrType uint16) []dns.RR {
	var records []dns.RR
	records = make([]dns.RR, 0)

	if rrType == dns.TypeANY {
		for _, rrs := range entry {
			records = append(records, rrs...)
		}
	} else if rrs, ok := entry[rrType]; ok {
		records = append(records, rrs...)
	}

	return records
}

// DomainEntry is used to hold a mapping of DNS request classes to RecordEntrys
type DomainEntry map[uint16]RecordsEntry

// AddEntry is used to add a new DNS record to this mapping
func (entry DomainEntry) AddEntry(record dns.RR) {
	var header *dns.RR_Header
	header = record.Header()

	if _, ok := entry[header.Class]; !ok {
		entry[header.Class] = make(RecordsEntry)
	}
	if _, ok := entry[header.Class][header.Rrtype]; !ok {
		entry[header.Class][header.Rrtype] = make([]dns.RR, 0)
	}

	entry[header.Class][header.Rrtype] = append(entry[header.Class][header.Rrtype], record)
}

// GetEntries is used to find the appropriate RecordEntrys for the requested DNS class
func (entry DomainEntry) GetEntries(rrClass uint16) []RecordsEntry {
	var entries []RecordsEntry
	entries = make([]RecordsEntry, 0)

	if rrClass == dns.ClassANY {
		for _, entry := range entry {
			entries = append(entries, entry)
		}
	} else if entry, ok := entry[rrClass]; ok {
		entries = append(entries, entry)
	}

	return entries
}

// Registry is a container for looking up DNS records for any request
type Registry struct {
	records map[string]DomainEntry
}

// NewRegistry will allocate and return a new *Registry
func NewRegistry() *Registry {
	return &Registry{
		records: make(map[string]DomainEntry),
	}
}

// addRecord is used to add a new DNS record to this registry
func (r *Registry) addRecord(record dns.RR) {
	var header *dns.RR_Header
	header = record.Header()

	var name string
	name = dns.Fqdn(header.Name)

	if _, ok := r.records[name]; !ok {
		r.records[name] = make(DomainEntry)
	}
	r.records[name].AddEntry(record)

	// If this record is an SOA record then also store under the Mbox name
	if header.Rrtype == dns.TypeSOA {
		var soa *dns.SOA
		soa = record.(*dns.SOA)

		if _, ok := r.records[soa.Mbox]; !ok {
			r.records[soa.Mbox] = make(DomainEntry)
		}
		r.records[soa.Mbox].AddEntry(record)
	}
}

// AddZone is used to add the records from a *Zone into this *Registry
func (r *Registry) AddZone(z *Zone) {
	for _, record := range z.Records() {
		r.addRecord(record)
	}
}

// Lookup will find all records which we should respond with for the given name, request type, and request class.
func (r *Registry) Lookup(name string, reqType uint16, reqClass uint16) []dns.RR {
	name = dns.Fqdn(name)
	var records []dns.RR
	records = make([]dns.RR, 0)

	var domainEntry DomainEntry
	var ok bool
	domainEntry, ok = r.records[name]
	if !ok {
		return records
	}

	var recordEntries []RecordsEntry
	recordEntries = domainEntry.GetEntries(reqClass)

	for _, recordEntry := range recordEntries {
		var rrs []dns.RR
		rrs = recordEntry.GetRecords(reqType)
		records = append(records, rrs...)

		if len(rrs) == 0 && reqType == dns.TypeA {
			rrs = recordEntry.GetRecords(dns.TypeCNAME)
			for _, rr := range rrs {
				records = append(records, rr)
				var header *dns.RR_Header
				header = rr.Header()
				if header.Rrtype == dns.TypeCNAME && reqType != dns.TypeCNAME {
					// Attempt to resolve this CNAME record
					var cname *dns.CNAME
					cname = rr.(*dns.CNAME)
					var cnameRecords []dns.RR
					cnameRecords = r.Lookup(dns.Fqdn(cname.Target), reqType, reqClass)
					records = append(records, cnameRecords...)
				}

			}
		}

	}
	return records
}
<|endoftext|>


<commit_before>package speed

import (
	"errors"
	"fmt"
	"regexp"
	"sync"

	"github.com/Sirupsen/logrus"
)

// Registry defines a valid set of instance domains and metrics
type Registry interface {
	// checks if an instance domain of the passed name is already present or not
	HasInstanceDomain(name string) bool

	// checks if an metric of the passed name is already present or not
	HasMetric(name string) bool

	// returns the number of Metrics in the current registry
	MetricCount() int

	// returns the number of Values in the current registry
	ValuesCount() int

	// returns the number of Instance Domains in the current registry
	InstanceDomainCount() int

	// returns the number of instances across all instance domains in the current registry
	InstanceCount() int

	// returns the number of non null strings initialized in the current registry
	StringCount() int

	// adds a InstanceDomain object to the writer
	AddInstanceDomain(InstanceDomain) error

	// adds a InstanceDomain object after constructing it using passed name and instances
	AddInstanceDomainByName(name string, instances []string) (InstanceDomain, error)

	// adds a Metric object to the writer
	AddMetric(Metric) error

	// adds a Metric object after parsing the passed string for Instances and InstanceDomains
	AddMetricByString(name string, val interface{}, s MetricSemantics, t MetricType, u MetricUnit) (Metric, error)
}

// PCPRegistry implements a registry for PCP as the client
type PCPRegistry struct {
	instanceDomains map[string]*PCPInstanceDomain // a cache for instanceDomains
	metrics         map[string]PCPMetric          // a cache for metrics

	// locks
	indomlock   sync.RWMutex
	metricslock sync.RWMutex

	// offsets
	instanceoffset int
	indomoffset    int
	metricsoffset  int
	valuesoffset   int
	stringsoffset  int

	// counts
	instanceCount int
	valueCount    int
	stringcount   int

	mapped bool
}

// NewPCPRegistry creates a new PCPRegistry object
func NewPCPRegistry() *PCPRegistry {
	return &PCPRegistry{
		instanceDomains: make(map[string]*PCPInstanceDomain),
		metrics:         make(map[string]PCPMetric),
	}
}

// InstanceCount returns the number of instances across all indoms in the registry
func (r *PCPRegistry) InstanceCount() int {
	r.indomlock.RLock()
	defer r.indomlock.RUnlock()

	return r.instanceCount
}

// InstanceDomainCount returns the number of instance domains in the registry
func (r *PCPRegistry) InstanceDomainCount() int {
	r.indomlock.RLock()
	defer r.indomlock.RUnlock()

	return len(r.instanceDomains)
}

// MetricCount returns the number of metrics in the registry
func (r *PCPRegistry) MetricCount() int {
	r.metricslock.RLock()
	defer r.metricslock.RUnlock()

	return len(r.metrics)
}

// ValuesCount returns the number of values in the registry
func (r *PCPRegistry) ValuesCount() int { return r.valueCount }

// StringCount returns the number of strings in the registry
func (r *PCPRegistry) StringCount() int { return r.stringcount }

// HasInstanceDomain returns true if the registry already has an indom of the specified name
func (r *PCPRegistry) HasInstanceDomain(name string) bool {
	r.indomlock.RLock()
	defer r.indomlock.RUnlock()

	_, present := r.instanceDomains[name]
	return present
}

// HasMetric returns true if the registry already has a metric of the specified name
func (r *PCPRegistry) HasMetric(name string) bool {
	r.metricslock.RLock()
	defer r.metricslock.RUnlock()

	_, present := r.metrics[name]
	return present
}

// AddInstanceDomain will add a new instance domain to the current registry
func (r *PCPRegistry) AddInstanceDomain(indom InstanceDomain) error {
	if r.HasInstanceDomain(indom.Name()) {
		return errors.New("InstanceDomain is already defined for the current registry")
	}

	r.indomlock.Lock()
	defer r.indomlock.Unlock()

	if r.mapped {
		return errors.New("Cannot add an indom when a mapping is active")
	}

	r.instanceDomains[indom.Name()] = indom.(*PCPInstanceDomain)
	r.instanceCount += indom.InstanceCount()

	log.WithFields(logrus.Fields{
		"prefix":        "registry",
		"name":          indom.Name(),
		"instanceCount": indom.InstanceCount(),
	}).Info("added new instance domain")

	if indom.(*PCPInstanceDomain).shortDescription.val != "" {
		r.stringcount++
	}

	if indom.(*PCPInstanceDomain).longDescription.val != "" {
		r.stringcount++
	}

	return nil
}

// AddMetric will add a new metric to the current registry
func (r *PCPRegistry) AddMetric(m Metric) error {
	if r.HasMetric(m.Name()) {
		return errors.New("Metric is already defined for the current registry")
	}

	pcpm := m.(PCPMetric)

	// if it is an indom metric
	if pcpm.Indom() != nil && !r.HasInstanceDomain(pcpm.Indom().Name()) {
		return errors.New("Instance Domain is not defined for current registry")
	}

	if r.mapped {
		return errors.New("Cannot add a metric when a mapping is active")
	}

	r.metricslock.Lock()
	defer r.metricslock.Unlock()

	log.WithFields(logrus.Fields{
		"prefix":    "registry",
		"name":      m.Name(),
		"type":      m.Type(),
		"unit":      m.Unit(),
		"semantics": m.Semantics(),
	}).Info("added new metric")

	r.metrics[m.Name()] = pcpm

	currentValues := 1
	if pcpm.Indom() != nil {
		currentValues = pcpm.Indom().InstanceCount()
	}

	r.valueCount += currentValues
	if pcpm.Type() == StringType {
		r.stringcount += currentValues
	}

	if pcpm.ShortDescription() != "" {
		r.stringcount++
	}

	if pcpm.LongDescription() != "" {
		r.stringcount++
	}

	return nil
}

// AddInstanceDomainByName adds an instance domain using passed parameters
func (r *PCPRegistry) AddInstanceDomainByName(name string, instances []string) (InstanceDomain, error) {
	if r.HasInstanceDomain(name) {
		return nil, errors.New("The InstanceDomain already exists for this registry")
	}

	indom, err := NewPCPInstanceDomain(name, instances)
	if err != nil {
		return nil, err
	}

	err = r.AddInstanceDomain(indom)
	if err != nil {
		return nil, err
	}

	return indom, nil
}

const id = "[\\p{L}\\p{N}]+"

var instancesPattern = fmt.Sprintf("(%v)((,\\s?(%v))*)", id, id)
var pattern = fmt.Sprintf("\\A((%v)(\\.%v)*?)(\\[(%v)\\])?((\\.%v)*)\\z", id, id, instancesPattern, id)

var ireg, _ = regexp.Compile(id)
var reg, _ = regexp.Compile(pattern)

func parseString(s string) (metric string, indom string, instances []string, err error) {
	if !reg.MatchString(s) {
		return "", "", nil, errors.New("Invalid String")
	}

	matches := reg.FindStringSubmatch(s)
	n := len(matches)

	indom = matches[1]
	metric = indom + matches[n-2]

	iarr := matches[5]
	if iarr != "" {
		instances = ireg.FindAllString(iarr, -1)
	} else {
		instances = nil
		indom = ""
	}

	return
}

// AddMetricByString dynamically creates a PCPMetric
func (r *PCPRegistry) AddMetricByString(str string, val interface{}, s MetricSemantics, t MetricType, u MetricUnit) (Metric, error) {
	metric, indom, instances, err := parseString(str)
	if err != nil {
		return nil, err
	}

	var m Metric

	if instances == nil {
		// singleton metric
		m, err = NewPCPSingletonMetric(val, metric, t, s, u)
		if err != nil {
			return nil, err
		}

		err = r.AddMetric(m)
		if err != nil {
			return nil, err
		}

		return m, nil
	}

	// instance metric
	mp, ok := val.(Instances)
	if !ok {
		return nil, errors.New("to define an instance metric, a Instances type is required")
	}

	id, err := r.AddInstanceDomainByName(indom, instances)
	if err != nil {
		return nil, err
	}

	m, err = NewPCPInstanceMetric(mp, metric, id.(*PCPInstanceDomain), t, s, u)
	if err != nil {
		return nil, err
	}

	err = r.AddMetric(m)
	if err != nil {
		return nil, err
	}

	return m, nil
}
<commit_msg>registry: AddMetricByString: use existing indom if present<commit_after>package speed

import (
	"errors"
	"fmt"
	"regexp"
	"sync"

	"github.com/Sirupsen/logrus"
)

// Registry defines a valid set of instance domains and metrics
type Registry interface {
	// checks if an instance domain of the passed name is already present or not
	HasInstanceDomain(name string) bool

	// checks if an metric of the passed name is already present or not
	HasMetric(name string) bool

	// returns the number of Metrics in the current registry
	MetricCount() int

	// returns the number of Values in the current registry
	ValuesCount() int

	// returns the number of Instance Domains in the current registry
	InstanceDomainCount() int

	// returns the number of instances across all instance domains in the current registry
	InstanceCount() int

	// returns the number of non null strings initialized in the current registry
	StringCount() int

	// adds a InstanceDomain object to the writer
	AddInstanceDomain(InstanceDomain) error

	// adds a InstanceDomain object after constructing it using passed name and instances
	AddInstanceDomainByName(name string, instances []string) (InstanceDomain, error)

	// adds a Metric object to the writer
	AddMetric(Metric) error

	// adds a Metric object after parsing the passed string for Instances and InstanceDomains
	AddMetricByString(name string, val interface{}, s MetricSemantics, t MetricType, u MetricUnit) (Metric, error)
}

// PCPRegistry implements a registry for PCP as the client
type PCPRegistry struct {
	instanceDomains map[string]*PCPInstanceDomain // a cache for instanceDomains
	metrics         map[string]PCPMetric          // a cache for metrics

	// locks
	indomlock   sync.RWMutex
	metricslock sync.RWMutex

	// offsets
	instanceoffset int
	indomoffset    int
	metricsoffset  int
	valuesoffset   int
	stringsoffset  int

	// counts
	instanceCount int
	valueCount    int
	stringcount   int

	mapped bool
}

// NewPCPRegistry creates a new PCPRegistry object
func NewPCPRegistry() *PCPRegistry {
	return &PCPRegistry{
		instanceDomains: make(map[string]*PCPInstanceDomain),
		metrics:         make(map[string]PCPMetric),
	}
}

// InstanceCount returns the number of instances across all indoms in the registry
func (r *PCPRegistry) InstanceCount() int {
	r.indomlock.RLock()
	defer r.indomlock.RUnlock()

	return r.instanceCount
}

// InstanceDomainCount returns the number of instance domains in the registry
func (r *PCPRegistry) InstanceDomainCount() int {
	r.indomlock.RLock()
	defer r.indomlock.RUnlock()

	return len(r.instanceDomains)
}

// MetricCount returns the number of metrics in the registry
func (r *PCPRegistry) MetricCount() int {
	r.metricslock.RLock()
	defer r.metricslock.RUnlock()

	return len(r.metrics)
}

// ValuesCount returns the number of values in the registry
func (r *PCPRegistry) ValuesCount() int { return r.valueCount }

// StringCount returns the number of strings in the registry
func (r *PCPRegistry) StringCount() int { return r.stringcount }

// HasInstanceDomain returns true if the registry already has an indom of the specified name
func (r *PCPRegistry) HasInstanceDomain(name string) bool {
	r.indomlock.RLock()
	defer r.indomlock.RUnlock()

	_, present := r.instanceDomains[name]
	return present
}

// HasMetric returns true if the registry already has a metric of the specified name
func (r *PCPRegistry) HasMetric(name string) bool {
	r.metricslock.RLock()
	defer r.metricslock.RUnlock()

	_, present := r.metrics[name]
	return present
}

// AddInstanceDomain will add a new instance domain to the current registry
func (r *PCPRegistry) AddInstanceDomain(indom InstanceDomain) error {
	if r.HasInstanceDomain(indom.Name()) {
		return errors.New("InstanceDomain is already defined for the current registry")
	}

	r.indomlock.Lock()
	defer r.indomlock.Unlock()

	if r.mapped {
		return errors.New("Cannot add an indom when a mapping is active")
	}

	r.instanceDomains[indom.Name()] = indom.(*PCPInstanceDomain)
	r.instanceCount += indom.InstanceCount()

	log.WithFields(logrus.Fields{
		"prefix":        "registry",
		"name":          indom.Name(),
		"instanceCount": indom.InstanceCount(),
	}).Info("added new instance domain")

	if indom.(*PCPInstanceDomain).shortDescription.val != "" {
		r.stringcount++
	}

	if indom.(*PCPInstanceDomain).longDescription.val != "" {
		r.stringcount++
	}

	return nil
}

// AddMetric will add a new metric to the current registry
func (r *PCPRegistry) AddMetric(m Metric) error {
	if r.HasMetric(m.Name()) {
		return errors.New("Metric is already defined for the current registry")
	}

	pcpm := m.(PCPMetric)

	// if it is an indom metric
	if pcpm.Indom() != nil && !r.HasInstanceDomain(pcpm.Indom().Name()) {
		return errors.New("Instance Domain is not defined for current registry")
	}

	if r.mapped {
		return errors.New("Cannot add a metric when a mapping is active")
	}

	r.metricslock.Lock()
	defer r.metricslock.Unlock()

	log.WithFields(logrus.Fields{
		"prefix":    "registry",
		"name":      m.Name(),
		"type":      m.Type(),
		"unit":      m.Unit(),
		"semantics": m.Semantics(),
	}).Info("added new metric")

	r.metrics[m.Name()] = pcpm

	currentValues := 1
	if pcpm.Indom() != nil {
		currentValues = pcpm.Indom().InstanceCount()
	}

	r.valueCount += currentValues
	if pcpm.Type() == StringType {
		r.stringcount += currentValues
	}

	if pcpm.ShortDescription() != "" {
		r.stringcount++
	}

	if pcpm.LongDescription() != "" {
		r.stringcount++
	}

	return nil
}

// AddInstanceDomainByName adds an instance domain using passed parameters
func (r *PCPRegistry) AddInstanceDomainByName(name string, instances []string) (InstanceDomain, error) {
	if r.HasInstanceDomain(name) {
		return nil, errors.New("The InstanceDomain already exists for this registry")
	}

	indom, err := NewPCPInstanceDomain(name, instances)
	if err != nil {
		return nil, err
	}

	err = r.AddInstanceDomain(indom)
	if err != nil {
		return nil, err
	}

	return indom, nil
}

const id = "[\\p{L}\\p{N}]+"

var instancesPattern = fmt.Sprintf("(%v)((,\\s?(%v))*)", id, id)
var pattern = fmt.Sprintf("\\A((%v)(\\.%v)*?)(\\[(%v)\\])?((\\.%v)*)\\z", id, id, instancesPattern, id)

var ireg, _ = regexp.Compile(id)
var reg, _ = regexp.Compile(pattern)

func parseString(s string) (metric string, indom string, instances []string, err error) {
	if !reg.MatchString(s) {
		return "", "", nil, errors.New("Invalid String")
	}

	matches := reg.FindStringSubmatch(s)
	n := len(matches)

	indom = matches[1]
	metric = indom + matches[n-2]

	iarr := matches[5]
	if iarr != "" {
		instances = ireg.FindAllString(iarr, -1)
	} else {
		instances = nil
		indom = ""
	}

	return
}

// AddMetricByString dynamically creates a PCPMetric
func (r *PCPRegistry) AddMetricByString(str string, val interface{}, s MetricSemantics, t MetricType, u MetricUnit) (Metric, error) {
	metric, indom, instances, err := parseString(str)
	if err != nil {
		return nil, err
	}

	var m Metric

	if instances == nil {
		// singleton metric
		m, err = NewPCPSingletonMetric(val, metric, t, s, u)
		if err != nil {
			return nil, err
		}

		err = r.AddMetric(m)
		if err != nil {
			return nil, err
		}

		return m, nil
	}

	// instance metric
	mp, ok := val.(Instances)
	if !ok {
		return nil, errors.New("to define an instance metric, a Instances type is required")
	}

	var id InstanceDomain

	if !r.HasInstanceDomain(indom) {
		id, err = r.AddInstanceDomainByName(indom, instances)
		if err != nil {
			return nil, err
		}
	} else {
		id = r.instanceDomains[indom]
	}

	m, err = NewPCPInstanceMetric(mp, metric, id.(*PCPInstanceDomain), t, s, u)
	if err != nil {
		return nil, err
	}

	err = r.AddMetric(m)
	if err != nil {
		return nil, err
	}

	return m, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/rbkreisberg/CloudForest"
	"log"
	"os"
)

func main() {
	fm := flag.String("fm", "featurematrix.afm", "AFM formated feature matrix to use.")
	rf := flag.String("rfpred", "rface.sf", "A predictor forest as outputed by rf-ace")
	outf := flag.String("splits", "splits.tsv", "a case by case sparse matrix of leaf cooccurance in tsv format")
	boutf := flag.String("branches", "branches.tsv", "a case by feature sparse matrix of case/splitter cooccurance in tsv format")
	rboutf := flag.String("relbranches", "relativeBranches.tsv", "a case by feature sparse matrix of split direction for each case/feature in tsv format")
	splitdistf := flag.String("splitlist", "splitList.tsv", "a list of values for each feature that was split on")

	flag.Parse()

	datafile, err := os.Open(*fm) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer datafile.Close()
	data := CloudForest.ParseAFM(datafile)
	nfeatures := len(data.Data)
	ncases := len(data.Data[0].Data)
	log.Print("Data file ", nfeatures, " by ", ncases)

	forestfile, err := os.Open(*rf) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer forestfile.Close()
	forest := CloudForest.ParseRfAcePredictor(forestfile)
	log.Print("Forest has ", len(forest.Trees), " trees ")

	featureCounts := make([]int, nfeatures) // a total count of the number of times each feature was used to split
	caseFeatureCounts := new(CloudForest.SparseCounter)
	relativeSplitCount := new(CloudForest.SparseCounter)

	splitValueList = make(map[int][]float64, nfeatures)

	for i := 0; i < len(forest.Trees); i++ {
		splits := forest.Trees[i].GetSplits(data, caseFeatureCounts, relativeSplitCount)

		for _, split := range splits {
			featureId := data.Map[split.Feature]
			featureCounts[featureId]++ //increment the count for the total # of times the feature was a splitter

			if split.Numerical == true {
				if splitValueList[featureId] == nil {
					splitValueList[featureId] = make([]float64, 0)
				}
				splitValueList[featureId] = append(splitValueList[featureId], split.Value)
			}
		}
	}

	log.Print("Outputing Split Feature/Case Cooccurance Counts")
	outfile, err := os.Create(*outf) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer outfile.Close()
	for feature, count := range featureCounts {
		if _, err := fmt.Fprintf(outfile, "%v\t%v\n", feature, count); err != nil {
			log.Fatal(err)
		}
	}

	log.Print("Outputing Split Distribution")
	splitdistfile, err := os.Create(*splitdistf) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer splitdistfile.Close()
	for feature, list := range splitValueList {
		if _, err := fmt.Fprintf(splitdistfile, "%v", feature); err != nil {
			log.Fatal(err)
		}
		for _, value := range list {
			if _, err := fmt.Fprintf(splitdistfile, "\t%v", value); err != nil {
				log.Fatal(err)
			}
		}
		fmt.Fprintf(splitdistfile, "\n")

	}

	log.Print("Outputing Case Feature Cooccurance Counts")
	boutfile, err := os.Create(*boutf) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer boutfile.Close()
	caseFeatureCounts.WriteTsv(boutfile)

	log.Print("Outputing Case Feature Splitter Direction")
	rboutfile, err := os.Create(*rboutf) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer rboutfile.Close()
	relativeSplitCount.WriteTsv(rboutfile)
}
<commit_msg>fix type handling<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/rbkreisberg/CloudForest"
	"log"
	"os"
)

func main() {
	fm := flag.String("fm", "featurematrix.afm", "AFM formated feature matrix to use.")
	rf := flag.String("rfpred", "rface.sf", "A predictor forest as outputed by rf-ace")
	outf := flag.String("splits", "splits.tsv", "a case by case sparse matrix of leaf cooccurance in tsv format")
	boutf := flag.String("branches", "branches.tsv", "a case by feature sparse matrix of case/splitter cooccurance in tsv format")
	rboutf := flag.String("relbranches", "relativeBranches.tsv", "a case by feature sparse matrix of split direction for each case/feature in tsv format")
	splitdistf := flag.String("splitlist", "splitList.tsv", "a list of values for each feature that was split on")

	flag.Parse()

	datafile, err := os.Open(*fm) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer datafile.Close()
	data := CloudForest.ParseAFM(datafile)
	nfeatures := len(data.Data)
	ncases := len(data.Data[0].Data)
	log.Print("Data file ", nfeatures, " by ", ncases)

	forestfile, err := os.Open(*rf) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer forestfile.Close()
	forest := CloudForest.ParseRfAcePredictor(forestfile)
	log.Print("Forest has ", len(forest.Trees), " trees ")

	featureCounts := make([]int, nfeatures) // a total count of the number of times each feature was used to split
	caseFeatureCounts := new(CloudForest.SparseCounter)
	relativeSplitCount := new(CloudForest.SparseCounter)

	splitValueList := make(map[int][]float64, nfeatures)

	for i := 0; i < len(forest.Trees); i++ {
		splits := forest.Trees[i].GetSplits(data, caseFeatureCounts, relativeSplitCount)

		for _, split := range splits {
			featureId := data.Map[split.Feature]
			featureCounts[featureId]++ //increment the count for the total # of times the feature was a splitter

			if split.Numerical == true {
				if splitValueList[featureId] == nil {
					splitValueList[featureId] = make([]float64, 0)
				}
				splitValueList[featureId] = append(splitValueList[featureId], float64(split.Value))
			}
		}
	}

	log.Print("Outputing Split Feature/Case Cooccurance Counts")
	outfile, err := os.Create(*outf) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer outfile.Close()
	for feature, count := range featureCounts {
		if _, err := fmt.Fprintf(outfile, "%v\t%v\n", feature, count); err != nil {
			log.Fatal(err)
		}
	}

	log.Print("Outputing Split Distribution")
	splitdistfile, err := os.Create(*splitdistf) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer splitdistfile.Close()
	for feature, list := range splitValueList {
		if _, err := fmt.Fprintf(splitdistfile, "%v", feature); err != nil {
			log.Fatal(err)
		}
		for _, value := range list {
			if _, err := fmt.Fprintf(splitdistfile, "\t%g", value); err != nil {
				log.Fatal(err)
			}
		}
		fmt.Fprintf(splitdistfile, "\n")

	}

	log.Print("Outputing Case Feature Cooccurance Counts")
	boutfile, err := os.Create(*boutf) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer boutfile.Close()
	caseFeatureCounts.WriteTsv(boutfile)

	log.Print("Outputing Case Feature Splitter Direction")
	rboutfile, err := os.Create(*rboutf) // For read access.
	if err != nil {
		log.Fatal(err)
	}
	defer rboutfile.Close()
	relativeSplitCount.WriteTsv(rboutfile)
}
<|endoftext|>


<commit_before>package shared

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

type ContainerState struct {
	Status     string                           `json:"status"`
	StatusCode StatusCode                       `json:"status_code"`
	CPU        ContainerStateCPU                `json:"cpu"`
	Disk       map[string]ContainerStateDisk    `json:"disk"`
	Memory     ContainerStateMemory             `json:"memory"`
	Network    map[string]ContainerStateNetwork `json:"network"`
	Pid        int64                            `json:"pid"`
	Processes  int64                            `json:"processes"`
}

type ContainerStateDisk struct {
	Usage int64 `json:"usage"`
}

type ContainerStateCPU struct {
	Usage int64 `json:"usage"`
}

type ContainerStateMemory struct {
	Usage         int64 `json:"usage"`
	UsagePeak     int64 `json:"usage_peak"`
	SwapUsage     int64 `json:"swap_usage"`
	SwapUsagePeak int64 `json:"swap_usage_peak"`
}

type ContainerStateNetwork struct {
	Addresses []ContainerStateNetworkAddress `json:"addresses"`
	Counters  ContainerStateNetworkCounters  `json:"counters"`
	Hwaddr    string                         `json:"hwaddr"`
	HostName  string                         `json:"host_name"`
	Mtu       int                            `json:"mtu"`
	State     string                         `json:"state"`
	Type      string                         `json:"type"`
}

type ContainerStateNetworkAddress struct {
	Family  string `json:"family"`
	Address string `json:"address"`
	Netmask string `json:"netmask"`
	Scope   string `json:"scope"`
}

type ContainerStateNetworkCounters struct {
	BytesReceived   int64 `json:"bytes_received"`
	BytesSent       int64 `json:"bytes_sent"`
	PacketsReceived int64 `json:"packets_received"`
	PacketsSent     int64 `json:"packets_sent"`
}

type ContainerExecControl struct {
	Command string            `json:"command"`
	Args    map[string]string `json:"args"`
}

type SnapshotInfo struct {
	Architecture    string            `json:"architecture"`
	Config          map[string]string `json:"config"`
	CreationDate    time.Time         `json:"created_at"`
	Devices         Devices           `json:"devices"`
	Ephemeral       bool              `json:"ephemeral"`
	ExpandedConfig  map[string]string `json:"expanded_config"`
	ExpandedDevices Devices           `json:"expanded_devices"`
	LastUsedDate    time.Time         `json:"last_used_at"`
	Name            string            `json:"name"`
	Profiles        []string          `json:"profiles"`
	Stateful        bool              `json:"stateful"`
}

type ContainerInfo struct {
	Architecture    string            `json:"architecture"`
	Config          map[string]string `json:"config"`
	CreationDate    time.Time         `json:"created_at"`
	Devices         Devices           `json:"devices"`
	Ephemeral       bool              `json:"ephemeral"`
	ExpandedConfig  map[string]string `json:"expanded_config"`
	ExpandedDevices Devices           `json:"expanded_devices"`
	LastUsedDate    time.Time         `json:"last_used_at"`
	Name            string            `json:"name"`
	Profiles        []string          `json:"profiles"`
	Stateful        bool              `json:"stateful"`
	Status          string            `json:"status"`
	StatusCode      StatusCode        `json:"status_code"`
}

func (c ContainerInfo) IsActive() bool {
	switch c.StatusCode {
	case Stopped:
		return false
	case Error:
		return false
	default:
		return true
	}
}

/*
 * BriefContainerState contains a subset of the fields in
 * ContainerState, namely those which a user may update
 */
type BriefContainerInfo struct {
	Name      string            `json:"name"`
	Profiles  []string          `json:"profiles"`
	Config    map[string]string `json:"config"`
	Devices   Devices           `json:"devices"`
	Ephemeral bool              `json:"ephemeral"`
}

func (c *ContainerInfo) Brief() BriefContainerInfo {
	retstate := BriefContainerInfo{Name: c.Name,
		Profiles:  c.Profiles,
		Config:    c.Config,
		Devices:   c.Devices,
		Ephemeral: c.Ephemeral}
	return retstate
}

func (c *ContainerInfo) BriefExpanded() BriefContainerInfo {
	retstate := BriefContainerInfo{Name: c.Name,
		Profiles:  c.Profiles,
		Config:    c.ExpandedConfig,
		Devices:   c.ExpandedDevices,
		Ephemeral: c.Ephemeral}
	return retstate
}

type ContainerAction string

const (
	Stop     ContainerAction = "stop"
	Start    ContainerAction = "start"
	Restart  ContainerAction = "restart"
	Freeze   ContainerAction = "freeze"
	Unfreeze ContainerAction = "unfreeze"
)

type ProfileConfig struct {
	Name        string            `json:"name"`
	Config      map[string]string `json:"config"`
	Description string            `json:"description"`
	Devices     Devices           `json:"devices"`
	UsedBy      []string          `json:"used_by"`
}

type NetworkConfig struct {
	Name    string            `json:"name"`
	Config  map[string]string `json:"config"`
	Managed bool              `json:"managed"`
	Type    string            `json:"type"`
	UsedBy  []string          `json:"used_by"`
}

func IsInt64(value string) error {
	if value == "" {
		return nil
	}

	_, err := strconv.ParseInt(value, 10, 64)
	if err != nil {
		return fmt.Errorf("Invalid value for an integer: %s", value)
	}

	return nil
}

func IsPriority(value string) error {
	if value == "" {
		return nil
	}

	valueInt, err := strconv.ParseInt(value, 10, 64)
	if err != nil {
		return fmt.Errorf("Invalid value for an integer: %s", value)
	}

	if valueInt < 0 || valueInt > 10 {
		return fmt.Errorf("Invalid value for a limit '%s'. Must be between 0 and 10.", value)
	}

	return nil
}

func IsBool(value string) error {
	if value == "" {
		return nil
	}

	if !StringInSlice(strings.ToLower(value), []string{"true", "false", "yes", "no", "1", "0", "on", "off"}) {
		return fmt.Errorf("Invalid value for a boolean: %s", value)
	}

	return nil
}

func IsOneOf(value string, valid []string) error {
	if value == "" {
		return nil
	}

	if !StringInSlice(value, valid) {
		return fmt.Errorf("Invalid value: %s (not one of %s)", value, valid)
	}

	return nil
}

func IsAny(value string) error {
	return nil
}

// KnownContainerConfigKeys maps all fully defined, well-known config keys
// to an appropriate checker function, which validates whether or not a
// given value is syntactically legal.
var KnownContainerConfigKeys = map[string]func(value string) error{
	"boot.autostart":             IsBool,
	"boot.autostart.delay":       IsInt64,
	"boot.autostart.priority":    IsInt64,
	"boot.host_shutdown_timeout": IsInt64,

	"limits.cpu":           IsAny,
	"limits.cpu.allowance": IsAny,
	"limits.cpu.priority":  IsPriority,

	"limits.disk.priority": IsPriority,

	"limits.memory": IsAny,
	"limits.memory.enforce": func(value string) error {
		return IsOneOf(value, []string{"soft", "hard"})
	},
	"limits.memory.swap":          IsBool,
	"limits.memory.swap.priority": IsPriority,

	"limits.network.priority": IsPriority,

	"limits.processes": IsInt64,

	"linux.kernel_modules": IsAny,

	"security.nesting":    IsBool,
	"security.privileged": IsBool,

	"security.syscalls.blacklist_default": IsBool,
	"security.syscalls.blacklist_compat":  IsBool,
	"security.syscalls.blacklist":         IsAny,
	"security.syscalls.whitelist":         IsAny,

	// Caller is responsible for full validation of any raw.* value
	"raw.apparmor": IsAny,
	"raw.lxc":      IsAny,
	"raw.seccomp":  IsAny,

	"volatile.apply_template":   IsAny,
	"volatile.base_image":       IsAny,
	"volatile.last_state.idmap": IsAny,
	"volatile.last_state.power": IsAny,
}

// ConfigKeyChecker returns a function that will check whether or not
// a provide value is valid for the associate config key.  Returns an
// error if the key is not known.  The checker function only performs
// syntactic checking of the value, semantic and usage checking must
// be done by the caller.  User defined keys are always considered to
// be valid, e.g. user.* and environment.* keys.
func ConfigKeyChecker(key string) (func(value string) error, error) {
	if f, ok := KnownContainerConfigKeys[key]; ok {
		return f, nil
	}

	if strings.HasPrefix(key, "volatile.") {
		if strings.HasSuffix(key, ".hwaddr") {
			return IsAny, nil
		}

		if strings.HasSuffix(key, ".name") {
			return IsAny, nil
		}
	}

	if strings.HasPrefix(key, "environment.") {
		return IsAny, nil
	}

	if strings.HasPrefix(key, "user.") {
		return IsAny, nil
	}

	return nil, fmt.Errorf("Bad key: %s", key)
}
<commit_msg>Properly validate memory limits<commit_after>package shared

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

type ContainerState struct {
	Status     string                           `json:"status"`
	StatusCode StatusCode                       `json:"status_code"`
	CPU        ContainerStateCPU                `json:"cpu"`
	Disk       map[string]ContainerStateDisk    `json:"disk"`
	Memory     ContainerStateMemory             `json:"memory"`
	Network    map[string]ContainerStateNetwork `json:"network"`
	Pid        int64                            `json:"pid"`
	Processes  int64                            `json:"processes"`
}

type ContainerStateDisk struct {
	Usage int64 `json:"usage"`
}

type ContainerStateCPU struct {
	Usage int64 `json:"usage"`
}

type ContainerStateMemory struct {
	Usage         int64 `json:"usage"`
	UsagePeak     int64 `json:"usage_peak"`
	SwapUsage     int64 `json:"swap_usage"`
	SwapUsagePeak int64 `json:"swap_usage_peak"`
}

type ContainerStateNetwork struct {
	Addresses []ContainerStateNetworkAddress `json:"addresses"`
	Counters  ContainerStateNetworkCounters  `json:"counters"`
	Hwaddr    string                         `json:"hwaddr"`
	HostName  string                         `json:"host_name"`
	Mtu       int                            `json:"mtu"`
	State     string                         `json:"state"`
	Type      string                         `json:"type"`
}

type ContainerStateNetworkAddress struct {
	Family  string `json:"family"`
	Address string `json:"address"`
	Netmask string `json:"netmask"`
	Scope   string `json:"scope"`
}

type ContainerStateNetworkCounters struct {
	BytesReceived   int64 `json:"bytes_received"`
	BytesSent       int64 `json:"bytes_sent"`
	PacketsReceived int64 `json:"packets_received"`
	PacketsSent     int64 `json:"packets_sent"`
}

type ContainerExecControl struct {
	Command string            `json:"command"`
	Args    map[string]string `json:"args"`
}

type SnapshotInfo struct {
	Architecture    string            `json:"architecture"`
	Config          map[string]string `json:"config"`
	CreationDate    time.Time         `json:"created_at"`
	Devices         Devices           `json:"devices"`
	Ephemeral       bool              `json:"ephemeral"`
	ExpandedConfig  map[string]string `json:"expanded_config"`
	ExpandedDevices Devices           `json:"expanded_devices"`
	LastUsedDate    time.Time         `json:"last_used_at"`
	Name            string            `json:"name"`
	Profiles        []string          `json:"profiles"`
	Stateful        bool              `json:"stateful"`
}

type ContainerInfo struct {
	Architecture    string            `json:"architecture"`
	Config          map[string]string `json:"config"`
	CreationDate    time.Time         `json:"created_at"`
	Devices         Devices           `json:"devices"`
	Ephemeral       bool              `json:"ephemeral"`
	ExpandedConfig  map[string]string `json:"expanded_config"`
	ExpandedDevices Devices           `json:"expanded_devices"`
	LastUsedDate    time.Time         `json:"last_used_at"`
	Name            string            `json:"name"`
	Profiles        []string          `json:"profiles"`
	Stateful        bool              `json:"stateful"`
	Status          string            `json:"status"`
	StatusCode      StatusCode        `json:"status_code"`
}

func (c ContainerInfo) IsActive() bool {
	switch c.StatusCode {
	case Stopped:
		return false
	case Error:
		return false
	default:
		return true
	}
}

/*
 * BriefContainerState contains a subset of the fields in
 * ContainerState, namely those which a user may update
 */
type BriefContainerInfo struct {
	Name      string            `json:"name"`
	Profiles  []string          `json:"profiles"`
	Config    map[string]string `json:"config"`
	Devices   Devices           `json:"devices"`
	Ephemeral bool              `json:"ephemeral"`
}

func (c *ContainerInfo) Brief() BriefContainerInfo {
	retstate := BriefContainerInfo{Name: c.Name,
		Profiles:  c.Profiles,
		Config:    c.Config,
		Devices:   c.Devices,
		Ephemeral: c.Ephemeral}
	return retstate
}

func (c *ContainerInfo) BriefExpanded() BriefContainerInfo {
	retstate := BriefContainerInfo{Name: c.Name,
		Profiles:  c.Profiles,
		Config:    c.ExpandedConfig,
		Devices:   c.ExpandedDevices,
		Ephemeral: c.Ephemeral}
	return retstate
}

type ContainerAction string

const (
	Stop     ContainerAction = "stop"
	Start    ContainerAction = "start"
	Restart  ContainerAction = "restart"
	Freeze   ContainerAction = "freeze"
	Unfreeze ContainerAction = "unfreeze"
)

type ProfileConfig struct {
	Name        string            `json:"name"`
	Config      map[string]string `json:"config"`
	Description string            `json:"description"`
	Devices     Devices           `json:"devices"`
	UsedBy      []string          `json:"used_by"`
}

type NetworkConfig struct {
	Name    string            `json:"name"`
	Config  map[string]string `json:"config"`
	Managed bool              `json:"managed"`
	Type    string            `json:"type"`
	UsedBy  []string          `json:"used_by"`
}

func IsInt64(value string) error {
	if value == "" {
		return nil
	}

	_, err := strconv.ParseInt(value, 10, 64)
	if err != nil {
		return fmt.Errorf("Invalid value for an integer: %s", value)
	}

	return nil
}

func IsPriority(value string) error {
	if value == "" {
		return nil
	}

	valueInt, err := strconv.ParseInt(value, 10, 64)
	if err != nil {
		return fmt.Errorf("Invalid value for an integer: %s", value)
	}

	if valueInt < 0 || valueInt > 10 {
		return fmt.Errorf("Invalid value for a limit '%s'. Must be between 0 and 10.", value)
	}

	return nil
}

func IsBool(value string) error {
	if value == "" {
		return nil
	}

	if !StringInSlice(strings.ToLower(value), []string{"true", "false", "yes", "no", "1", "0", "on", "off"}) {
		return fmt.Errorf("Invalid value for a boolean: %s", value)
	}

	return nil
}

func IsOneOf(value string, valid []string) error {
	if value == "" {
		return nil
	}

	if !StringInSlice(value, valid) {
		return fmt.Errorf("Invalid value: %s (not one of %s)", value, valid)
	}

	return nil
}

func IsAny(value string) error {
	return nil
}

// KnownContainerConfigKeys maps all fully defined, well-known config keys
// to an appropriate checker function, which validates whether or not a
// given value is syntactically legal.
var KnownContainerConfigKeys = map[string]func(value string) error{
	"boot.autostart":             IsBool,
	"boot.autostart.delay":       IsInt64,
	"boot.autostart.priority":    IsInt64,
	"boot.host_shutdown_timeout": IsInt64,

	"limits.cpu":           IsAny,
	"limits.cpu.allowance": IsAny,
	"limits.cpu.priority":  IsPriority,

	"limits.disk.priority": IsPriority,

	"limits.memory": func(value string) error {
		if value == "" {
			return nil
		}

		if strings.HasSuffix(value, "%") {
			_, err := strconv.ParseInt(strings.TrimSuffix(value, "%"), 10, 64)
			if err != nil {
				return err
			}

			return nil
		}

		_, err := ParseByteSizeString(value)
		if err != nil {
			return err
		}

		return nil
	},
	"limits.memory.enforce": func(value string) error {
		return IsOneOf(value, []string{"soft", "hard"})
	},
	"limits.memory.swap":          IsBool,
	"limits.memory.swap.priority": IsPriority,

	"limits.network.priority": IsPriority,

	"limits.processes": IsInt64,

	"linux.kernel_modules": IsAny,

	"security.nesting":    IsBool,
	"security.privileged": IsBool,

	"security.syscalls.blacklist_default": IsBool,
	"security.syscalls.blacklist_compat":  IsBool,
	"security.syscalls.blacklist":         IsAny,
	"security.syscalls.whitelist":         IsAny,

	// Caller is responsible for full validation of any raw.* value
	"raw.apparmor": IsAny,
	"raw.lxc":      IsAny,
	"raw.seccomp":  IsAny,

	"volatile.apply_template":   IsAny,
	"volatile.base_image":       IsAny,
	"volatile.last_state.idmap": IsAny,
	"volatile.last_state.power": IsAny,
}

// ConfigKeyChecker returns a function that will check whether or not
// a provide value is valid for the associate config key.  Returns an
// error if the key is not known.  The checker function only performs
// syntactic checking of the value, semantic and usage checking must
// be done by the caller.  User defined keys are always considered to
// be valid, e.g. user.* and environment.* keys.
func ConfigKeyChecker(key string) (func(value string) error, error) {
	if f, ok := KnownContainerConfigKeys[key]; ok {
		return f, nil
	}

	if strings.HasPrefix(key, "volatile.") {
		if strings.HasSuffix(key, ".hwaddr") {
			return IsAny, nil
		}

		if strings.HasSuffix(key, ".name") {
			return IsAny, nil
		}
	}

	if strings.HasPrefix(key, "environment.") {
		return IsAny, nil
	}

	if strings.HasPrefix(key, "user.") {
		return IsAny, nil
	}

	return nil, fmt.Errorf("Bad key: %s", key)
}
<|endoftext|>


<commit_before>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// OpenPGP CFB Mode. http://tools.ietf.org/html/rfc4880#section-13.9

package cipher

type ocfb struct {
	b       Block
	fre     []byte
	outUsed int
}

// NewOCFBEncrypter returns a Stream which encrypts data with OpenPGP's cipher
// feedback mode using the given Block, and an initial amount of ciphertext.
// randData must be random bytes and be the same length as the Block's block
// size.
func NewOCFBEncrypter(block Block, randData []byte) (Stream, []byte) {
	blockSize := block.BlockSize()
	if len(randData) != blockSize {
		return nil, nil
	}

	x := &ocfb{
		b:       block,
		fre:     make([]byte, blockSize),
		outUsed: 0,
	}
	prefix := make([]byte, blockSize+2)

	block.Encrypt(x.fre, x.fre)
	for i := 0; i < blockSize; i++ {
		prefix[i] = randData[i] ^ x.fre[i]
	}

	block.Encrypt(x.fre, prefix[:blockSize])
	prefix[blockSize] = x.fre[0] ^ randData[blockSize-2]
	prefix[blockSize+1] = x.fre[1] ^ randData[blockSize-1]

	block.Encrypt(x.fre, prefix[2:])
	return x, prefix
}

// NewOCFBDecrypter returns a Stream which decrypts data with OpenPGP's cipher
// feedback mode using the given Block. Prefix must be the first blockSize + 2
// bytes of the ciphertext, where blockSize is the Block's block size. If an
// incorrect key is detected then nil is returned.
func NewOCFBDecrypter(block Block, prefix []byte) Stream {
	blockSize := block.BlockSize()
	if len(prefix) != blockSize+2 {
		return nil
	}

	x := &ocfb{
		b:       block,
		fre:     make([]byte, blockSize),
		outUsed: 0,
	}
	prefixCopy := make([]byte, len(prefix))
	copy(prefixCopy, prefix)

	block.Encrypt(x.fre, x.fre)
	for i := 0; i < blockSize; i++ {
		prefixCopy[i] ^= x.fre[i]
	}

	block.Encrypt(x.fre, prefix[:blockSize])
	prefixCopy[blockSize] ^= x.fre[0]
	prefixCopy[blockSize+1] ^= x.fre[1]

	if prefixCopy[blockSize-2] != prefixCopy[blockSize] ||
		prefixCopy[blockSize-1] != prefixCopy[blockSize+1] {
		return nil
	}

	block.Encrypt(x.fre, prefix[2:])
	return x
}

func (x *ocfb) XORKeyStream(dst, src []byte) {
	for i := 0; i < len(src); i++ {
		if x.outUsed == len(x.fre) {
			x.b.Encrypt(x.fre, x.fre)
			x.outUsed = 0
		}

		dst[i] = x.fre[x.outUsed] ^ src[i]
		x.outUsed++
	}
}
<commit_msg>crypto/cipher: fix OCFB<commit_after>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// OpenPGP CFB Mode. http://tools.ietf.org/html/rfc4880#section-13.9

package cipher

type ocfbEncrypter struct {
	b       Block
	fre     []byte
	outUsed int
}

// NewOCFBEncrypter returns a Stream which encrypts data with OpenPGP's cipher
// feedback mode using the given Block, and an initial amount of ciphertext.
// randData must be random bytes and be the same length as the Block's block
// size.
func NewOCFBEncrypter(block Block, randData []byte) (Stream, []byte) {
	blockSize := block.BlockSize()
	if len(randData) != blockSize {
		return nil, nil
	}

	x := &ocfbEncrypter{
		b:       block,
		fre:     make([]byte, blockSize),
		outUsed: 0,
	}
	prefix := make([]byte, blockSize+2)

	block.Encrypt(x.fre, x.fre)
	for i := 0; i < blockSize; i++ {
		prefix[i] = randData[i] ^ x.fre[i]
	}

	block.Encrypt(x.fre, prefix[:blockSize])
	prefix[blockSize] = x.fre[0] ^ randData[blockSize-2]
	prefix[blockSize+1] = x.fre[1] ^ randData[blockSize-1]

	block.Encrypt(x.fre, prefix[2:])
	return x, prefix
}

func (x *ocfbEncrypter) XORKeyStream(dst, src []byte) {
	for i := 0; i < len(src); i++ {
		if x.outUsed == len(x.fre) {
			x.b.Encrypt(x.fre, x.fre)
			x.outUsed = 0
		}

		x.fre[x.outUsed] ^= src[i]
		dst[i] = x.fre[x.outUsed]
		x.outUsed++
	}
}

type ocfbDecrypter struct {
	b       Block
	fre     []byte
	outUsed int
}

// NewOCFBDecrypter returns a Stream which decrypts data with OpenPGP's cipher
// feedback mode using the given Block. Prefix must be the first blockSize + 2
// bytes of the ciphertext, where blockSize is the Block's block size. If an
// incorrect key is detected then nil is returned.
func NewOCFBDecrypter(block Block, prefix []byte) Stream {
	blockSize := block.BlockSize()
	if len(prefix) != blockSize+2 {
		return nil
	}

	x := &ocfbDecrypter{
		b:       block,
		fre:     make([]byte, blockSize),
		outUsed: 0,
	}
	prefixCopy := make([]byte, len(prefix))
	copy(prefixCopy, prefix)

	block.Encrypt(x.fre, x.fre)
	for i := 0; i < blockSize; i++ {
		prefixCopy[i] ^= x.fre[i]
	}

	block.Encrypt(x.fre, prefix[:blockSize])
	prefixCopy[blockSize] ^= x.fre[0]
	prefixCopy[blockSize+1] ^= x.fre[1]

	if prefixCopy[blockSize-2] != prefixCopy[blockSize] ||
		prefixCopy[blockSize-1] != prefixCopy[blockSize+1] {
		return nil
	}

	block.Encrypt(x.fre, prefix[2:])
	return x
}

func (x *ocfbDecrypter) XORKeyStream(dst, src []byte) {
	for i := 0; i < len(src); i++ {
		if x.outUsed == len(x.fre) {
			x.b.Encrypt(x.fre, x.fre)
			x.outUsed = 0
		}

		c := src[i]
		dst[i] = x.fre[x.outUsed] ^ src[i]
		x.fre[x.outUsed] = c
		x.outUsed++
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/google/go-github/github"
	"github.com/parkr/auto-reply/ctx"
	"github.com/parkr/auto-reply/jekyll"
	"github.com/parkr/auto-reply/releases"
	"github.com/parkr/auto-reply/sentry"
	"golang.org/x/sync/errgroup"
)

var (
	defaultRepos = jekyll.DefaultRepos

	twoMonthsAgoUnix = time.Now().AddDate(0, -2, 0).Unix()

	issueLabels = []string{"release"}
)

func main() {
	var perform bool
	flag.BoolVar(&perform, "f", false, "Whether to actually file issues.")
	var inputRepos string
	flag.StringVar(&inputRepos, "repos", "", "Specify a list of comma-separated repo name/owner pairs, e.g. 'jekyll/jekyll-admin'.")
	flag.Parse()

	// Get latest 10 releases.
	// Sort releases by semver version, taking highest one.
	//
	// Has there been 100 commits since this release? If so, make an issue.
	// Has at least 1 commit been made since this release & is this release at least 2 month old? If so, make an issue.

	var repos []jekyll.Repository
	if inputRepos == "" {
		repos = defaultRepos
	}

	log.SetPrefix("nudge-maintainers-to-release: ")

	sentryClient, err := sentry.NewClient(map[string]string{
		"app":          "nudge-maintainers-to-release",
		"inputRepos":   inputRepos,
		"actuallyDoIt": fmt.Sprintf("%t", perform),
	})
	if err != nil {
		panic(err)
	}
	sentryClient.Recover(func() error {
		context := ctx.NewDefaultContext()
		if context.GitHub == nil {
			return errors.New("cannot proceed without github client")
		}

		if inputRepos != "" {
			repos = []jekyll.Repository{}
			for _, inputRepo := range strings.Split(inputRepos, ",") {
				pieces := strings.Split(inputRepo, "/")
				if len(pieces) != 2 {
					return fmt.Errorf("input repo %q is improperly formed", inputRepo)
				}
				repos = append(repos, jekyll.NewRepository(pieces[0], pieces[1]))
			}
		}

		wg, _ := errgroup.WithContext(context.Context())
		for _, repo := range repos {
			repo := repo
			wg.Go(func() error {
				latestRelease, err := releases.LatestRelease(context, repo)
				if err != nil {
					log.Printf("%s error fetching latest release: %+v", repo, err)
					return err
				}
				if latestRelease == nil {
					log.Printf("%s has no releases", repo)
					return nil
				}

				commitsSinceLatestRelease, err := releases.CommitsSinceRelease(context, repo, latestRelease)
				if err != nil {
					log.Printf("%s error fetching commits since latest release: %+v", repo, err)
					return err
				}

				if commitsSinceLatestRelease > 100 {
					if perform {
						return fileIssue(context, repo, latestRelease, "Over 100 commits have been made since the last release.")
					} else {
						log.Printf("%s would have been nudged for commits=%d", repo, commitsSinceLatestRelease)
					}
				} else if commitsSinceLatestRelease >= 1 && latestRelease.GetCreatedAt().Unix() <= twoMonthsAgoUnix {
					if perform {
						return fileIssue(context, repo, latestRelease, "The last release was over 2 months ago and there are unreleased commits on master.")
					} else {
						log.Printf("%s would have been nudged for date=%s commits=%d", repo, latestRelease.GetCreatedAt(), commitsSinceLatestRelease)
					}
				} else {
					log.Printf("%s is not in need of a release", repo)
				}

				return nil
			})
		}
		return wg.Wait()
	})
}

func fileIssue(context *ctx.Context, repo jekyll.Repository, latestRelease *github.RepositoryRelease, reason string) error {
	// TODO: does one already exist?
	issue, _, err := context.GitHub.Issues.Create(
		context.Context(),
		repo.Owner(), repo.Name(),
		&github.IssueRequest{
			Title:  github.String("Time for a new release!"),
			Labels: &issueLabels,
			Body: github.String(strings.TrimSpace(fmt.Sprintf(`
Hello, fine maintainers!

You've made some wonderful progress! %s

Would you mind shipping a new release soon so our users can enjoy the optimizations the community has made on master?

Thanks! :revolving_hearts: :sparkles:
`, reason))),
		},
	)
	if err != nil {
		return err
	}

	log.Printf("%s filed %s", repo, issue.HTMLURL)
	return nil
}
<commit_msg>nudge-maintainers-to-release: be idempotent & use an html/template.Template<commit_after>package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"html/template"
	"log"
	"strings"
	"time"

	"github.com/google/go-github/github"
	"github.com/parkr/auto-reply/ctx"
	"github.com/parkr/auto-reply/jekyll"
	"github.com/parkr/auto-reply/releases"
	"github.com/parkr/auto-reply/search"
	"github.com/parkr/auto-reply/sentry"
	"github.com/parkr/githubapi/githubsearch"
	"golang.org/x/sync/errgroup"
)

var (
	defaultRepos = jekyll.DefaultRepos

	twoMonthsAgoUnix = time.Now().AddDate(0, -2, 0).Unix()

	issueTitle  = "Time for a new release"
	issueLabels = []string{"release"}

	issueBodyTemplate = template.Must(template.New("issueBodyTemplate").Parse(`
Hello, maintainers! :wave:

By my calculations, it's time for a new release of {{.Repo.Name}}. {{if gt .CommitsOnMasterSinceLatestRelease 100}}There have been {{.CommitsOnMasterSinceLatestRelease}} commits{{else}}It's been over 2 months{{end}} since the last release, {{.LatestRelease.TagName}}.

What else is left to be done before a new release can be made? Please make sure to update History.markdown too if it's not already updated.

Thanks! :revolving_hearts: :sparkles:
`))
)

type templateInfo struct {
	Repo                              jekyll.Repository
	CommitsOnMasterSinceLatestRelease int
	LatestRelease                     *github.RepositoryRelease
}

func main() {
	var perform bool
	flag.BoolVar(&perform, "f", false, "Whether to actually file issues.")
	var inputRepos string
	flag.StringVar(&inputRepos, "repos", "", "Specify a list of comma-separated repo name/owner pairs, e.g. 'jekyll/jekyll-admin'.")
	flag.Parse()

	// Get latest 10 releases.
	// Sort releases by semver version, taking highest one.
	//
	// Has there been 100 commits since this release? If so, make an issue.
	// Has at least 1 commit been made since this release & is this release at least 2 month old? If so, make an issue.

	var repos []jekyll.Repository
	if inputRepos == "" {
		repos = defaultRepos
	}

	log.SetPrefix("nudge-maintainers-to-release: ")

	sentryClient, err := sentry.NewClient(map[string]string{
		"app":          "nudge-maintainers-to-release",
		"inputRepos":   inputRepos,
		"actuallyDoIt": fmt.Sprintf("%t", perform),
	})
	if err != nil {
		panic(err)
	}
	sentryClient.Recover(func() error {
		context := ctx.NewDefaultContext()
		if context.GitHub == nil {
			return errors.New("cannot proceed without github client")
		}

		if inputRepos != "" {
			repos = []jekyll.Repository{}
			for _, inputRepo := range strings.Split(inputRepos, ",") {
				pieces := strings.Split(inputRepo, "/")
				if len(pieces) != 2 {
					return fmt.Errorf("input repo %q is improperly formed", inputRepo)
				}
				repos = append(repos, jekyll.NewRepository(pieces[0], pieces[1]))
			}
		}

		wg, _ := errgroup.WithContext(context.Context())
		for _, repo := range repos {
			repo := repo
			wg.Go(func() error {
				latestRelease, err := releases.LatestRelease(context, repo)
				if err != nil {
					log.Printf("%s error fetching latest release: %+v", repo, err)
					return err
				}
				if latestRelease == nil {
					log.Printf("%s has no releases", repo)
					return nil
				}

				commitsSinceLatestRelease, err := releases.CommitsSinceRelease(context, repo, latestRelease)
				if err != nil {
					log.Printf("%s error fetching commits since latest release: %+v", repo, err)
					return err
				}

				if commitsSinceLatestRelease > 100 || (commitsSinceLatestRelease >= 1 && latestRelease.GetCreatedAt().Unix() <= twoMonthsAgoUnix) {
					if perform {
						err := fileIssue(context, templateInfo{
							Repo:                              repo,
							LatestRelease:                     latestRelease,
							CommitsOnMasterSinceLatestRelease: commitsSinceLatestRelease,
						})
						if err != nil {
							log.Printf("%s: nudged maintainers (release=%s commits=%d released_on=%s)",
								repo,
								latestRelease.GetTagName(),
								commitsSinceLatestRelease,
								latestRelease.GetCreatedAt(),
							)
						}
						return err
					} else {
						log.Printf("%s is in need of a nudge (release=%s commits=%d released_on=%s)",
							repo,
							latestRelease.GetTagName(),
							commitsSinceLatestRelease,
							latestRelease.GetCreatedAt(),
						)
					}
				} else {
					log.Printf("%s is NOT in need of a nudge: (release=%s commits=%d released_on=%s)",
						repo,
						latestRelease.GetTagName(),
						commitsSinceLatestRelease,
						latestRelease.GetCreatedAt(),
					)
				}

				return nil
			})
		}
		return wg.Wait()
	})
}

func fileIssue(context *ctx.Context, issueInfo templateInfo) error {
	if issue := getReleaseNudgeIssue(context, issueInfo.Repo); issue != nil {
		return fmt.Errorf("%s: issue already exists: %s", issueInfo.Repo, issue.GetHTMLURL())
	}

	var body bytes.Buffer
	if err := issueBodyTemplate.Execute(&body, issueInfo); err != nil {
		return fmt.Errorf("%s: error executing template: %+v", issueInfo.Repo, err)
	}

	issue, _, err := context.GitHub.Issues.Create(
		context.Context(),
		issueInfo.Repo.Owner(), issueInfo.Repo.Name(),
		&github.IssueRequest{
			Title:  &issueTitle,
			Labels: &issueLabels,
			Body:   github.String(body.String()),
		},
	)
	if err != nil {
		return fmt.Errorf("%s: error filing issue: %+v", issueInfo.Repo, err)
	}

	log.Printf("%s filed %s", issueInfo.Repo, issue.GetHTMLURL())
	return nil
}

func getReleaseNudgeIssue(context *ctx.Context, repo jekyll.Repository) *github.Issue {
	query := githubsearch.IssueSearchParameters{
		Type:       githubsearch.Issue,
		Scope:      githubsearch.TitleScope,
		Author:     context.CurrentlyAuthedGitHubUser().GetLogin(),
		State:      githubsearch.Open,
		Repository: &githubsearch.RepositoryName{Owner: repo.Owner(), Name: repo.Name()},
		Query:      issueTitle,
	}
	issues, err := search.GitHubIssues(context, query)
	if err != nil {
		log.Printf("%s: error searching %s: %+v", repo, query, err)
		return nil
	}
	if len(issues) > 0 {
		return &(issues[0])
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"golang.org/x/oauth2"
	"net/http"
	"os"
	"testing"
)

func handler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "Hi there, I love %s!", r.URL.Path[1:])
}

func init() {
	http.HandleFunc("/", handler)
	go http.ListenAndServe(":12345", nil)
}

var testOauth2Config = oauth2.Config{
	ClientID:     "foo",
	ClientSecret: "bar",
	Endpoint: oauth2.Endpoint{
		AuthURL:  "http://localhost:12345/auth",
		TokenURL: "http://localhost:12345/token"},
	RedirectURL: "https://example.com" + redirectPath,
	Scopes:      []string{"openidc", "email"},
}

func TestOauth2DoRedirectoToProviderHandlerSuccess(t *testing.T) {
	state, passwdFile, err := setupValidRuntimeStateSigner()
	if err != nil {
		t.Fatal(err)
	}
	defer os.Remove(passwdFile.Name()) // clean up
	state.authCookie = make(map[string]authInfo)
	state.pendingOauth2 = make(map[string]pendingAuth2Request)
	state.Config.Oauth2.Config = &testOauth2Config

	req, err := http.NewRequest("GET", oauth2LoginBeginPath, nil)
	if err != nil {
		t.Fatal(err)
	}
	_, err = checkRequestHandlerCode(req, state.oauth2DoRedirectoToProviderHandler, http.StatusFound)
	if err != nil {
		t.Fatal(err)
	}
	// Todo Check for the response contents
}
<commit_msg>start of main workhorse test<commit_after>package main

import (
	"fmt"
	"golang.org/x/oauth2"
	"net/http"
	"os"
	"testing"
)

func handler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "Hi there, I love %s!", r.URL.Path[1:])
}

func init() {
	http.HandleFunc("/", handler)
	go http.ListenAndServe(":12345", nil)
}

var testOauth2Config = oauth2.Config{
	ClientID:     "foo",
	ClientSecret: "bar",
	Endpoint: oauth2.Endpoint{
		AuthURL:  "http://localhost:12345/auth",
		TokenURL: "http://localhost:12345/token"},
	RedirectURL: "https://example.com" + redirectPath,
	Scopes:      []string{"openidc", "email"},
}

func TestOauth2DoRedirectoToProviderHandlerSuccess(t *testing.T) {
	state, passwdFile, err := setupValidRuntimeStateSigner()
	if err != nil {
		t.Fatal(err)
	}
	defer os.Remove(passwdFile.Name()) // clean up
	state.authCookie = make(map[string]authInfo)
	state.pendingOauth2 = make(map[string]pendingAuth2Request)
	state.Config.Oauth2.Config = &testOauth2Config

	req, err := http.NewRequest("GET", oauth2LoginBeginPath, nil)
	if err != nil {
		t.Fatal(err)
	}
	_, err = checkRequestHandlerCode(req, state.oauth2DoRedirectoToProviderHandler, http.StatusFound)
	if err != nil {
		t.Fatal(err)
	}
	// Todo Check for the response contents
}

func TestOauth2RedirectPathHandlerSuccess(t *testing.T) {
	state, passwdFile, err := setupValidRuntimeStateSigner()
	if err != nil {
		t.Fatal(err)
	}
	defer os.Remove(passwdFile.Name()) // clean up
	state.authCookie = make(map[string]authInfo)
	state.pendingOauth2 = make(map[string]pendingAuth2Request)
	state.Config.Oauth2.Config = &testOauth2Config

	//initially the request should fail for lack of preconditions
	req, err := http.NewRequest("GET", redirectPath, nil)
	if err != nil {
		t.Fatal(err)
	}
	// request has no cookies
	_, err = checkRequestHandlerCode(req, state.oauth2RedirectPathHandler, http.StatusBadRequest)
	if err != nil {
		t.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Package cli is a library to help creating command line tools.
package cli

import (
	"flag"
	"fmt"
	"os"
)

// To add a module, implement this interface. Definition is the command
// definition. Exec is the behaviour that you want to implement as a command
type Command interface {
	Definition() string // usually it's the output for --help
	Exec() error
}

// Module is the shared structure of commands and sub-commands.
type Module struct {
	children map[string]*Module // Non-nil if sub-command
	command  Command            // Non-nil if command
}

// NewCLI returns a root Module that you can add commands and
// another modules (sub-commands).
func NewCLI() *Module {
	return &Module{
		children: make(map[string]*Module, 0),
	}
}

// AddCommand adds a new command this module.
func (m *Module) AddCommand(name string, command Command) *Module {
	child := &Module{
		command: command,
	}
	m.children[name] = child
	return child
}

// AddSubCommand adds a new sub-command this module.
func (m *Module) AddSubCommand(name string) *Module {
	child := &Module{
		children: make(map[string]*Module, 0),
	}
	m.children[name] = child
	return child
}

// Run is the function that is intended to be run from main().
func (m *Module) Run() {
	flag.Parse()
	args := flag.Args()

	command, err := m.findCommand(args)
	if err != nil {
		exitErr(err)
	}

	err = command.Exec()
	if err != nil {
		exitErr(err)
	}

	os.Exit(0)
}

func (m *Module) findCommand(args []string) (Command, error) {
	// Iterate over args and update the module pointer "m"
	for i, arg := range args {
		// Treat m as a module (sub-command)
		sub := m.children[arg]
		if sub == nil {
			return nil, fmt.Errorf("Command not found")
		}

		// sub is another module here
		if sub.command == nil {
			m = m.children[arg]
			continue
		}

		// command behaves like a subprocess, it will parse arguments again
		// so we re discarding parsed arguments
		temp := os.Args
		os.Args = []string{temp[0]}
		os.Args = append(os.Args, temp[i+2:]...)

		// Returning command module
		return sub, nil
	}

	// Returning sub-command module
	return m, nil
}

func exitErr(err error) {
	fmt.Fprintf(os.Stderr, "%s\n", err.Error())
	os.Exit(1)
}

////////////////////////////////////////////////////////////////////////
// Methods below implement Command interface for Module (sub-command) //
////////////////////////////////////////////////////////////////////////

func (m *Module) Definition() string {
	if m.command != nil {
		return m.command.Definition()
	}

	return fmt.Sprintf("Run to see sub-commands")
}

func (m *Module) Exec() error {
	if m.command != nil {
		return m.command.Exec()
	}

	// Print command list
	fmt.Println("Possible commands:")
	for n, module := range m.children {
		fmt.Printf("  %-10s  ", n)

		if module.command != nil {
			fmt.Println(module.command.Definition())
		} else {
			fmt.Println(module.Definition())
		}
	}

	return nil
}
<commit_msg>do not modify os.Args<commit_after>// Package cli is a library to help creating command line tools.
package cli

import (
	"flag"
	"fmt"
	"os"
)

// To add a module, implement this interface. Definition is the command
// definition. Exec is the behaviour that you want to implement as a command
type Command interface {
	Definition() string // usually it's the output for --help
	Exec() error
}

// Module is the shared structure of commands and sub-commands.
type Module struct {
	children map[string]*Module // Non-nil if sub-command
	command  Command            // Non-nil if command
}

// NewCLI returns a root Module that you can add commands and
// another modules (sub-commands).
func NewCLI() *Module {
	return &Module{
		children: make(map[string]*Module, 0),
	}
}

// AddCommand adds a new command this module.
func (m *Module) AddCommand(name string, command Command) {
	child := &Module{
		command: command,
	}
	m.children[name] = child
}

// AddSubCommand adds a new sub-command this module.
func (m *Module) AddSubCommand(name string) *Module {
	child := &Module{
		children: make(map[string]*Module, 0),
	}
	m.children[name] = child
	return child
}

// Run is the function that is intended to be run from main().
func (m *Module) Run() {
	flag.Parse()
	args := flag.Args()

	command, err := m.findCommand(args)
	if err != nil {
		exitErr(err)
	}

	err = command.Exec()
	if err != nil {
		exitErr(err)
	}

	os.Exit(0)
}

func (m *Module) findCommand(args []string) (Command, error) {
	// Iterate over args and update the module pointer "m"
	for _, arg := range args {
		// Treat m as a module (sub-command)
		sub := m.children[arg]
		if sub == nil {
			return nil, fmt.Errorf("Command not found")
		}

		// sub is another module here
		if sub.command == nil {
			m = m.children[arg]
			continue
		}

		args = args[1:]

		// Returning command module
		return sub, nil
	}

	// Returning sub-command module
	return m, nil
}

func exitErr(err error) {
	fmt.Fprintf(os.Stderr, "%s\n", err.Error())
	os.Exit(1)
}

////////////////////////////////////////////////////////////////////////
// Methods below implement Command interface for Module (sub-command) //
////////////////////////////////////////////////////////////////////////

func (m *Module) Definition() string {
	if m.command != nil {
		return m.command.Definition()
	}

	return fmt.Sprintf("Run to see sub-commands")
}

func (m *Module) Exec() error {
	if m.command != nil {
		return m.command.Exec()
	}

	// Print command list
	fmt.Println("Possible commands:")
	for n, module := range m.children {
		fmt.Printf("  %-10s  ", n)

		if module.command != nil {
			fmt.Println(module.command.Definition())
		} else {
			fmt.Println(module.Definition())
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/xml"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/fatih/color"
)

type Reporter interface {
	Report(result TestResult)
	Flush()
}

type ConsoleReporter struct {
	total  int
	failed int
}

func (r *ConsoleReporter) Report(result TestResult) {
	r.total = r.total + 1
	if result.Cause != nil {
		r.failed = r.failed + 1
		r.reportError(result)
	} else {
		r.reportSuccess(result)
	}
}

func (r ConsoleReporter) reportSuccess(result TestResult) {
	c := color.New(color.FgGreen).Add(color.Bold)
	fmt.Printf("[")
	c.Print("PASSED")
	fmt.Printf("] %s\n", result.Case.Description)
}

func (r ConsoleReporter) reportError(result TestResult) {
	c := color.New(color.FgRed).Add(color.Bold)
	fmt.Printf("[")
	c.Print("FAILED")
	fmt.Printf("] %s\n", result.Case.Description)
	lines := strings.Split(result.Cause.Error(), "\n")

	for _, line := range lines {
		fmt.Printf("\t\t%s \n", line)
	}
}

func (r ConsoleReporter) Flush() {
	fmt.Println("~~~ Summary ~~~")
	fmt.Printf("# of test cases : %v\n", r.total)
	fmt.Printf("# Errors: %v\n", r.failed)

	if r.failed > 0 {
		fmt.Println("~~~ Test run FAILURE! ~~~")
		os.Exit(1)
		return
	} // test run failed

	fmt.Println("~~~ Test run SUCCESS ~~~")
	os.Exit(0)
}

// NewConsoleReporter returns new instance of console reporter
func NewConsoleReporter() Reporter {
	return &ConsoleReporter{}
}

// JUnitXMLReporter produces separate xml file for each test sute
type JUnitXMLReporter struct {
	// output directory
	OutPath string

	// current suite
	// when suite is being changed, flush previous one
	suite *suite
}

type suite struct {
	XMLName     string `xml:"testsuite"`
	ID          int    `xml:"id,attr"`
	Name        string `xml:"name,attr"`
	PackageName string `xml:"package,attr"`
	TimeStamp   string `xml:"timestamp,attr"`
	Time        uint16 `xml:"time,attr"`
	HostName    string `xml:"hostname,attr"`

	Tests    int `xml:"tests,attr"`
	Failures int `xml:"failures,attr"`
	Errors   int `xml:"errors,attr"`

	Properties properties `xml:"properties"`
	Cases      []tc       `xml:"testcase"`

	SystemOut string `xml:"system-out"`
	SystemErr string `xml:"system-err"`
}

type properties struct {
}

type tc struct {
	Name      string   `xml:"name,attr"`
	ClassName string   `xml:"classname,attr"`
	Time      uint16   `xml:"time,attr"`
	Failure   *failure `xml:"failure,omitempty"`
}

type failure struct {
	// not clear what type is but it's required
	Type    string `xml:"type,attr"`
	Message string `xml:"message,attr"`
}

func (r *JUnitXMLReporter) Report(result TestResult) {
	if r.suite == nil {
		r.suite = newSuite(result)
	}

	if r.suite.Name != result.Suite.Name {
		r.flushSuite()
		r.suite = newSuite(result)
	}

	testCase := tc{Name: result.Case.Description}
	if result.Cause != nil {
		testCase.Failure = &failure{Type: result.Cause.Error(), Message: result.Cause.Error()}
		r.suite.Failures = r.suite.Failures + 1
	}
	r.suite.Tests = r.suite.Tests + 1
	r.suite.ID = r.suite.ID + 1
	r.suite.Cases = append(r.suite.Cases, testCase)
}

func (r JUnitXMLReporter) flushSuite() {
	fileName := strings.Replace(r.suite.PackageName, "/", "_", -1) + r.suite.Name + ".xml"
	fp := filepath.Join(r.OutPath, fileName)
	err := os.MkdirAll(r.OutPath, 0777)
	if err != nil {
		panic(err)
	}
	f, err := os.Create(fp)
	if err != nil {
		panic(err)
	}

	data, err := xml.Marshal(r.suite)
	if err != nil {
		panic(err)
	}

	f.Write(data)
}

func newSuite(result TestResult) *suite {
	return &suite{
		ID:          0,
		Name:        result.Suite.Name,
		PackageName: result.Suite.PackageName,
		TimeStamp:   time.Now().UTC().Format("2006-01-02T15:04:05"),
		HostName:    "test",
	}
}

func (r JUnitXMLReporter) Flush() {
	// r.flushSuite()
}

func NewJUnitReporter(outdir string) Reporter {
	return &JUnitXMLReporter{OutPath: outdir}
}
<commit_msg>flush last suite<commit_after>package main

import (
	"encoding/xml"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/fatih/color"
)

type Reporter interface {
	Report(result TestResult)
	Flush()
}

type ConsoleReporter struct {
	total  int
	failed int
}

func (r *ConsoleReporter) Report(result TestResult) {
	r.total = r.total + 1
	if result.Cause != nil {
		r.failed = r.failed + 1
		r.reportError(result)
	} else {
		r.reportSuccess(result)
	}
}

func (r ConsoleReporter) reportSuccess(result TestResult) {
	c := color.New(color.FgGreen).Add(color.Bold)
	fmt.Printf("[")
	c.Print("PASSED")
	fmt.Printf("] %s\n", result.Case.Description)
}

func (r ConsoleReporter) reportError(result TestResult) {
	c := color.New(color.FgRed).Add(color.Bold)
	fmt.Printf("[")
	c.Print("FAILED")
	fmt.Printf("] %s\n", result.Case.Description)
	lines := strings.Split(result.Cause.Error(), "\n")

	for _, line := range lines {
		fmt.Printf("\t\t%s \n", line)
	}
}

func (r ConsoleReporter) Flush() {
	fmt.Println("~~~ Summary ~~~")
	fmt.Printf("# of test cases : %v\n", r.total)
	fmt.Printf("# Errors: %v\n", r.failed)

	if r.failed > 0 {
		fmt.Println("~~~ Test run FAILURE! ~~~")
		os.Exit(1)
		return
	} // test run failed

	fmt.Println("~~~ Test run SUCCESS ~~~")
	os.Exit(0)
}

// NewConsoleReporter returns new instance of console reporter
func NewConsoleReporter() Reporter {
	return &ConsoleReporter{}
}

// JUnitXMLReporter produces separate xml file for each test sute
type JUnitXMLReporter struct {
	// output directory
	OutPath string

	// current suite
	// when suite is being changed, flush previous one
	suite *suite
}

type suite struct {
	XMLName     string `xml:"testsuite"`
	ID          int    `xml:"id,attr"`
	Name        string `xml:"name,attr"`
	PackageName string `xml:"package,attr"`
	TimeStamp   string `xml:"timestamp,attr"`
	Time        uint16 `xml:"time,attr"`
	HostName    string `xml:"hostname,attr"`

	Tests    int `xml:"tests,attr"`
	Failures int `xml:"failures,attr"`
	Errors   int `xml:"errors,attr"`

	Properties properties `xml:"properties"`
	Cases      []tc       `xml:"testcase"`

	SystemOut string `xml:"system-out"`
	SystemErr string `xml:"system-err"`
}

type properties struct {
}

type tc struct {
	Name      string   `xml:"name,attr"`
	ClassName string   `xml:"classname,attr"`
	Time      uint16   `xml:"time,attr"`
	Failure   *failure `xml:"failure,omitempty"`
}

type failure struct {
	// not clear what type is but it's required
	Type    string `xml:"type,attr"`
	Message string `xml:"message,attr"`
}

func (r *JUnitXMLReporter) Report(result TestResult) {
	if r.suite == nil {
		r.suite = newSuite(result)
	}

	if r.suite.Name != result.Suite.Name {
		r.flushSuite()
		r.suite = newSuite(result)
	}

	testCase := tc{Name: result.Case.Description}
	if result.Cause != nil {
		testCase.Failure = &failure{Type: result.Cause.Error(), Message: result.Cause.Error()}
		r.suite.Failures = r.suite.Failures + 1
	}
	r.suite.Tests = r.suite.Tests + 1
	r.suite.ID = r.suite.ID + 1
	r.suite.Cases = append(r.suite.Cases, testCase)
}

func (r JUnitXMLReporter) flushSuite() {
	fileName := strings.Replace(r.suite.PackageName, "/", "_", -1) + r.suite.Name + ".xml"
	fp := filepath.Join(r.OutPath, fileName)
	err := os.MkdirAll(r.OutPath, 0777)
	if err != nil {
		panic(err)
	}
	f, err := os.Create(fp)
	if err != nil {
		panic(err)
	}

	data, err := xml.Marshal(r.suite)
	if err != nil {
		panic(err)
	}

	f.Write(data)
}

func newSuite(result TestResult) *suite {
	return &suite{
		ID:          0,
		Name:        result.Suite.Name,
		PackageName: result.Suite.PackageName,
		TimeStamp:   time.Now().UTC().Format("2006-01-02T15:04:05"),
		HostName:    "test",
	}
}

func (r JUnitXMLReporter) Flush() {
	r.flushSuite()
}

func NewJUnitReporter(outdir string) Reporter {
	return &JUnitXMLReporter{OutPath: outdir}
}
<|endoftext|>