<commit_before>package kafka

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"reflect"
)

var errShortRead = errors.New("not enough bytes available to load the response")

func peekRead(r *bufio.Reader, sz int, n int, f func([]byte)) (int, error) {
	if n > sz {
		return sz, errShortRead
	}
	b, err := r.Peek(n)
	if err != nil {
		return sz, err
	}
	f(b)
	return discardN(r, sz, n)
}

func readInt8(r *bufio.Reader, sz int, v *int8) (int, error) {
	return peekRead(r, sz, 1, func(b []byte) { *v = makeInt8(b) })
}

func readInt16(r *bufio.Reader, sz int, v *int16) (int, error) {
	return peekRead(r, sz, 2, func(b []byte) { *v = makeInt16(b) })
}

func readInt32(r *bufio.Reader, sz int, v *int32) (int, error) {
	return peekRead(r, sz, 4, func(b []byte) { *v = makeInt32(b) })
}

func readInt64(r *bufio.Reader, sz int, v *int64) (int, error) {
	return peekRead(r, sz, 8, func(b []byte) { *v = makeInt64(b) })
}

func readString(r *bufio.Reader, sz int, v *string) (int, error) {
	return readStringWith(r, sz, func(r *bufio.Reader, sz int, len int16) (int, error) {
		var b []byte
		var n = int(len)

		if n > sz {
			return sz, errShortRead
		}
		b = make([]byte, n)

		n, err := io.ReadFull(r, b)
		if err == nil {
			*v = string(b)
		}

		return sz - n, err
	})
}

func readStringWith(r *bufio.Reader, sz int, cb func(*bufio.Reader, int, int16) (int, error)) (int, error) {
	var err error
	var len int16

	if sz, err = readInt16(r, sz, &len); err != nil {
		return sz, err
	}

	if len < 0 {
		len = 0
	}

	return cb(r, sz, len)
}

func readBytes(r *bufio.Reader, sz int, v *[]byte) (int, error) {
	return readStringWith(r, sz, func(r *bufio.Reader, sz int, len int16) (int, error) {
		var b []byte
		var n = int(len)

		if n > sz {
			return sz, errShortRead
		}
		b = make([]byte, n)

		n, err := io.ReadFull(r, b)
		if err == nil {
			*v = b
		}

		return sz - n, err
	})
}

func readBytesWith(r *bufio.Reader, sz int, cb func(*bufio.Reader, int, int32) (int, error)) (int, error) {
	var err error
	var len int32

	if sz, err = readInt32(r, sz, &len); err != nil {
		return sz, err
	}

	if len < 0 {
		len = 0
	}

	return cb(r, sz, len)
}

func readArrayWith(r *bufio.Reader, sz int, cb func(*bufio.Reader, int) (int, error)) (int, error) {
	var err error
	var len int32

	if sz, err = readInt32(r, sz, &len); err != nil {
		return sz, err
	}

	for n := int(len); n > 0; n-- {
		if sz, err = cb(r, sz); err != nil {
			break
		}
	}

	return sz, err
}

func read(r *bufio.Reader, sz int, a interface{}) (int, error) {
	switch v := a.(type) {
	case *int8:
		return readInt8(r, sz, v)
	case *int16:
		return readInt16(r, sz, v)
	case *int32:
		return readInt32(r, sz, v)
	case *int64:
		return readInt64(r, sz, v)
	case *string:
		return readString(r, sz, v)
	case *[]byte:
		return readBytes(r, sz, v)
	}
	switch v := reflect.ValueOf(a).Elem(); v.Kind() {
	case reflect.Struct:
		return readStruct(r, sz, v)
	case reflect.Slice:
		return readSlice(r, sz, v)
	default:
		panic(fmt.Sprintf("unsupported type: %T", a))
	}
}

func readStruct(r *bufio.Reader, sz int, v reflect.Value) (int, error) {
	var err error
	for i, n := 0, v.NumField(); i != n; i++ {
		if sz, err = read(r, sz, v.Field(i).Addr().Interface()); err != nil {
			return sz, err
		}
	}
	return sz, nil
}

func readSlice(r *bufio.Reader, sz int, v reflect.Value) (int, error) {
	var err error
	var len int32

	if sz, err = readInt32(r, sz, &len); err != nil {
		return sz, err
	}

	if n := int(len); n < 0 {
		v.Set(reflect.Zero(v.Type()))
	} else {
		v.Set(reflect.MakeSlice(v.Type(), n, n))

		for i := 0; i != n; i++ {
			if sz, err = read(r, sz, v.Index(i).Addr().Interface()); err != nil {
				return sz, err
			}
		}
	}

	return sz, nil
}

func readFetchResponseHeader(r *bufio.Reader, size int) (throttle int32, remain int, err error) {
	var n int32
	var p struct {
		Partition           int32
		ErrorCode           int16
		HighwaterMarkOffset int64
		MessageSetSize      int32
	}

	if remain, err = readInt32(r, size, &throttle); err != nil {
		return
	}

	if remain, err = readInt32(r, remain, &n); err != nil {
		return
	}

	// This error should never trigger, unless there's a bug in the kafka client
	// or server.
	if n != 1 {
		err = fmt.Errorf("1 kafka topic was expected in the fetch response but the client received %d", n)
		return
	}

	// We ignore the topic name because we've requests messages for a single
	// topic, unless there's a bug in the kafka server we will have received
	// the name of the topic that we requested.
	if remain, err = discardString(r, remain); err != nil {
		return
	}

	if remain, err = readInt32(r, remain, &n); err != nil {
		return
	}

	// This error should never trigger, unless there's a bug in the kafka client
	// or server.
	if n != 1 {
		err = fmt.Errorf("1 kafka partition was expected in the fetch response but the client received %d", n)
		return
	}

	if remain, err = read(r, remain, &p); err != nil {
		return
	}

	if p.ErrorCode != 0 {
		err = Error(p.ErrorCode)
		return
	}

	// This error should never trigger, unless there's a bug in the kafka client
	// or server.
	if remain != int(p.MessageSetSize) {
		err = fmt.Errorf("the size of the message set in a fetch response doesn't match the number of remaining bytes (message set size = %d, remaining bytes = %d)", p.MessageSetSize, remain)
		return
	}

	return
}

func readMessageHeader(r *bufio.Reader, sz int) (offset int64, attributes int8, timestamp int64, remain int, err error) {
	var version int8
	var msgsize int32

	if remain, err = readInt64(r, sz, &offset); err != nil {
		return
	}

	// Not sure why kafka gives the message size here, we already have the
	// number of remaining bytes in the response and kafka should only truncate
	// the trailing message. We still do a consistency check in case the server
	// returns something weird.
	if remain, err = readInt32(r, remain, &msgsize); err != nil {
		return
	}

	if msgsize < 0 {
		err = fmt.Errorf("invalid negative message size of %d in the fetch response", msgsize)
		return
	}

	if int(msgsize) > remain {
		err = fmt.Errorf("a message size of %d bytes in the fetch response exceeds the remaining size of %d bytes", msgsize, remain)
		return
	}

	// TCP is already taking care of ensuring data integrirty, no need to waste
	// resources doing it a second time so we just skip the message CRC.
	if remain, err = discardInt32(r, remain); err != nil {
		return
	}

	if remain, err = readInt8(r, remain, &version); err != nil {
		return
	}

	if remain, err = readInt8(r, remain, &attributes); err != nil {
		return
	}

	switch version {
	case 0:
	case 1:
		remain, err = readInt64(r, remain, &timestamp)
	default:
		err = fmt.Errorf("unsupported message version %d found in fetch response", version)
	}

	return
}

func readMessageBytes(r *bufio.Reader, sz int, read func(*bufio.Reader, int, int) (int, error)) (int, error) {
	return readBytesWith(r, sz, func(r *bufio.Reader, sz int, len int32) (int, error) {
		return read(r, sz, int(len))
	})
}

func readMessage(r *bufio.Reader, sz int,
	key func(*bufio.Reader, int, int) (int, error),
	val func(*bufio.Reader, int, int) (int, error),
) (offset int64, attributes int8, timestamp int64, remain int, err error) {
	if offset, attributes, timestamp, remain, err = readMessageHeader(r, sz); err != nil {
		return
	}
	if remain, err = readMessageBytes(r, remain, key); err != nil {
		return
	}
	remain, err = readMessageBytes(r, remain, val)
	return
}

func readResponse(r *bufio.Reader, sz int, res interface{}) error {
	sz, err := read(r, sz, res)
	switch err.(type) {
	case Error:
		sz, err = discardN(r, sz, sz)
	}
	return expectZeroSize(sz, err)
}
<commit_msg>return errShortBuffer when reading a string or bytes with not enough remaining bytes<commit_after>package kafka

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"reflect"
)

var errShortRead = errors.New("not enough bytes available to load the response")

func peekRead(r *bufio.Reader, sz int, n int, f func([]byte)) (int, error) {
	if n > sz {
		return sz, errShortRead
	}
	b, err := r.Peek(n)
	if err != nil {
		return sz, err
	}
	f(b)
	return discardN(r, sz, n)
}

func readInt8(r *bufio.Reader, sz int, v *int8) (int, error) {
	return peekRead(r, sz, 1, func(b []byte) { *v = makeInt8(b) })
}

func readInt16(r *bufio.Reader, sz int, v *int16) (int, error) {
	return peekRead(r, sz, 2, func(b []byte) { *v = makeInt16(b) })
}

func readInt32(r *bufio.Reader, sz int, v *int32) (int, error) {
	return peekRead(r, sz, 4, func(b []byte) { *v = makeInt32(b) })
}

func readInt64(r *bufio.Reader, sz int, v *int64) (int, error) {
	return peekRead(r, sz, 8, func(b []byte) { *v = makeInt64(b) })
}

func readString(r *bufio.Reader, sz int, v *string) (int, error) {
	return readStringWith(r, sz, func(r *bufio.Reader, sz int, len int16) (int, error) {
		var b []byte
		var n = int(len)

		if n > sz {
			return sz, errShortRead
		}
		b = make([]byte, n)

		n, err := io.ReadFull(r, b)
		if err == nil {
			*v = string(b)
		}

		return sz - n, err
	})
}

func readStringWith(r *bufio.Reader, sz int, cb func(*bufio.Reader, int, int16) (int, error)) (int, error) {
	var err error
	var len int16

	if sz, err = readInt16(r, sz, &len); err != nil {
		return sz, err
	}

	if len < 0 {
		len = 0
	}

	if sz < int(len) {
		return sz, errShortRead
	}

	return cb(r, sz, len)
}

func readBytes(r *bufio.Reader, sz int, v *[]byte) (int, error) {
	return readStringWith(r, sz, func(r *bufio.Reader, sz int, len int16) (int, error) {
		var b []byte
		var n = int(len)

		if n > sz {
			return sz, errShortRead
		}
		b = make([]byte, n)

		n, err := io.ReadFull(r, b)
		if err == nil {
			*v = b
		}

		return sz - n, err
	})
}

func readBytesWith(r *bufio.Reader, sz int, cb func(*bufio.Reader, int, int32) (int, error)) (int, error) {
	var err error
	var len int32

	if sz, err = readInt32(r, sz, &len); err != nil {
		return sz, err
	}

	if len < 0 {
		len = 0
	}

	if sz < int(len) {
		return sz, errShortRead
	}

	return cb(r, sz, len)
}

func readArrayWith(r *bufio.Reader, sz int, cb func(*bufio.Reader, int) (int, error)) (int, error) {
	var err error
	var len int32

	if sz, err = readInt32(r, sz, &len); err != nil {
		return sz, err
	}

	for n := int(len); n > 0; n-- {
		if sz, err = cb(r, sz); err != nil {
			break
		}
	}

	return sz, err
}

func read(r *bufio.Reader, sz int, a interface{}) (int, error) {
	switch v := a.(type) {
	case *int8:
		return readInt8(r, sz, v)
	case *int16:
		return readInt16(r, sz, v)
	case *int32:
		return readInt32(r, sz, v)
	case *int64:
		return readInt64(r, sz, v)
	case *string:
		return readString(r, sz, v)
	case *[]byte:
		return readBytes(r, sz, v)
	}
	switch v := reflect.ValueOf(a).Elem(); v.Kind() {
	case reflect.Struct:
		return readStruct(r, sz, v)
	case reflect.Slice:
		return readSlice(r, sz, v)
	default:
		panic(fmt.Sprintf("unsupported type: %T", a))
	}
}

func readStruct(r *bufio.Reader, sz int, v reflect.Value) (int, error) {
	var err error
	for i, n := 0, v.NumField(); i != n; i++ {
		if sz, err = read(r, sz, v.Field(i).Addr().Interface()); err != nil {
			return sz, err
		}
	}
	return sz, nil
}

func readSlice(r *bufio.Reader, sz int, v reflect.Value) (int, error) {
	var err error
	var len int32

	if sz, err = readInt32(r, sz, &len); err != nil {
		return sz, err
	}

	if n := int(len); n < 0 {
		v.Set(reflect.Zero(v.Type()))
	} else {
		v.Set(reflect.MakeSlice(v.Type(), n, n))

		for i := 0; i != n; i++ {
			if sz, err = read(r, sz, v.Index(i).Addr().Interface()); err != nil {
				return sz, err
			}
		}
	}

	return sz, nil
}

func readFetchResponseHeader(r *bufio.Reader, size int) (throttle int32, remain int, err error) {
	var n int32
	var p struct {
		Partition           int32
		ErrorCode           int16
		HighwaterMarkOffset int64
		MessageSetSize      int32
	}

	if remain, err = readInt32(r, size, &throttle); err != nil {
		return
	}

	if remain, err = readInt32(r, remain, &n); err != nil {
		return
	}

	// This error should never trigger, unless there's a bug in the kafka client
	// or server.
	if n != 1 {
		err = fmt.Errorf("1 kafka topic was expected in the fetch response but the client received %d", n)
		return
	}

	// We ignore the topic name because we've requests messages for a single
	// topic, unless there's a bug in the kafka server we will have received
	// the name of the topic that we requested.
	if remain, err = discardString(r, remain); err != nil {
		return
	}

	if remain, err = readInt32(r, remain, &n); err != nil {
		return
	}

	// This error should never trigger, unless there's a bug in the kafka client
	// or server.
	if n != 1 {
		err = fmt.Errorf("1 kafka partition was expected in the fetch response but the client received %d", n)
		return
	}

	if remain, err = read(r, remain, &p); err != nil {
		return
	}

	if p.ErrorCode != 0 {
		err = Error(p.ErrorCode)
		return
	}

	// This error should never trigger, unless there's a bug in the kafka client
	// or server.
	if remain != int(p.MessageSetSize) {
		err = fmt.Errorf("the size of the message set in a fetch response doesn't match the number of remaining bytes (message set size = %d, remaining bytes = %d)", p.MessageSetSize, remain)
		return
	}

	return
}

func readMessageHeader(r *bufio.Reader, sz int) (offset int64, attributes int8, timestamp int64, remain int, err error) {
	var version int8
	var msgsize int32

	if remain, err = readInt64(r, sz, &offset); err != nil {
		return
	}

	// Not sure why kafka gives the message size here, we already have the
	// number of remaining bytes in the response and kafka should only truncate
	// the trailing message. We still do a consistency check in case the server
	// returns something weird.
	if remain, err = readInt32(r, remain, &msgsize); err != nil {
		return
	}

	if msgsize < 0 {
		err = fmt.Errorf("invalid negative message size of %d in the fetch response", msgsize)
		return
	}

	if int(msgsize) > remain {
		err = fmt.Errorf("a message size of %d bytes in the fetch response exceeds the remaining size of %d bytes", msgsize, remain)
		return
	}

	// TCP is already taking care of ensuring data integrirty, no need to waste
	// resources doing it a second time so we just skip the message CRC.
	if remain, err = discardInt32(r, remain); err != nil {
		return
	}

	if remain, err = readInt8(r, remain, &version); err != nil {
		return
	}

	if remain, err = readInt8(r, remain, &attributes); err != nil {
		return
	}

	switch version {
	case 0:
	case 1:
		remain, err = readInt64(r, remain, &timestamp)
	default:
		err = fmt.Errorf("unsupported message version %d found in fetch response", version)
	}

	return
}

func readMessageBytes(r *bufio.Reader, sz int, read func(*bufio.Reader, int, int) (int, error)) (int, error) {
	return readBytesWith(r, sz, func(r *bufio.Reader, sz int, len int32) (int, error) {
		return read(r, sz, int(len))
	})
}

func readMessage(r *bufio.Reader, sz int,
	key func(*bufio.Reader, int, int) (int, error),
	val func(*bufio.Reader, int, int) (int, error),
) (offset int64, attributes int8, timestamp int64, remain int, err error) {
	if offset, attributes, timestamp, remain, err = readMessageHeader(r, sz); err != nil {
		return
	}
	if remain, err = readMessageBytes(r, remain, key); err != nil {
		return
	}
	remain, err = readMessageBytes(r, remain, val)
	return
}

func readResponse(r *bufio.Reader, sz int, res interface{}) error {
	sz, err := read(r, sz, res)
	switch err.(type) {
	case Error:
		sz, err = discardN(r, sz, sz)
	}
	return expectZeroSize(sz, err)
}
<|endoftext|>


<commit_before>// RTLAMR - An rtl-sdr receiver for smart meters operating in the 900MHz ISM band.
// Copyright (C) 2014 Douglas Hall
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	"github.com/bemasher/rtlamr/csv"
	"github.com/bemasher/rtltcp"
)

const (
	CenterFreq = 920299072
	TimeFormat = "2006-01-02T15:04:05.000"
)

var rcvr Receiver

type Receiver struct {
	rtltcp.SDR
	d Decoder
	p Parser
}

func (rcvr *Receiver) NewReceiver() {
	switch strings.ToLower(*msgType) {
	case "scm":
		rcvr.d = NewDecoder(NewSCMPacketConfig(*symbolLength))
		rcvr.p = NewSCMParser()
	case "idm":
		rcvr.d = NewDecoder(NewIDMPacketConfig(*symbolLength))
		rcvr.p = NewIDMParser()
	default:
		log.Fatalf("Invalid message type: %q\n", *msgType)
	}

	if !*quiet {
		rcvr.d.cfg.Log()
		log.Println("CRC:", rcvr.p)
	}

	// Connect to rtl_tcp server.
	if err := rcvr.Connect(nil); err != nil {
		log.Fatal(err)
	}

	rcvr.HandleFlags()

	// Tell the user how many gain settings were reported by rtl_tcp.
	if !*quiet {
		log.Println("GainCount:", rcvr.SDR.Info.GainCount)
	}

	// Set some parameters for listening.
	rcvr.SetSampleRate(uint32(rcvr.d.cfg.SampleRate))
	rcvr.SetGainMode(true)

	return
}

func (rcvr *Receiver) Run() {
	// Setup signal channel for interruption.
	sigint := make(chan os.Signal, 1)
	signal.Notify(sigint, os.Kill, os.Interrupt)

	// Setup time limit channel
	tLimit := make(<-chan time.Time, 1)
	if *timeLimit != 0 {
		tLimit = time.After(*timeLimit)
	}

	block := make([]byte, rcvr.d.cfg.BlockSize2)

	start := time.Now()
	for {
		// Exit on interrupt or time limit, otherwise receive.
		select {
		case <-sigint:
			return
		case <-tLimit:
			fmt.Println("Time Limit Reached:", time.Since(start))
			return
		default:
			// Read new sample block.
			_, err := rcvr.Read(block)
			if err != nil {
				log.Fatal("Error reading samples: ", err)
			}

			pktFound := false
			for _, pkt := range rcvr.d.Decode(block) {
				scm, err := rcvr.p.Parse(NewDataFromBytes(pkt))
				if err != nil {
					// log.Println(err)
					continue
				}

				if *meterID != 0 && uint32(*meterID) != scm.MeterID() {
					continue
				}

				if *meterType != 0 && uint8(*meterType) != scm.MeterType() {
					continue
				}

				msg := NewLogMessage(scm)
				pktFound = true

				if encoder == nil {
					// A nil encoder is just plain-text output.
					fmt.Fprintln(logFile, msg)
				} else {
					err = encoder.Encode(msg)
					if err != nil {
						log.Fatal("Error encoding message: ", err)
					}

					// The XML encoder doesn't write new lines after each
					// element, add them.
					if strings.ToLower(*format) == "xml" {
						fmt.Fprintln(logFile)
					}
				}

				if *single {
					return
				}
			}

			if pktFound {
				_, err = sampleFile.Write(rcvr.d.iq)
				if err != nil {
					log.Fatal("Error writing raw samples to file:", err)
				}
			}
		}
	}
}

type Data struct {
	Bits  string
	Bytes []byte
}

func NewDataFromBytes(data []byte) (d Data) {
	d.Bytes = data
	for _, b := range data {
		d.Bits += fmt.Sprintf("%08b", b)
	}
	return
}

func NewDataFromBits(data string) (d Data) {
	d.Bits = data
	for idx := 0; idx < len(data); idx += 8 {
		b, _ := strconv.ParseUint(d.Bits[idx:idx+8], 2, 8)
		d.Bytes[idx>>3] = uint8(b)
	}
	return
}

type Parser interface {
	Parse(Data) (Message, error)
}

type Message interface {
	MsgType() string
	MeterID() uint32
	MeterType() uint8
	csv.Recorder
}

type LogMessage struct {
	Time   time.Time
	Offset int64
	Length int
	Body   Message
}

func NewLogMessage(body Message) (msg LogMessage) {
	msg.Time = time.Now()
	msg.Offset, _ = sampleFile.Seek(0, os.SEEK_CUR)
	msg.Length = rcvr.d.cfg.BufferLength << 1
	msg.Body = body

	return
}

func (msg LogMessage) String() string {
	if *sampleFilename == os.DevNull {
		return fmt.Sprintf("{Time:%s %s:%s}", msg.Time.Format(TimeFormat), msg.Body.MsgType(), msg.Body)
	}

	return fmt.Sprintf("{Time:%s Offset:%d Length:%d %s:%s}",
		msg.Time.Format(TimeFormat), msg.Offset, msg.Length, msg.Body.MsgType(), msg.Body,
	)
}

func (msg LogMessage) Record() (r []string) {
	r = append(r, msg.Time.Format(time.RFC3339Nano))
	r = append(r, strconv.FormatInt(msg.Offset, 10))
	r = append(r, strconv.FormatInt(int64(msg.Length), 10))
	r = append(r, msg.Body.Record()...)
	return r
}

func init() {
	log.SetFlags(log.Lshortfile | log.Lmicroseconds)
}

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to this file")

func main() {
	rcvr.RegisterFlags()
	RegisterFlags()

	flag.Parse()
	HandleFlags()

	rcvr.NewReceiver()

	defer logFile.Close()
	defer sampleFile.Close()
	defer rcvr.Close()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	rcvr.Run()
}
<commit_msg>Add duplicate packet detection, refactor LogMessage for clarity.<commit_after>// RTLAMR - An rtl-sdr receiver for smart meters operating in the 900MHz ISM band.
// Copyright (C) 2014 Douglas Hall
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"encoding/xml"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	"github.com/bemasher/rtlamr/csv"
	"github.com/bemasher/rtltcp"
)

const (
	CenterFreq = 920299072
	TimeFormat = "2006-01-02T15:04:05.000"
)

var rcvr Receiver

type Receiver struct {
	rtltcp.SDR
	d Decoder
	p Parser
}

func (rcvr *Receiver) NewReceiver() {
	switch strings.ToLower(*msgType) {
	case "scm":
		rcvr.d = NewDecoder(NewSCMPacketConfig(*symbolLength))
		rcvr.p = NewSCMParser()
	case "idm":
		rcvr.d = NewDecoder(NewIDMPacketConfig(*symbolLength))
		rcvr.p = NewIDMParser()
	default:
		log.Fatalf("Invalid message type: %q\n", *msgType)
	}

	if !*quiet {
		rcvr.d.cfg.Log()
		log.Println("CRC:", rcvr.p)
	}

	// Connect to rtl_tcp server.
	if err := rcvr.Connect(nil); err != nil {
		log.Fatal(err)
	}

	rcvr.HandleFlags()

	// Tell the user how many gain settings were reported by rtl_tcp.
	if !*quiet {
		log.Println("GainCount:", rcvr.SDR.Info.GainCount)
	}

	// Set some parameters for listening.
	rcvr.SetSampleRate(uint32(rcvr.d.cfg.SampleRate))
	rcvr.SetGainMode(true)

	return
}

func (rcvr *Receiver) Run() {
	// Setup signal channel for interruption.
	sigint := make(chan os.Signal, 1)
	signal.Notify(sigint, os.Kill, os.Interrupt)

	// Setup time limit channel
	tLimit := make(<-chan time.Time, 1)
	if *timeLimit != 0 {
		tLimit = time.After(*timeLimit)
	}

	block := make([]byte, rcvr.d.cfg.BlockSize2)

	start := time.Now()
	for {
		// Exit on interrupt or time limit, otherwise receive.
		select {
		case <-sigint:
			return
		case <-tLimit:
			fmt.Println("Time Limit Reached:", time.Since(start))
			return
		default:
			// Read new sample block.
			_, err := rcvr.Read(block)
			if err != nil {
				log.Fatal("Error reading samples: ", err)
			}

			seen := make(map[string]bool)
			for _, pkt := range rcvr.d.Decode(block) {
				data := NewDataFromBytes(pkt)

				if !seen[data.Bits] {
					seen[data.Bits] = true
				} else {
					continue
				}

				scm, err := rcvr.p.Parse(NewDataFromBytes(pkt))
				if err != nil {
					// log.Println(err)
					continue
				}

				if *meterID != 0 && uint32(*meterID) != scm.MeterID() {
					continue
				}

				if *meterType != 0 && uint8(*meterType) != scm.MeterType() {
					continue
				}

				msg := NewLogMessage(scm)

				if encoder == nil {
					// A nil encoder is just plain-text output.
					fmt.Fprintln(logFile, msg)
				} else {
					err = encoder.Encode(msg)
					if err != nil {
						log.Fatal("Error encoding message: ", err)
					}

					// The XML encoder doesn't write new lines after each
					// element, add them.
					if _, ok := encoder.(*xml.Encoder); ok {
						fmt.Fprintln(logFile)
					}
				}

				if *single {
					return
				}
			}

			if *sampleFilename != os.DevNull {
				_, err = sampleFile.Write(rcvr.d.iq)
				if err != nil {
					log.Fatal("Error writing raw samples to file:", err)
				}
			}
		}
	}
}

type Data struct {
	Bits  string
	Bytes []byte
}

func NewDataFromBytes(data []byte) (d Data) {
	d.Bytes = data
	for _, b := range data {
		d.Bits += fmt.Sprintf("%08b", b)
	}
	return
}

func NewDataFromBits(data string) (d Data) {
	d.Bits = data
	for idx := 0; idx < len(data); idx += 8 {
		b, _ := strconv.ParseUint(d.Bits[idx:idx+8], 2, 8)
		d.Bytes[idx>>3] = uint8(b)
	}
	return
}

type Parser interface {
	Parse(Data) (Message, error)
}

type Message interface {
	MsgType() string
	MeterID() uint32
	MeterType() uint8
	csv.Recorder
}

type LogMessage struct {
	Time   time.Time
	Offset int64
	Length int
	Message
}

func NewLogMessage(msg Message) (logMsg LogMessage) {
	logMsg.Time = time.Now()
	logMsg.Offset, _ = sampleFile.Seek(0, os.SEEK_CUR)
	logMsg.Length = rcvr.d.cfg.BufferLength << 1
	logMsg.Message = msg

	return
}

func (msg LogMessage) String() string {
	if *sampleFilename == os.DevNull {
		return fmt.Sprintf("{Time:%s %s:%s}", msg.Time.Format(TimeFormat), msg.MsgType(), msg.Message)
	}

	return fmt.Sprintf("{Time:%s Offset:%d Length:%d %s:%s}",
		msg.Time.Format(TimeFormat), msg.Offset, msg.Length, msg.MsgType(), msg.Message,
	)
}

func (msg LogMessage) Record() (r []string) {
	r = append(r, msg.Time.Format(time.RFC3339Nano))
	r = append(r, strconv.FormatInt(msg.Offset, 10))
	r = append(r, strconv.FormatInt(int64(msg.Length), 10))
	r = append(r, msg.Message.Record()...)
	return r
}

func init() {
	log.SetFlags(log.Lshortfile | log.Lmicroseconds)
}

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to this file")

func main() {
	rcvr.RegisterFlags()
	RegisterFlags()

	flag.Parse()
	HandleFlags()

	rcvr.NewReceiver()

	defer logFile.Close()
	defer sampleFile.Close()
	defer rcvr.Close()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	rcvr.Run()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
)

// Repo represents a repository of information yaml files.
type Repo struct {
	Key     string
	Info    map[string]Info
	Control map[string]Info
	root    string
	wg      sync.WaitGroup
}

func (r Repo) String() string {
	return fmt.Sprintf("R: %s (%d articles)", r.Key, len(r.Info))
}

// LoadRepos loads multiple repositories and stores them
func LoadRepos(p string) (repos map[string]Repo) {
	repos = make(map[string]Repo)
	p = getPath(p)
	wg := sync.WaitGroup{}

	files, _ := ioutil.ReadDir(p)
	for _, file := range files {
		fn := filepath.Join(p, file.Name())

		if _, err := os.Stat(filepath.Join(fn, "_repo.yml")); os.IsNotExist(err) {
			log.Println(fmt.Sprintf("Skipping repo %s: no _repo.yml found.", file.Name()))
			continue
		}

		wg.Add(1)
		go func(repos map[string]Repo, fn string, wg *sync.WaitGroup) {
			defer wg.Done()
			repos[asKey(fn)] = NewRepo(fn)
		}(repos, fn, &wg)
	}

	wg.Wait()
	return
}

// NewRepo loads a repository on a path
func NewRepo(p string) (r Repo) {
	r = Repo{Key: asKey(p), root: getPath(p)}
	r.Info = make(map[string]Info)
	r.Control = make(map[string]Info)

	filepath.Walk(r.root, r.walk)
	r.wg.Wait()

	return
}

// Keys returns a sorted list of the info keys in the repository
func (r *Repo) Keys() []string {
	keys := make([]string, 0, len(r.Info))
	for _, info := range r.Info {
		keys = append(keys, info.ID)
	}

	sort.Strings(keys)

	return keys
}

func (r *Repo) walk(path string, info os.FileInfo, err error) error {
	if err != nil {
		log.Println("walk error: ", err)
		return err
	}

	if strings.HasSuffix(path, ".yml") {
		r.wg.Add(1)
		go r.loadInfo(path)
	}

	return nil
}

func (r *Repo) loadInfo(path string) {
	defer r.wg.Done()

	info, err := LoadInfo(path)
	if err != nil {
		log.Println("Failed to load info: ", err)
	}

	// Control files start with an underscore and should not be stored as
	// normal Info documents.
	if strings.HasPrefix(asKey(path), "_") {
		r.Control[info.ID] = info
	} else {
		r.Info[info.ID] = info
	}
}
<commit_msg>Add git() helper to Repo{}<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"sync"
)

// Repo represents a repository of information yaml files.
type Repo struct {
	Key     string
	Info    map[string]Info
	Control map[string]Info
	root    string
	wg      sync.WaitGroup
}

func (r Repo) String() string {
	return fmt.Sprintf("R: %s (%d articles)", r.Key, len(r.Info))
}

// LoadRepos loads multiple repositories and stores them
func LoadRepos(p string) (repos map[string]Repo) {
	repos = make(map[string]Repo)
	p = getPath(p)
	wg := sync.WaitGroup{}

	files, _ := ioutil.ReadDir(p)
	for _, file := range files {
		fn := filepath.Join(p, file.Name())

		if _, err := os.Stat(filepath.Join(fn, "_repo.yml")); os.IsNotExist(err) {
			log.Println(fmt.Sprintf("Skipping repo %s: no _repo.yml found.", file.Name()))
			continue
		}

		wg.Add(1)
		go func(repos map[string]Repo, fn string, wg *sync.WaitGroup) {
			defer wg.Done()
			repos[asKey(fn)] = NewRepo(fn)
		}(repos, fn, &wg)
	}

	wg.Wait()
	return
}

// NewRepo loads a repository on a path
func NewRepo(p string) (r Repo) {
	r = Repo{Key: asKey(p), root: getPath(p)}
	r.Info = make(map[string]Info)
	r.Control = make(map[string]Info)

	filepath.Walk(r.root, r.walk)
	r.wg.Wait()

	return
}

// Keys returns a sorted list of the info keys in the repository
func (r *Repo) Keys() []string {
	keys := make([]string, 0, len(r.Info))
	for _, info := range r.Info {
		keys = append(keys, info.ID)
	}

	sort.Strings(keys)

	return keys
}

func (r *Repo) walk(path string, info os.FileInfo, err error) error {
	if err != nil {
		log.Println("walk error: ", err)
		return err
	}

	if strings.HasSuffix(path, ".yml") {
		r.wg.Add(1)
		go r.loadInfo(path)
	}

	return nil
}

func (r *Repo) loadInfo(path string) {
	defer r.wg.Done()

	info, err := LoadInfo(path)
	if err != nil {
		log.Println("Failed to load info: ", err)
	}

	// Control files start with an underscore and should not be stored as
	// normal Info documents.
	if strings.HasPrefix(asKey(path), "_") {
		r.Control[info.ID] = info
	} else {
		r.Info[info.ID] = info
	}
}

// Helper to run git commands inside of a repository
func (r *Repo) git(args ...string) {
	git, err := exec.LookPath("git")
	if err != nil {
		log.Fatal("no git :'(   ", err)
	}

	// why................
	args = append([]string{git}, args...)

	cmd := exec.Cmd{
		Path:   git,
		Args:   args,
		Dir:    r.root,
		Stdout: os.Stdout,
		Stderr: os.Stderr,
	}

	err = cmd.Run()
	if err != nil {
		log.Println("git command failed - aborting")
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"strings"

	"github.com/mikkeloscar/maze/repo"
)

// Repo defines a local copy of a remote arch repository.
type Repo struct {
	url      string
	dbCached bool
	local    repo.Repo
}

// Fetch db, downloads it first if needed.
func (r *Repo) fetchDB() error {
	if r.dbCached {
		return nil
	}

	fileName := fmt.Sprintf("%s.db.tar.gz", r.local.Name)

	if strings.HasPrefix(r.url, "http://") {
		// TODO: handle more db naming
		_, err := r.httpDownload(fileName)
		if err != nil {
			return err
		}

		r.dbCached = true
		return nil
	}

	// else local repo
	if r.url != r.local.Path() {
		err := copyFile(r.local.DB(), path.Join(r.url, fileName))
		if err != nil {
			return err
		}
	}

	r.dbCached = true
	return nil
}

// copy file
func copyFile(dst, src string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, in)
	if err != nil {
		return err
	}

	return nil
}

// download db file over http.
func (r *Repo) httpDownload(file string) (string, error) {
	filePath := path.Join(r.local.Path(), file)
	out, err := os.Create(filePath)
	if err != nil {
		return "", err
	}
	defer out.Close()

	// Get data
	resp, err := http.Get(r.url)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	// Write the data to file
	_, err = io.Copy(out, resp.Body)
	if err != nil {
		return "", err
	}

	return filePath, nil
}

// GetUpdated takes a list of package sources and returns a sorted list of the
// packages that need to be build (because the source is newer than what's in
// the repo).
func (r *Repo) GetUpdated(pkgs []*SrcPkg) ([]*SrcPkg, error) {
	err := r.fetchDB()
	if err != nil {
		return nil, err
	}

	updated := make([]*SrcPkg, 0, len(pkgs))
	for _, pkg := range pkgs {
		new, err := r.local.IsNew(pkg.PKGBUILD.Pkgbase, pkg.PKGBUILD.CompleteVersion())
		if err != nil {
			return nil, err
		}

		if new {
			updated = append(updated, pkg)
		}
	}

	return TopologicalSort(updated)
}
<commit_msg>Fix repo url<commit_after>package main

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"strings"

	"github.com/mikkeloscar/maze/repo"
)

// Repo defines a local copy of a remote arch repository.
type Repo struct {
	url      string
	dbCached bool
	local    repo.Repo
}

// Fetch db, downloads it first if needed.
func (r *Repo) fetchDB() error {
	if r.dbCached {
		return nil
	}

	fileName := fmt.Sprintf("%s.db.tar.gz", r.local.Name)

	if strings.HasPrefix(r.url, "http://") || strings.HasPrefix(r.url, "https://") {
		// TODO: handle more db naming
		_, err := r.httpDownload(fileName)
		if err != nil {
			return err
		}

		r.dbCached = true
		return nil
	}

	// else local repo
	if r.url != r.local.Path() {
		err := copyFile(r.local.DB(), path.Join(r.url, fileName))
		if err != nil {
			return err
		}
	}

	r.dbCached = true
	return nil
}

// copy file
func copyFile(dst, src string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, in)
	if err != nil {
		return err
	}

	return nil
}

// download db file over http.
func (r *Repo) httpDownload(file string) (string, error) {
	filePath := path.Join(r.local.Path(), file)
	out, err := os.Create(filePath)
	if err != nil {
		return "", err
	}
	defer out.Close()

	// Get data
	resp, err := http.Get(fmt.Sprintf("%s/%s", r.url, file))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	// Write the data to file
	_, err = io.Copy(out, resp.Body)
	if err != nil {
		return "", err
	}

	return filePath, nil
}

// GetUpdated takes a list of package sources and returns a sorted list of the
// packages that need to be build (because the source is newer than what's in
// the repo).
func (r *Repo) GetUpdated(pkgs []*SrcPkg) ([]*SrcPkg, error) {
	err := r.fetchDB()
	if err != nil {
		return nil, err
	}

	updated := make([]*SrcPkg, 0, len(pkgs))
	for _, pkg := range pkgs {
		new, err := r.local.IsNew(pkg.PKGBUILD.Pkgbase, pkg.PKGBUILD.CompleteVersion())
		if err != nil {
			return nil, err
		}

		if new {
			updated = append(updated, pkg)
		}
	}

	return TopologicalSort(updated)
}
<|endoftext|>


<commit_before>package fluidbackup

import "testing"
import "time"
import "io/ioutil"
import "math/rand"
import "os"
import "bytes"
import "strings"
import "fmt"

func TestBasic(t *testing.T) {
	fileContents := make([]byte, 1500*1024)
	for i := range fileContents {
		fileContents[i] = byte(rand.Int() % 256)
	}
	err := ioutil.WriteFile("orange.txt", fileContents, 0644)
	if err != nil {
		t.Error("failed to write random contents to orange.txt: " + err.Error())
	}

	mainPort := 19838
	bMain := MakeFluidBackup(mainPort)
	bMain.fileStore.RegisterFile("orange.txt")
	bOther := make([]*FluidBackup, 11)
	for i := range bOther {
		port := 19839 + i
		bOther[i] = MakeFluidBackup(port)

		if i < 5 {
			bMain.peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: port})
		} else {
			bOther[i].peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: mainPort})
		}
	}

	time.Sleep(5 * time.Second)

	err = os.Remove("orange.txt")

	if err != nil {
		t.Error("Encountered error while removing file: " + err.Error())
		return
	}

	bMain.fileStore.RecoverFile("orange.txt")
	newFileContents, err := ioutil.ReadFile("orange.txt")
	if err != nil {
		t.Error("failed to read from recovered orange.txt: " + err.Error())
	}

	if !bytes.Equal(fileContents, newFileContents) {
		t.Error("recovered file does not match original")
	}

	// shut down fluidbackup instances
	bMain.Stop()
	for _, inst := range bOther {
		inst.Stop()
	}
}

// Attempts to kill the local peer
// and ensure that the file can be recovered.
func TestStable(t *testing.T) {
	fileContents := make([]byte, 1500*1024)
	for i := range fileContents {
		fileContents[i] = byte(rand.Int() % 256)
	}
	err := ioutil.WriteFile("orange.txt", fileContents, 0644)
	if err != nil {
		t.Error("failed to write random contents to orange.txt: " + err.Error())
	}

	mainPort := 19838
	bMain := MakeFluidBackup(mainPort)
	bMain.fileStore.RegisterFile("orange.txt")
	bOther := make([]*FluidBackup, 11)
	for i := range bOther {
		port := 19839 + i
		bOther[i] = MakeFluidBackup(port)

		if i < 5 {
			bMain.peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: port})
		} else {
			bOther[i].peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: mainPort})
		}
	}

	time.Sleep(5 * time.Second)

	err = os.Remove("orange.txt")

	if err != nil {
		t.Error("Encountered error while removing file: " + err.Error())
		return
	}

	bMain.Save()
	bMain.Stop()
	time.Sleep(5 * time.Second)
	bMain = MakeFluidBackup(mainPort)
	bMain.Load()

	bMain.fileStore.RecoverFile("orange.txt")
	newFileContents, err := ioutil.ReadFile("orange.txt")
	if err != nil {
		t.Error("failed to read from recovered orange.txt: " + err.Error())
	}

	if !bytes.Equal(fileContents, newFileContents) {
		t.Error("recovered file does not match original")
	}

	// shut down fluidbackup instances
	bMain.Stop()
	for _, inst := range bOther {
		inst.Stop()
	}
}

func TestReread(t *testing.T) {
	fileContents := make([]byte, 1500*1024)
	for i := range fileContents {
		fileContents[i] = byte(rand.Int() % 256)
	}
	err := ioutil.WriteFile("orange.txt", fileContents, 0644)
	if err != nil {
		t.Error("failed to write random contents to orange.txt: " + err.Error())
	}

	mainPort := 19838
	bMain := MakeFluidBackup(mainPort)
	bMain.fileStore.RegisterFile("orange.txt")
	time.Sleep(time.Second)
	bMain.Save()
	bMain.Stop()
	time.Sleep(2 * time.Second)
	bMain = MakeFluidBackup(mainPort)
	bMain.Load()

	bOther := make([]*FluidBackup, 11)
	for i := range bOther {
		port := 19839 + i
		bOther[i] = MakeFluidBackup(port)

		if i < 5 {
			bMain.peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: port})
		} else {
			bOther[i].peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: mainPort})
		}
	}

	time.Sleep(5 * time.Second)

	err = os.Remove("orange.txt")

	if err != nil {
		t.Error("Encountered error while removing file: " + err.Error())
		return
	}

	bMain.fileStore.RecoverFile("orange.txt")
	newFileContents, err := ioutil.ReadFile("orange.txt")
	if err != nil {
		t.Error("failed to read from recovered orange.txt: " + err.Error())
	}

	if !bytes.Equal(fileContents, newFileContents) {
		t.Error("recovered file does not match original")
	}

	// shut down fluidbackup instances
	bMain.Stop()
	for _, inst := range bOther {
		inst.Stop()
	}
}

func TestFailVerify(t *testing.T) {
	fileContents := make([]byte, 1500*1024)
	for i := range fileContents {
		fileContents[i] = byte(rand.Int() % 256)
	}
	err := ioutil.WriteFile("orange.txt", fileContents, 0644)
	if err != nil {
		t.Error("failed to write random contents to orange.txt: " + err.Error())
	}

	mainPort := 19838
	bMain := MakeFluidBackup(mainPort)
	bMain.fileStore.RegisterFile("orange.txt")
	bOther := make([]*FluidBackup, 11)
	for i := range bOther {
		port := 19839 + i
		bOther[i] = MakeFluidBackup(port)

		if i < 5 {
			bMain.peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: port})
		} else {
			bOther[i].peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: mainPort})
		}
	}

	time.Sleep(5 * time.Second)
	files, _ := ioutil.ReadDir("store/")
	for _, f := range files {
		if strings.HasSuffix(f.Name(), ".obj") {
			os.Remove("store/" + f.Name())
		}
	}
	time.Sleep(25 * time.Second)
	err = os.Remove("orange.txt")

	if err != nil {
		t.Error("Encountered error while removing file: " + err.Error())
		return
	}

	bMain.fileStore.RecoverFile("orange.txt")
	newFileContents, err := ioutil.ReadFile("orange.txt")
	if err != nil {
		t.Error("failed to read from recovered orange.txt: " + err.Error())
	}

	if !bytes.Equal(fileContents, newFileContents) {
		t.Error("recovered file does not match original")
	}

	// shut down fluidbackup instances
	bMain.Stop()
	for _, inst := range bOther {
		inst.Stop()
	}
}

func TestBasicPeerSharing(t *testing.T) {
	// todo: need to implement:
	// - trust ranking
	// - how many peers do we want?
	mainPort := 19838
	bMain := MakeFluidBackup(mainPort)

	// Make a giant centroid peer network first
	bOther := make([]*FluidBackup, 11)
	for i := range bOther {
		port := 19839 + i
		bOther[i] = MakeFluidBackup(port)
		if i < 5 {
			bMain.peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: port})
		} else {
			bOther[i].peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: mainPort})
		}
	}

	time.Sleep(time.Second)

	// Now, connect anoter peer to the centroid
	newDiscoverer := MakeFluidBackup(mainPort - 1)
	newDiscoverer.peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: mainPort})
	// No need for next line: they are now autodiscovered
	// numPeers := newDiscoverer.peerList.FindNewPeers(len(bOther))

	time.Sleep(time.Second * 5)

	// check the number of discovered peers for one on centroid
	numPeers := newDiscoverer.peerList.NumPeers()
	desiredNum := PeerListDefaultDesiredNumPeers
	fmt.Printf("Discoverer Found %v peers!\n", numPeers)
	if numPeers < desiredNum {
		t.Fatalf("Num peers %v, expected at least %v", numPeers, desiredNum)
	}

	// check the number of discovered peers for each bOther
	for i := range bOther {
		fbOther := bOther[i]
		numPeers := fbOther.peerList.NumPeers()
		fmt.Printf("%v Found %v peers!\n", fbOther.protocol.port, numPeers)
		if numPeers < desiredNum {
			t.Fatalf("%v: Num peers %v, expected at least %v", fbOther.protocol.port, numPeers, desiredNum)
		}
	}

	// shut down fluidbackup instances
	bMain.Stop()
	newDiscoverer.Stop()
	for _, inst := range bOther {
		inst.Stop()
	}
}
<commit_msg>A very little bit of clean-up;<commit_after>package fluidbackup

import "testing"
import "time"
import "io/ioutil"
import "math/rand"
import "os"
import "bytes"
import "strings"
import "fmt"

func TestBasic(t *testing.T) {
	fileContents := make([]byte, 1500*1024)
	for i := range fileContents {
		fileContents[i] = byte(rand.Int() % 256)
	}
	err := ioutil.WriteFile("orange.txt", fileContents, 0644)
	if err != nil {
		t.Error("failed to write random contents to orange.txt: " + err.Error())
	}

	mainPort := 19838
	bMain := MakeFluidBackup(mainPort)
	bMain.fileStore.RegisterFile("orange.txt")
	bOther := make([]*FluidBackup, 11)
	for i := range bOther {
		port := 19839 + i
		bOther[i] = MakeFluidBackup(port)

		if i < 5 {
			bMain.peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: port})
		} else {
			bOther[i].peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: mainPort})
		}
	}

	time.Sleep(5 * time.Second)

	err = os.Remove("orange.txt")

	if err != nil {
		t.Error("Encountered error while removing file: " + err.Error())
		return
	}

	bMain.fileStore.RecoverFile("orange.txt")
	newFileContents, err := ioutil.ReadFile("orange.txt")
	if err != nil {
		t.Error("failed to read from recovered orange.txt: " + err.Error())
	}

	if !bytes.Equal(fileContents, newFileContents) {
		t.Error("recovered file does not match original")
	}

	// shut down fluidbackup instances
	bMain.Stop()
	for _, inst := range bOther {
		inst.Stop()
	}
}

// Attempts to kill the local peer
// and ensure that the file can be recovered.
func TestStable(t *testing.T) {
	fileContents := make([]byte, 1500*1024)
	for i := range fileContents {
		fileContents[i] = byte(rand.Int() % 256)
	}
	err := ioutil.WriteFile("orange.txt", fileContents, 0644)
	if err != nil {
		t.Error("failed to write random contents to orange.txt: " + err.Error())
	}

	mainPort := 19838
	bMain := MakeFluidBackup(mainPort)
	bMain.fileStore.RegisterFile("orange.txt")
	bOther := make([]*FluidBackup, 11)
	for i := range bOther {
		port := 19839 + i
		bOther[i] = MakeFluidBackup(port)

		if i < 5 {
			bMain.peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: port})
		} else {
			bOther[i].peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: mainPort})
		}
	}

	time.Sleep(5 * time.Second)

	err = os.Remove("orange.txt")

	if err != nil {
		t.Error("Encountered error while removing file: " + err.Error())
		return
	}

	bMain.Save()
	bMain.Stop()
	time.Sleep(5 * time.Second)
	bMain = MakeFluidBackup(mainPort)
	bMain.Load()

	bMain.fileStore.RecoverFile("orange.txt")
	newFileContents, err := ioutil.ReadFile("orange.txt")
	if err != nil {
		t.Error("failed to read from recovered orange.txt: " + err.Error())
	}

	if !bytes.Equal(fileContents, newFileContents) {
		t.Error("recovered file does not match original")
	}

	// shut down fluidbackup instances
	bMain.Stop()
	for _, inst := range bOther {
		inst.Stop()
	}
}

func TestReread(t *testing.T) {
	fileContents := make([]byte, 1500*1024)
	for i := range fileContents {
		fileContents[i] = byte(rand.Int() % 256)
	}
	err := ioutil.WriteFile("orange.txt", fileContents, 0644)
	if err != nil {
		t.Error("failed to write random contents to orange.txt: " + err.Error())
	}

	mainPort := 19838
	bMain := MakeFluidBackup(mainPort)
	bMain.fileStore.RegisterFile("orange.txt")
	time.Sleep(time.Second)
	bMain.Save()
	bMain.Stop()
	time.Sleep(2 * time.Second)
	bMain = MakeFluidBackup(mainPort)
	bMain.Load()

	bOther := make([]*FluidBackup, 11)
	for i := range bOther {
		port := 19839 + i
		bOther[i] = MakeFluidBackup(port)

		if i < 5 {
			bMain.peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: port})
		} else {
			bOther[i].peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: mainPort})
		}
	}

	time.Sleep(5 * time.Second)

	err = os.Remove("orange.txt")

	if err != nil {
		t.Error("Encountered error while removing file: " + err.Error())
		return
	}

	bMain.fileStore.RecoverFile("orange.txt")
	newFileContents, err := ioutil.ReadFile("orange.txt")
	if err != nil {
		t.Error("failed to read from recovered orange.txt: " + err.Error())
	}

	if !bytes.Equal(fileContents, newFileContents) {
		t.Error("recovered file does not match original")
	}

	// shut down fluidbackup instances
	bMain.Stop()
	for _, inst := range bOther {
		inst.Stop()
	}
}

func TestFailVerify(t *testing.T) {
	fileContents := make([]byte, 1500*1024)
	for i := range fileContents {
		fileContents[i] = byte(rand.Int() % 256)
	}
	err := ioutil.WriteFile("orange.txt", fileContents, 0644)
	if err != nil {
		t.Error("failed to write random contents to orange.txt: " + err.Error())
	}

	mainPort := 19838
	bMain := MakeFluidBackup(mainPort)
	bMain.fileStore.RegisterFile("orange.txt")
	bOther := make([]*FluidBackup, 11)
	for i := range bOther {
		port := 19839 + i
		bOther[i] = MakeFluidBackup(port)

		if i < 5 {
			bMain.peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: port})
		} else {
			bOther[i].peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: mainPort})
		}
	}

	time.Sleep(5 * time.Second)
	files, _ := ioutil.ReadDir("store/")
	for _, f := range files {
		if strings.HasSuffix(f.Name(), ".obj") {
			os.Remove("store/" + f.Name())
		}
	}
	time.Sleep(25 * time.Second)
	err = os.Remove("orange.txt")

	if err != nil {
		t.Error("Encountered error while removing file: " + err.Error())
		return
	}

	bMain.fileStore.RecoverFile("orange.txt")
	newFileContents, err := ioutil.ReadFile("orange.txt")
	if err != nil {
		t.Error("failed to read from recovered orange.txt: " + err.Error())
	}

	if !bytes.Equal(fileContents, newFileContents) {
		t.Error("recovered file does not match original")
	}

	// shut down fluidbackup instances
	bMain.Stop()
	for _, inst := range bOther {
		inst.Stop()
	}
}

func TestBasicPeerSharing(t *testing.T) {
	mainPort := 19838
	bMain := MakeFluidBackup(mainPort)

	// Make a giant centroid peer network first
	bOther := make([]*FluidBackup, 11)
	for i := range bOther {
		port := 19839 + i
		bOther[i] = MakeFluidBackup(port)
		if i < 5 {
			bMain.peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: port})
		} else {
			bOther[i].peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: mainPort})
		}
	}

	time.Sleep(time.Second)

	// Now, connect anoter peer to the centroid
	newDiscoverer := MakeFluidBackup(mainPort - 1)
	newDiscoverer.peerList.DiscoveredPeer(PeerId{Address: "127.0.0.1", Port: mainPort})
	// No need for next line: they are now autodiscovered
	// numPeers := newDiscoverer.peerList.FindNewPeers(len(bOther))

	time.Sleep(time.Second * 5)

	// check the number of discovered peers for one on centroid
	numPeers := newDiscoverer.peerList.NumPeers()
	desiredNum := PeerListDefaultDesiredNumPeers
	fmt.Printf("Discoverer Found %v peers!\n", numPeers)
	if numPeers < desiredNum {
		t.Fatalf("Num peers %v, expected at least %v", numPeers, desiredNum)
	}

	// check the number of discovered peers for each bOther
	for i := range bOther {
		fbOther := bOther[i]
		numPeers := fbOther.peerList.NumPeers()
		fmt.Printf("%v Found %v peers!\n", fbOther.protocol.port, numPeers)
		if numPeers < desiredNum {
			t.Fatalf("%v: Num peers %v, expected at least %v", fbOther.protocol.port, numPeers, desiredNum)
		}
	}

	// shut down fluidbackup instances
	bMain.Stop()
	newDiscoverer.Stop()
	for _, inst := range bOther {
		inst.Stop()
	}
}

func TestBenchmark(t *testing.T) {

}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"gofire/command"
	"gofire/message"
	"gofire/user"
	//"regexp"
	//"strings"
	"code.google.com/p/go.net/websocket"
)

type Connection struct {
	//User who sends
	Usr *user.User
	//websocket-connection
	Conn *websocket.Conn
	// Buffered channel of outbound messages.
	send chan *command.Command
	// chatroom in which the user is
	chatRoom *ChatRoom
}

func (c *Connection) Read() {
	for {
		//raw incoming command
		var rawIncome []byte

		err := websocket.Message.Receive(c.Conn, &rawIncome)
		if err != nil {
			fmt.Println(err)
			break
		}

		var cmd *command.Command

		//var userName string
		var found bool

		err = json.Unmarshal([]byte(rawIncome), &cmd)

		if err == nil {
			//REGISTER -> {REGISTER, USERNAME}
			if cmd.Type == command.REGISTER {
				//check if user name already is taken.
				for d := range server.registeredConnections {
					if d.Usr != nil && d.Usr.Name == string(cmd.Value) {
						found = true
					}
				}
				if found {
					break
				} else {
					c.Usr = &user.User{Name: string(cmd.Value)}
				}
			}

			if cmd.Type == command.BMESSAGE {
				m, errm := json.Marshal(message.Message{c.Usr, cmd.Value})
				if errm == nil {
					c.chatRoom.broadcast <- &command.Command{command.BMESSAGE, m}
				} else {
					fmt.Println(errm)
				}
			}

			/*if cmd.Type == command.MESSAGE {
			//check if the message is a private message.
			r, _ := regexp.MatchString("@", string(cmd.Value))
			if r {
				found = false
				userName = strings.Split(string(cmd.Value), " ")[0]
				userName = strings.Split(userName, "@")[1]
				for d := range server.registeredConnections {
					if d.Usr.Name == userName {
						d.send <- &Message{c.Usr, cmd.Value}
						found = true
					}
				}
				if found {
					c.send <- &Message{c.Usr, cmd.Value}
				} else {
					c.send <- &Message{&User{"From server"}, []byte("user not found")}
				}
				//else send the message to everyone
			} else {
				c.chatRoom.broadcast <- &Message{c.Usr, cmd.Value}
			}*/
		}

		if cmd.Type == command.BLOGIN {
			cmd,err := command.PrepareMessage(command.BMESSAGE, c.Usr, []byte("Logged In"))
			if err == nil {
				c.chatRoom.broadcast <- cmd
			} else {
				fmt.Println(err)
			}

		}

		if cmd.Type == command.BLOGOUT {
			cmd,err := command.PrepareMessage(command.BMESSAGE, c.Usr, []byte("Logged out"))
			if err == nil {
				c.chatRoom.broadcast <- cmd
			} else {
				fmt.Println(err)
			}

		}
	}

	c.Conn.Close()
}

func (c *Connection) Write() {
	for command := range c.send {
		//marshal in

		jsonC, _ := json.Marshal(command)

		err := websocket.Message.Send(c.Conn, string(jsonC))
		if err != nil {
			break
		}
	}
}
<commit_msg>print user<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"gofire/command"
	"gofire/message"
	"gofire/user"
	//"regexp"
	//"strings"
	"code.google.com/p/go.net/websocket"
)

type Connection struct {
	//User who sends
	Usr *user.User
	//websocket-connection
	Conn *websocket.Conn
	// Buffered channel of outbound messages.
	send chan *command.Command
	// chatroom in which the user is
	chatRoom *ChatRoom
}

func (c *Connection) Read() {
	for {
		//raw incoming command
		var rawIncome []byte

		err := websocket.Message.Receive(c.Conn, &rawIncome)
		if err != nil {
			fmt.Println(err)
			break
		}

		var cmd *command.Command

		//var userName string
		var found bool

		err = json.Unmarshal([]byte(rawIncome), &cmd)

		if err == nil {
			//REGISTER -> {REGISTER, USERNAME}
			if cmd.Type == command.REGISTER {
				//check if user name already is taken.
				for d := range server.registeredConnections {
					if d.Usr != nil && d.Usr.Name == string(cmd.Value) {
						found = true
					}
				}
				if found {
					break
				} else {
					c.Usr = &user.User{Name: string(cmd.Value)}
				}
			}

			if cmd.Type == command.BMESSAGE {
				m, errm := json.Marshal(message.Message{c.Usr, cmd.Value})
				if errm == nil {
					c.chatRoom.broadcast <- &command.Command{command.BMESSAGE, m}
				} else {
					fmt.Println(errm)
				}
			}

			/*if cmd.Type == command.MESSAGE {
			//check if the message is a private message.
			r, _ := regexp.MatchString("@", string(cmd.Value))
			if r {
				found = false
				userName = strings.Split(string(cmd.Value), " ")[0]
				userName = strings.Split(userName, "@")[1]
				for d := range server.registeredConnections {
					if d.Usr.Name == userName {
						d.send <- &Message{c.Usr, cmd.Value}
						found = true
					}
				}
				if found {
					c.send <- &Message{c.Usr, cmd.Value}
				} else {
					c.send <- &Message{&User{"From server"}, []byte("user not found")}
				}
				//else send the message to everyone
			} else {
				c.chatRoom.broadcast <- &Message{c.Usr, cmd.Value}
			}*/
		}

		if cmd.Type == command.BLOGIN {
			fmt.Println(c.Usr)
			cmd,err := command.PrepareMessage(command.BMESSAGE, c.Usr, []byte("Logged In"))
			if err == nil {
				c.chatRoom.broadcast <- cmd
			} else {
				fmt.Println(err)
			}

		}

		if cmd.Type == command.BLOGOUT {
			cmd,err := command.PrepareMessage(command.BMESSAGE, c.Usr, []byte("Logged out"))
			if err == nil {
				c.chatRoom.broadcast <- cmd
			} else {
				fmt.Println(err)
			}

		}
	}

	c.Conn.Close()
}

func (c *Connection) Write() {
	for command := range c.send {
		//marshal in

		jsonC, _ := json.Marshal(command)

		err := websocket.Message.Send(c.Conn, string(jsonC))
		if err != nil {
			break
		}
	}
}
<|endoftext|>


<commit_before>/*
   Hockeypuck - OpenPGP key server
   Copyright (C) 2012-2014  Casey Marshall

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published by
   the Free Software Foundation, version 3.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

package hkp

import (
	"bytes"
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"net/url"
	"path/filepath"
	"strings"
	"time"

	"github.com/julienschmidt/httprouter"
	"golang.org/x/crypto/openpgp/armor"
	"gopkg.in/errgo.v1"

	"hockeypuck/conflux/recon"
	"hockeypuck/hkp/sks"
	"hockeypuck/hkp/storage"
	log "hockeypuck/logrus"
	"hockeypuck/openpgp"
)

const (
	shortKeyIDLen       = 8
	longKeyIDLen        = 16
	fingerprintKeyIDLen = 40
)

var errKeywordSearchNotAvailable = errgo.New("keyword search is not available")

func httpError(w http.ResponseWriter, statusCode int, err error) {
	if statusCode != http.StatusNotFound {
		log.Errorf("HTTP %d: %v", statusCode, errgo.Details(err))
	}
	http.Error(w, http.StatusText(statusCode), statusCode)
}

type Handler struct {
	storage storage.Storage

	indexWriter  IndexFormat
	vindexWriter IndexFormat

	statsTemplate *template.Template
	statsFunc     func() (interface{}, error)

	selfSignedOnly  bool
	fingerprintOnly bool
}

type HandlerOption func(h *Handler) error

func IndexTemplate(path string, extra ...string) HandlerOption {
	return func(h *Handler) error {
		tw, err := NewHTMLFormat(path, extra)
		if err != nil {
			return errgo.Mask(err)
		}
		h.indexWriter = tw
		return nil
	}
}

func VIndexTemplate(path string, extra ...string) HandlerOption {
	return func(h *Handler) error {
		tw, err := NewHTMLFormat(path, extra)
		if err != nil {
			return errgo.Mask(err)
		}
		h.vindexWriter = tw
		return nil
	}
}

func StatsTemplate(path string, extra ...string) HandlerOption {
	return func(h *Handler) error {
		t := template.New(filepath.Base(path)).Funcs(template.FuncMap{
			"url": func(u *url.URL) template.URL {
				return template.URL(u.String())
			},
			"day": func(t time.Time) string {
				return t.Format("2006-01-02")
			},
			"hour": func(t time.Time) string {
				return t.Format("2006-01-02 15")
			},
		})
		var err error
		if len(extra) > 0 {
			t, err = t.ParseFiles(append([]string{path}, extra...)...)
		} else {
			t, err = t.ParseGlob(path)
		}
		if err != nil {
			return errgo.Mask(err)
		}
		h.statsTemplate = t
		return nil
	}
}

func StatsFunc(f func() (interface{}, error)) HandlerOption {
	return func(h *Handler) error {
		h.statsFunc = f
		return nil
	}
}

func SelfSignedOnly(selfSignedOnly bool) HandlerOption {
	return func(h *Handler) error {
		h.selfSignedOnly = selfSignedOnly
		return nil
	}
}

func FingerprintOnly(fingerprintOnly bool) HandlerOption {
	return func(h *Handler) error {
		h.fingerprintOnly = fingerprintOnly
		return nil
	}
}

func NewHandler(storage storage.Storage, options ...HandlerOption) (*Handler, error) {
	h := &Handler{
		storage: storage,
	}
	for _, option := range options {
		err := option(h)
		if err != nil {
			return nil, errgo.Mask(err)
		}
	}
	return h, nil
}

func (h *Handler) Register(r *httprouter.Router) {
	r.GET("/pks/lookup", h.Lookup)
	r.POST("/pks/add", h.Add)
	r.POST("/pks/hashquery", h.HashQuery)
}

func (h *Handler) Lookup(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	l, err := ParseLookup(r)
	if err != nil {
		httpError(w, http.StatusBadRequest, err)
		return
	}
	switch l.Op {
	case OperationGet, OperationHGet:
		h.get(w, l)
	case OperationIndex:
		h.index(w, l, h.indexWriter)
	case OperationVIndex:
		h.index(w, l, h.vindexWriter)
	case OperationStats:
		h.stats(w, l)
	default:
		httpError(w, http.StatusNotFound, errgo.Newf("operation not found: %v", l.Op))
		return
	}
}

func (h *Handler) HashQuery(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	hq, err := ParseHashQuery(r)
	if err != nil {
		httpError(w, http.StatusBadRequest, errgo.Mask(err))
		return
	}
	var result []*openpgp.PrimaryKey
	for _, digest := range hq.Digests {
		rfps, err := h.storage.MatchMD5([]string{digest})
		if err != nil {
			log.Errorf("error resolving hashquery digest %q", digest)
			continue
		}
		keys, err := h.storage.FetchKeys(rfps)
		if err != nil {
			log.Errorf("error fetching hashquery key %q", digest)
			continue
		}
		result = append(result, keys...)
	}

	w.Header().Set("Content-Type", "pgp/keys")

	// Write the number of keys
	err = recon.WriteInt(w, len(result))
	for _, key := range result {
		// Write each key in binary packet format, prefixed with length
		err = writeHashqueryKey(w, key)
		if err != nil {
			log.Errorf("error writing hashquery key %q: %v", key.RFingerprint, err)
			return
		}
	}

	// SKS expects hashquery response to terminate with a CRLF
	_, err = w.Write([]byte{0x0d, 0x0a})
	if err != nil {
		log.Errorf("error writing hashquery terminator: %v", err)
	}
}

func writeHashqueryKey(w http.ResponseWriter, key *openpgp.PrimaryKey) error {
	var buf bytes.Buffer
	err := openpgp.WritePackets(&buf, key)
	if err != nil {
		return errgo.Mask(err)
	}
	err = recon.WriteInt(w, buf.Len())
	if err != nil {
		return errgo.Mask(err)
	}
	_, err = w.Write(buf.Bytes())
	if err != nil {
		return errgo.Mask(err)
	}
	return nil
}

func (h *Handler) resolve(l *Lookup) ([]string, error) {
	if l.Op == OperationHGet {
		return h.storage.MatchMD5([]string{l.Search})
	}
	if strings.HasPrefix(l.Search, "0x") {
		keyID := openpgp.Reverse(strings.ToLower(l.Search[2:]))
		switch len(keyID) {
		case shortKeyIDLen, longKeyIDLen, fingerprintKeyIDLen:
			return h.storage.Resolve([]string{keyID})
		}
	}
	if h.fingerprintOnly {
		return nil, errKeywordSearchNotAvailable
	}
	return h.storage.MatchKeyword([]string{l.Search})
}

func (h *Handler) keys(l *Lookup) ([]*openpgp.PrimaryKey, error) {
	rfps, err := h.resolve(l)
	if err != nil {
		return nil, err
	}
	keys, err := h.storage.FetchKeys(rfps)
	if err != nil {
		return nil, errgo.Mask(err)
	}
	if h.selfSignedOnly {
		for _, key := range keys {
			if err := openpgp.SelfSignedOnly(key); err != nil {
				return nil, errgo.Mask(err)
			}
		}
	}
	return keys, nil
}

func (h *Handler) get(w http.ResponseWriter, l *Lookup) {
	keys, err := h.keys(l)
	if err == errKeywordSearchNotAvailable {
		httpError(w, http.StatusBadRequest, errgo.Mask(err))
		return
	} else if err != nil {
		httpError(w, http.StatusInternalServerError, errgo.Mask(err))
		return
	}
	if len(keys) == 0 {
		httpError(w, http.StatusNotFound, errgo.New("not found"))
		return
	}

	// Drop malformed packets, since these break GPG imports.
	for _, key := range keys {
		var others []*openpgp.Packet
		for _, other := range key.Others {
			if other.Malformed {
				continue
			}
			others = append(others, other)
		}
		key.Others = others
	}

	w.Header().Set("Content-Type", "text/plain")
	err = openpgp.WriteArmoredPackets(w, keys)
	if err != nil {
		log.Errorf("get %q: error writing armored keys: %v", l.Search, err)
	}
	// Write a trailing newline as required by the HKP spec
	// (§3.1.2.1) and as expected by many tools, e.g. RPM.
	_, err = w.Write([]byte("\n"))
	if err != nil {
		log.Errorf("get %q: failed to write trailing newline: %v", l.Search, err)
	}
}

func (h *Handler) index(w http.ResponseWriter, l *Lookup, f IndexFormat) {
	keys, err := h.keys(l)
	if err != nil {
		httpError(w, http.StatusInternalServerError, errgo.Mask(err))
		return
	}
	if len(keys) == 0 {
		httpError(w, http.StatusNotFound, errgo.New("not found"))
		return
	}

	if l.Options[OptionMachineReadable] {
		f = mrFormat
	} else if l.Options[OptionJSON] || f == nil {
		f = jsonFormat
	}

	err = f.Write(w, l, keys)
	if err != nil {
		httpError(w, http.StatusInternalServerError, errgo.Mask(err))
		return
	}
}

func (h *Handler) indexJSON(w http.ResponseWriter, keys []*openpgp.PrimaryKey) {
	w.Header().Set("Content-Type", "application/json")
	enc := json.NewEncoder(w)
	err := enc.Encode(&keys)
	if err != nil {
		httpError(w, http.StatusInternalServerError, errgo.Mask(err))
		return
	}
}

func mrTimeString(t time.Time) string {
	if t.IsZero() {
		return ""
	}
	return fmt.Sprintf("%d", t.Unix())
}

type StatsResponse struct {
	Info  interface{}
	Stats *sks.Stats
}

func (h *Handler) stats(w http.ResponseWriter, l *Lookup) {
	if h.statsFunc == nil {
		httpError(w, http.StatusBadRequest, errgo.New("stats not configured"))
		fmt.Fprintln(w, "stats not configured")
		return
	}
	data, err := h.statsFunc()
	if err != nil {
		httpError(w, http.StatusInternalServerError, errgo.Mask(err))
		return
	}

	if h.statsTemplate != nil && !(l.Options[OptionJSON] || l.Options[OptionMachineReadable]) {
		err = h.statsTemplate.Execute(w, data)
	} else {
		err = json.NewEncoder(w).Encode(data)
	}
	if err != nil {
		httpError(w, http.StatusInternalServerError, errgo.Mask(err))
	}
}

type AddResponse struct {
	Inserted []string `json:"inserted"`
	Updated  []string `json:"updated"`
	Ignored  []string `json:"ignored"`
}

func (h *Handler) Add(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	add, err := ParseAdd(r)
	if err != nil {
		httpError(w, http.StatusBadRequest, errgo.Mask(err))
		return
	}

	// Check and decode the armor
	armorBlock, err := armor.Decode(bytes.NewBufferString(add.Keytext))
	if err != nil {
		httpError(w, http.StatusBadRequest, errgo.Mask(err))
		return
	}

	var result AddResponse
	for readKey := range openpgp.ReadKeys(armorBlock.Body) {
		if readKey.Error != nil {
			httpError(w, http.StatusBadRequest, errgo.Mask(err))
			return
		}
		err := openpgp.DropDuplicates(readKey.PrimaryKey)
		if err != nil {
			httpError(w, http.StatusInternalServerError, errgo.Mask(err))
			return
		}
		change, err := storage.UpsertKey(h.storage, readKey.PrimaryKey)
		if err != nil {
			httpError(w, http.StatusInternalServerError, errgo.Mask(err))
			return
		}

		fp := readKey.PrimaryKey.QualifiedFingerprint()
		switch change.(type) {
		case storage.KeyAdded:
			result.Inserted = append(result.Inserted, fp)
		case storage.KeyReplaced:
			result.Updated = append(result.Updated, fp)
		case storage.KeyNotChanged:
			result.Ignored = append(result.Ignored, fp)
		}
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	enc := json.NewEncoder(w)
	enc.Encode(&result)
}
<commit_msg>Fix 400 response for op=index and keyword search disabled.<commit_after>/*
   Hockeypuck - OpenPGP key server
   Copyright (C) 2012-2014  Casey Marshall

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published by
   the Free Software Foundation, version 3.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

package hkp

import (
	"bytes"
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"net/url"
	"path/filepath"
	"strings"
	"time"

	"github.com/julienschmidt/httprouter"
	"golang.org/x/crypto/openpgp/armor"
	"gopkg.in/errgo.v1"

	"hockeypuck/conflux/recon"
	"hockeypuck/hkp/sks"
	"hockeypuck/hkp/storage"
	log "hockeypuck/logrus"
	"hockeypuck/openpgp"
)

const (
	shortKeyIDLen       = 8
	longKeyIDLen        = 16
	fingerprintKeyIDLen = 40
)

var errKeywordSearchNotAvailable = errgo.New("keyword search is not available")

func httpError(w http.ResponseWriter, statusCode int, err error) {
	if statusCode != http.StatusNotFound {
		log.Errorf("HTTP %d: %v", statusCode, errgo.Details(err))
	}
	http.Error(w, http.StatusText(statusCode), statusCode)
}

type Handler struct {
	storage storage.Storage

	indexWriter  IndexFormat
	vindexWriter IndexFormat

	statsTemplate *template.Template
	statsFunc     func() (interface{}, error)

	selfSignedOnly  bool
	fingerprintOnly bool
}

type HandlerOption func(h *Handler) error

func IndexTemplate(path string, extra ...string) HandlerOption {
	return func(h *Handler) error {
		tw, err := NewHTMLFormat(path, extra)
		if err != nil {
			return errgo.Mask(err)
		}
		h.indexWriter = tw
		return nil
	}
}

func VIndexTemplate(path string, extra ...string) HandlerOption {
	return func(h *Handler) error {
		tw, err := NewHTMLFormat(path, extra)
		if err != nil {
			return errgo.Mask(err)
		}
		h.vindexWriter = tw
		return nil
	}
}

func StatsTemplate(path string, extra ...string) HandlerOption {
	return func(h *Handler) error {
		t := template.New(filepath.Base(path)).Funcs(template.FuncMap{
			"url": func(u *url.URL) template.URL {
				return template.URL(u.String())
			},
			"day": func(t time.Time) string {
				return t.Format("2006-01-02")
			},
			"hour": func(t time.Time) string {
				return t.Format("2006-01-02 15")
			},
		})
		var err error
		if len(extra) > 0 {
			t, err = t.ParseFiles(append([]string{path}, extra...)...)
		} else {
			t, err = t.ParseGlob(path)
		}
		if err != nil {
			return errgo.Mask(err)
		}
		h.statsTemplate = t
		return nil
	}
}

func StatsFunc(f func() (interface{}, error)) HandlerOption {
	return func(h *Handler) error {
		h.statsFunc = f
		return nil
	}
}

func SelfSignedOnly(selfSignedOnly bool) HandlerOption {
	return func(h *Handler) error {
		h.selfSignedOnly = selfSignedOnly
		return nil
	}
}

func FingerprintOnly(fingerprintOnly bool) HandlerOption {
	return func(h *Handler) error {
		h.fingerprintOnly = fingerprintOnly
		return nil
	}
}

func NewHandler(storage storage.Storage, options ...HandlerOption) (*Handler, error) {
	h := &Handler{
		storage: storage,
	}
	for _, option := range options {
		err := option(h)
		if err != nil {
			return nil, errgo.Mask(err)
		}
	}
	return h, nil
}

func (h *Handler) Register(r *httprouter.Router) {
	r.GET("/pks/lookup", h.Lookup)
	r.POST("/pks/add", h.Add)
	r.POST("/pks/hashquery", h.HashQuery)
}

func (h *Handler) Lookup(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	l, err := ParseLookup(r)
	if err != nil {
		httpError(w, http.StatusBadRequest, err)
		return
	}
	switch l.Op {
	case OperationGet, OperationHGet:
		h.get(w, l)
	case OperationIndex:
		h.index(w, l, h.indexWriter)
	case OperationVIndex:
		h.index(w, l, h.vindexWriter)
	case OperationStats:
		h.stats(w, l)
	default:
		httpError(w, http.StatusNotFound, errgo.Newf("operation not found: %v", l.Op))
		return
	}
}

func (h *Handler) HashQuery(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	hq, err := ParseHashQuery(r)
	if err != nil {
		httpError(w, http.StatusBadRequest, errgo.Mask(err))
		return
	}
	var result []*openpgp.PrimaryKey
	for _, digest := range hq.Digests {
		rfps, err := h.storage.MatchMD5([]string{digest})
		if err != nil {
			log.Errorf("error resolving hashquery digest %q", digest)
			continue
		}
		keys, err := h.storage.FetchKeys(rfps)
		if err != nil {
			log.Errorf("error fetching hashquery key %q", digest)
			continue
		}
		result = append(result, keys...)
	}

	w.Header().Set("Content-Type", "pgp/keys")

	// Write the number of keys
	err = recon.WriteInt(w, len(result))
	for _, key := range result {
		// Write each key in binary packet format, prefixed with length
		err = writeHashqueryKey(w, key)
		if err != nil {
			log.Errorf("error writing hashquery key %q: %v", key.RFingerprint, err)
			return
		}
	}

	// SKS expects hashquery response to terminate with a CRLF
	_, err = w.Write([]byte{0x0d, 0x0a})
	if err != nil {
		log.Errorf("error writing hashquery terminator: %v", err)
	}
}

func writeHashqueryKey(w http.ResponseWriter, key *openpgp.PrimaryKey) error {
	var buf bytes.Buffer
	err := openpgp.WritePackets(&buf, key)
	if err != nil {
		return errgo.Mask(err)
	}
	err = recon.WriteInt(w, buf.Len())
	if err != nil {
		return errgo.Mask(err)
	}
	_, err = w.Write(buf.Bytes())
	if err != nil {
		return errgo.Mask(err)
	}
	return nil
}

func (h *Handler) resolve(l *Lookup) ([]string, error) {
	if l.Op == OperationHGet {
		return h.storage.MatchMD5([]string{l.Search})
	}
	if strings.HasPrefix(l.Search, "0x") {
		keyID := openpgp.Reverse(strings.ToLower(l.Search[2:]))
		switch len(keyID) {
		case shortKeyIDLen, longKeyIDLen, fingerprintKeyIDLen:
			return h.storage.Resolve([]string{keyID})
		}
	}
	if h.fingerprintOnly {
		return nil, errKeywordSearchNotAvailable
	}
	return h.storage.MatchKeyword([]string{l.Search})
}

func (h *Handler) keys(l *Lookup) ([]*openpgp.PrimaryKey, error) {
	rfps, err := h.resolve(l)
	if err != nil {
		return nil, err
	}
	keys, err := h.storage.FetchKeys(rfps)
	if err != nil {
		return nil, errgo.Mask(err)
	}
	if h.selfSignedOnly {
		for _, key := range keys {
			if err := openpgp.SelfSignedOnly(key); err != nil {
				return nil, errgo.Mask(err)
			}
		}
	}
	return keys, nil
}

func (h *Handler) get(w http.ResponseWriter, l *Lookup) {
	keys, err := h.keys(l)
	if err == errKeywordSearchNotAvailable {
		httpError(w, http.StatusBadRequest, errgo.Mask(err))
		return
	} else if err != nil {
		httpError(w, http.StatusInternalServerError, errgo.Mask(err))
		return
	}
	if len(keys) == 0 {
		httpError(w, http.StatusNotFound, errgo.New("not found"))
		return
	}

	// Drop malformed packets, since these break GPG imports.
	for _, key := range keys {
		var others []*openpgp.Packet
		for _, other := range key.Others {
			if other.Malformed {
				continue
			}
			others = append(others, other)
		}
		key.Others = others
	}

	w.Header().Set("Content-Type", "text/plain")
	err = openpgp.WriteArmoredPackets(w, keys)
	if err != nil {
		log.Errorf("get %q: error writing armored keys: %v", l.Search, err)
	}
	// Write a trailing newline as required by the HKP spec
	// (§3.1.2.1) and as expected by many tools, e.g. RPM.
	_, err = w.Write([]byte("\n"))
	if err != nil {
		log.Errorf("get %q: failed to write trailing newline: %v", l.Search, err)
	}
}

func (h *Handler) index(w http.ResponseWriter, l *Lookup, f IndexFormat) {
	keys, err := h.keys(l)
	if err == errKeywordSearchNotAvailable {
		httpError(w, http.StatusBadRequest, errgo.Mask(err))
		return
	} else if err != nil {
		httpError(w, http.StatusInternalServerError, errgo.Mask(err))
		return
	}
	if len(keys) == 0 {
		httpError(w, http.StatusNotFound, errgo.New("not found"))
		return
	}

	if l.Options[OptionMachineReadable] {
		f = mrFormat
	} else if l.Options[OptionJSON] || f == nil {
		f = jsonFormat
	}

	err = f.Write(w, l, keys)
	if err != nil {
		httpError(w, http.StatusInternalServerError, errgo.Mask(err))
		return
	}
}

func (h *Handler) indexJSON(w http.ResponseWriter, keys []*openpgp.PrimaryKey) {
	w.Header().Set("Content-Type", "application/json")
	enc := json.NewEncoder(w)
	err := enc.Encode(&keys)
	if err != nil {
		httpError(w, http.StatusInternalServerError, errgo.Mask(err))
		return
	}
}

func mrTimeString(t time.Time) string {
	if t.IsZero() {
		return ""
	}
	return fmt.Sprintf("%d", t.Unix())
}

type StatsResponse struct {
	Info  interface{}
	Stats *sks.Stats
}

func (h *Handler) stats(w http.ResponseWriter, l *Lookup) {
	if h.statsFunc == nil {
		httpError(w, http.StatusBadRequest, errgo.New("stats not configured"))
		fmt.Fprintln(w, "stats not configured")
		return
	}
	data, err := h.statsFunc()
	if err != nil {
		httpError(w, http.StatusInternalServerError, errgo.Mask(err))
		return
	}

	if h.statsTemplate != nil && !(l.Options[OptionJSON] || l.Options[OptionMachineReadable]) {
		err = h.statsTemplate.Execute(w, data)
	} else {
		err = json.NewEncoder(w).Encode(data)
	}
	if err != nil {
		httpError(w, http.StatusInternalServerError, errgo.Mask(err))
	}
}

type AddResponse struct {
	Inserted []string `json:"inserted"`
	Updated  []string `json:"updated"`
	Ignored  []string `json:"ignored"`
}

func (h *Handler) Add(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	add, err := ParseAdd(r)
	if err != nil {
		httpError(w, http.StatusBadRequest, errgo.Mask(err))
		return
	}

	// Check and decode the armor
	armorBlock, err := armor.Decode(bytes.NewBufferString(add.Keytext))
	if err != nil {
		httpError(w, http.StatusBadRequest, errgo.Mask(err))
		return
	}

	var result AddResponse
	for readKey := range openpgp.ReadKeys(armorBlock.Body) {
		if readKey.Error != nil {
			httpError(w, http.StatusBadRequest, errgo.Mask(err))
			return
		}
		err := openpgp.DropDuplicates(readKey.PrimaryKey)
		if err != nil {
			httpError(w, http.StatusInternalServerError, errgo.Mask(err))
			return
		}
		change, err := storage.UpsertKey(h.storage, readKey.PrimaryKey)
		if err != nil {
			httpError(w, http.StatusInternalServerError, errgo.Mask(err))
			return
		}

		fp := readKey.PrimaryKey.QualifiedFingerprint()
		switch change.(type) {
		case storage.KeyAdded:
			result.Inserted = append(result.Inserted, fp)
		case storage.KeyReplaced:
			result.Updated = append(result.Updated, fp)
		case storage.KeyNotChanged:
			result.Ignored = append(result.Ignored, fp)
		}
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	enc := json.NewEncoder(w)
	enc.Encode(&result)
}
<|endoftext|>


<commit_before>package vm

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestDictionarySet(t *testing.T) {
	for _, th := range []*Thunk{
		True, False, Nil, NewNumber(42), NewString("risp"),
	} {
		_, ok := EmptyDictionary.Eval().(dictionaryType).Set(th.Eval(), Nil).(dictionaryType)
		assert.True(t, ok)
	}
}

func TestDictionaryGet(t *testing.T) {
	for _, kvs := range [][][]*Thunk{
		{{True, False}},
		{{Nil, NewNumber(42)}},
		{{True, False}, {False, NewNumber(42)}},
		{{True, False}, {False, NewNumber(42)}, {Nil, NewString("Mr. Value")}},
	} {
		d := EmptyDictionary.Eval().(dictionaryType)

		for i, kv := range kvs {
			t.Logf("Setting a %vth key...\n", i)

			k, v := extractKeyAndValue(kv)

			var ok bool
			d, ok = d.Set(k, v).(dictionaryType)

			assert.True(t, ok)
		}

		assert.Equal(t, len(kvs), int(d.hashMap.Size()))

		for i, kv := range kvs {
			t.Logf("Getting a %vth value...\n", i)

			k, v := extractKeyAndValue(kv)

			t.Log(k)

			if e, ok := d.Get(k).Eval().(errorType); ok {
				t.Log(e.message.Eval())
			}

			assert.True(t, bool(App(Equal, d.Get(k), v).Eval().(boolType)))
		}
	}
}

func extractKeyAndValue(kv []*Thunk) (Object, *Thunk) {
	if len(kv) != 2 {
		panic("Invalid number of arguments to extractKeyAndValue.")
	}

	return kv[0].Eval(), kv[1]
}
<commit_msg>Rename test functions for Dictionary<commit_after>package vm

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestDictionarySetMethod(t *testing.T) {
	for _, th := range []*Thunk{
		True, False, Nil, NewNumber(42), NewString("risp"),
	} {
		_, ok := EmptyDictionary.Eval().(dictionaryType).Set(th.Eval(), Nil).(dictionaryType)
		assert.True(t, ok)
	}
}

func TestDictionaryGetMethod(t *testing.T) {
	for _, kvs := range [][][]*Thunk{
		{{True, False}},
		{{Nil, NewNumber(42)}},
		{{True, False}, {False, NewNumber(42)}},
		{{True, False}, {False, NewNumber(42)}, {Nil, NewString("Mr. Value")}},
	} {
		d := EmptyDictionary.Eval().(dictionaryType)

		for i, kv := range kvs {
			t.Logf("Setting a %vth key...\n", i)

			k, v := extractKeyAndValue(kv)

			var ok bool
			d, ok = d.Set(k, v).(dictionaryType)

			assert.True(t, ok)
		}

		assert.Equal(t, len(kvs), int(d.hashMap.Size()))

		for i, kv := range kvs {
			t.Logf("Getting a %vth value...\n", i)

			k, v := extractKeyAndValue(kv)

			t.Log(k)

			if e, ok := d.Get(k).Eval().(errorType); ok {
				t.Log(e.message.Eval())
			}

			assert.True(t, bool(App(Equal, d.Get(k), v).Eval().(boolType)))
		}
	}
}

func extractKeyAndValue(kv []*Thunk) (Object, *Thunk) {
	if len(kv) != 2 {
		panic("Invalid number of arguments to extractKeyAndValue.")
	}

	return kv[0].Eval(), kv[1]
}
<|endoftext|>


<commit_before>package version

import (
	"fmt"

	pb "github.com/pachyderm/pachyderm/src/client/version/versionpb"
)

const (
	// MajorVersion is the current major version for pachyderm.
	MajorVersion = 1
	// MinorVersion is the current minor version for pachyderm.
	MinorVersion = 7
	// MicroVersion is the patch number for pachyderm.
	MicroVersion = 0
)

var (
	// AdditionalVersion is the string provided at release time
	// The value is passed to the linker at build time
	// DO NOT set the value of this variable here
	AdditionalVersion string = "rc1"
	// Version is the current version for pachyderm.
	Version = &pb.Version{
		Major:      MajorVersion,
		Minor:      MinorVersion,
		Micro:      MicroVersion,
		Additional: AdditionalVersion,
	}
)

// PrettyPrintVersion returns a version string optionally tagged with metadata.
// For example: "1.2.3", or "1.2.3rc1" if version.Additional is "rc1".
func PrettyPrintVersion(version *pb.Version) string {
	result := PrettyPrintVersionNoAdditional(version)
	if version.Additional != "" {
		result += fmt.Sprintf("%s", version.Additional)
	}
	return result
}

// PrettyPrintVersionNoAdditional returns a version string without
// version.Additional.
func PrettyPrintVersionNoAdditional(version *pb.Version) string {
	return fmt.Sprintf("%d.%d.%d", version.Major, version.Minor, version.Micro)
}
<commit_msg>Fix the linter.<commit_after>package version

import (
	"fmt"

	pb "github.com/pachyderm/pachyderm/src/client/version/versionpb"
)

const (
	// MajorVersion is the current major version for pachyderm.
	MajorVersion = 1
	// MinorVersion is the current minor version for pachyderm.
	MinorVersion = 7
	// MicroVersion is the patch number for pachyderm.
	MicroVersion = 0
)

var (
	// AdditionalVersion is the string provided at release time
	// The value is passed to the linker at build time
	// DO NOT set the value of this variable here
	AdditionalVersion = "rc1"
	// Version is the current version for pachyderm.
	Version = &pb.Version{
		Major:      MajorVersion,
		Minor:      MinorVersion,
		Micro:      MicroVersion,
		Additional: AdditionalVersion,
	}
)

// PrettyPrintVersion returns a version string optionally tagged with metadata.
// For example: "1.2.3", or "1.2.3rc1" if version.Additional is "rc1".
func PrettyPrintVersion(version *pb.Version) string {
	result := PrettyPrintVersionNoAdditional(version)
	if version.Additional != "" {
		result += fmt.Sprintf("%s", version.Additional)
	}
	return result
}

// PrettyPrintVersionNoAdditional returns a version string without
// version.Additional.
func PrettyPrintVersionNoAdditional(version *pb.Version) string {
	return fmt.Sprintf("%d.%d.%d", version.Major, version.Minor, version.Micro)
}
<|endoftext|>