<commit_before>package wkb

import (
	"encoding/binary"
	"fmt"
	"io"
)

type wkbReader func(io.Reader, binary.ByteOrder) (Geom, error)

var wkbReaders = map[uint32]wkbReader{
	wkbPoint:        pointReader,
	wkbPointZ:       pointZReader,
	wkbPointM:       pointMReader,
	wkbPointZM:      pointZMReader,
	wkbLineString:   lineStringReader,
	wkbLineStringZ:  lineStringZReader,
	wkbLineStringM:  lineStringMReader,
	wkbLineStringZM: lineStringZMReader,
	wkbPolygon:      polygonReader,
	wkbPolygonZ:     polygonZReader,
	wkbPolygonM:     polygonMReader,
	wkbPolygonZM:    polygonZMReader,
}

func readLinearRing(r io.Reader, byteOrder binary.ByteOrder) ([]Point, error) {
	var numPoints uint32
	if err := binary.Read(r, byteOrder, &numPoints); err != nil {
		return nil, err
	}
	points := make([]Point, numPoints)
	if err := binary.Read(r, byteOrder, &points); err != nil {
		return nil, err
	}
	return points, nil
}

func readLinearRingZ(r io.Reader, byteOrder binary.ByteOrder) ([]PointZ, error) {
	var numPoints uint32
	if err := binary.Read(r, byteOrder, &numPoints); err != nil {
		return nil, err
	}
	pointZs := make([]PointZ, numPoints)
	if err := binary.Read(r, byteOrder, &pointZs); err != nil {
		return nil, err
	}
	return pointZs, nil
}

func readLinearRingM(r io.Reader, byteOrder binary.ByteOrder) ([]PointM, error) {
	var numPoints uint32
	if err := binary.Read(r, byteOrder, &numPoints); err != nil {
		return nil, err
	}
	pointMs := make([]PointM, numPoints)
	if err := binary.Read(r, byteOrder, &pointMs); err != nil {
		return nil, err
	}
	return pointMs, nil
}

func readLinearRingZM(r io.Reader, byteOrder binary.ByteOrder) ([]PointZM, error) {
	var numPoints uint32
	if err := binary.Read(r, byteOrder, &numPoints); err != nil {
		return nil, err
	}
	pointZMs := make([]PointZM, numPoints)
	if err := binary.Read(r, byteOrder, &pointZMs); err != nil {
		return nil, err
	}
	return pointZMs, nil
}

func pointReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	point := Point{}
	if err := binary.Read(r, byteOrder, &point); err != nil {
		return nil, err
	}
	return point, nil
}

func pointZReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	pointZ := PointZ{}
	if err := binary.Read(r, byteOrder, &pointZ); err != nil {
		return nil, err
	}
	return pointZ, nil
}

func pointMReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	pointM := PointM{}
	if err := binary.Read(r, byteOrder, &pointM); err != nil {
		return nil, err
	}
	return pointM, nil
}

func pointZMReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	pointZM := PointZM{}
	if err := binary.Read(r, byteOrder, &pointZM); err != nil {
		return nil, err
	}
	return pointZM, nil
}

func lineStringReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	points, err := readLinearRing(r, byteOrder)
	if err != nil {
		return nil, err
	}
	return LineString{points}, nil
}

func lineStringZReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	pointZs, err := readLinearRingZ(r, byteOrder)
	if err != nil {
		return nil, err
	}
	return LineStringZ{pointZs}, nil
}

func lineStringMReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	pointMs, err := readLinearRingM(r, byteOrder)
	if err != nil {
		return nil, err
	}
	return LineStringM{pointMs}, nil
}

func lineStringZMReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	pointZMs, err := readLinearRingZM(r, byteOrder)
	if err != nil {
		return nil, err
	}
	return LineStringZM{pointZMs}, nil
}

func polygonReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	var numRings uint32
	if err := binary.Read(r, byteOrder, &numRings); err != nil {
		return nil, err
	}
	rings := make([]LinearRing, numRings)
	for i := uint32(0); i < numRings; i++ {
		if points, err := readLinearRing(r, byteOrder); err != nil {
			return nil, err
		} else {
			rings[i] = points
		}
	}
	return Polygon{rings}, nil
}

func polygonZReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	var numRings uint32
	if err := binary.Read(r, byteOrder, &numRings); err != nil {
		return nil, err
	}
	rings := make([]LinearRingZ, numRings)
	for i := uint32(0); i < numRings; i++ {
		if pointZs, err := readLinearRingZ(r, byteOrder); err != nil {
			return nil, err
		} else {
			rings[i] = pointZs
		}
	}
	return PolygonZ{rings}, nil
}

func polygonMReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	var numRings uint32
	if err := binary.Read(r, byteOrder, &numRings); err != nil {
		return nil, err
	}
	rings := make([]LinearRingM, numRings)
	for i := uint32(0); i < numRings; i++ {
		if pointMs, err := readLinearRingM(r, byteOrder); err != nil {
			return nil, err
		} else {
			rings[i] = pointMs
		}
	}
	return PolygonM{rings}, nil
}

func polygonZMReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	var numRings uint32
	if err := binary.Read(r, byteOrder, &numRings); err != nil {
		return nil, err
	}
	rings := make([]LinearRingZM, numRings)
	for i := uint32(0); i < numRings; i++ {
		if pointZMs, err := readLinearRingZM(r, byteOrder); err != nil {
			return nil, err
		} else {
			rings[i] = pointZMs
		}
	}
	return PolygonZM{rings}, nil
}

func Read(r io.Reader) (Geom, error) {

	var wkbByteOrder uint8
	if err := binary.Read(r, binary.LittleEndian, &wkbByteOrder); err != nil {
		return nil, err
	}
	var byteOrder binary.ByteOrder
	switch wkbByteOrder {
	case wkbXDR:
		byteOrder = binary.BigEndian
	case wkbNDR:
		byteOrder = binary.LittleEndian
	default:
		return nil, fmt.Errorf("invalid byte order %u", wkbByteOrder)
	}

	var wkbGeometryType uint32
	if err := binary.Read(r, byteOrder, &wkbGeometryType); err != nil {
		return nil, err
	}

	if reader, ok := wkbReaders[wkbGeometryType]; ok {
		return reader(r, byteOrder)
	} else {
		return nil, fmt.Errorf("unsupported geometry type %u", wkbGeometryType)
	}

}
<commit_msg>Rename variables for consistency<commit_after>package wkb

import (
	"encoding/binary"
	"fmt"
	"io"
)

type wkbReader func(io.Reader, binary.ByteOrder) (Geom, error)

var wkbReaders = map[uint32]wkbReader{
	wkbPoint:        pointReader,
	wkbPointZ:       pointZReader,
	wkbPointM:       pointMReader,
	wkbPointZM:      pointZMReader,
	wkbLineString:   lineStringReader,
	wkbLineStringZ:  lineStringZReader,
	wkbLineStringM:  lineStringMReader,
	wkbLineStringZM: lineStringZMReader,
	wkbPolygon:      polygonReader,
	wkbPolygonZ:     polygonZReader,
	wkbPolygonM:     polygonMReader,
	wkbPolygonZM:    polygonZMReader,
}

func readLinearRing(r io.Reader, byteOrder binary.ByteOrder) ([]Point, error) {
	var numPoints uint32
	if err := binary.Read(r, byteOrder, &numPoints); err != nil {
		return nil, err
	}
	points := make([]Point, numPoints)
	if err := binary.Read(r, byteOrder, &points); err != nil {
		return nil, err
	}
	return points, nil
}

func readLinearRingZ(r io.Reader, byteOrder binary.ByteOrder) ([]PointZ, error) {
	var numPoints uint32
	if err := binary.Read(r, byteOrder, &numPoints); err != nil {
		return nil, err
	}
	pointZs := make([]PointZ, numPoints)
	if err := binary.Read(r, byteOrder, &pointZs); err != nil {
		return nil, err
	}
	return pointZs, nil
}

func readLinearRingM(r io.Reader, byteOrder binary.ByteOrder) ([]PointM, error) {
	var numPoints uint32
	if err := binary.Read(r, byteOrder, &numPoints); err != nil {
		return nil, err
	}
	pointMs := make([]PointM, numPoints)
	if err := binary.Read(r, byteOrder, &pointMs); err != nil {
		return nil, err
	}
	return pointMs, nil
}

func readLinearRingZM(r io.Reader, byteOrder binary.ByteOrder) ([]PointZM, error) {
	var numPoints uint32
	if err := binary.Read(r, byteOrder, &numPoints); err != nil {
		return nil, err
	}
	pointZMs := make([]PointZM, numPoints)
	if err := binary.Read(r, byteOrder, &pointZMs); err != nil {
		return nil, err
	}
	return pointZMs, nil
}

func pointReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	point := Point{}
	if err := binary.Read(r, byteOrder, &point); err != nil {
		return nil, err
	}
	return point, nil
}

func pointZReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	pointZ := PointZ{}
	if err := binary.Read(r, byteOrder, &pointZ); err != nil {
		return nil, err
	}
	return pointZ, nil
}

func pointMReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	pointM := PointM{}
	if err := binary.Read(r, byteOrder, &pointM); err != nil {
		return nil, err
	}
	return pointM, nil
}

func pointZMReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	pointZM := PointZM{}
	if err := binary.Read(r, byteOrder, &pointZM); err != nil {
		return nil, err
	}
	return pointZM, nil
}

func lineStringReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	points, err := readLinearRing(r, byteOrder)
	if err != nil {
		return nil, err
	}
	return LineString{points}, nil
}

func lineStringZReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	pointZs, err := readLinearRingZ(r, byteOrder)
	if err != nil {
		return nil, err
	}
	return LineStringZ{pointZs}, nil
}

func lineStringMReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	pointMs, err := readLinearRingM(r, byteOrder)
	if err != nil {
		return nil, err
	}
	return LineStringM{pointMs}, nil
}

func lineStringZMReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	pointZMs, err := readLinearRingZM(r, byteOrder)
	if err != nil {
		return nil, err
	}
	return LineStringZM{pointZMs}, nil
}

func polygonReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	var numRings uint32
	if err := binary.Read(r, byteOrder, &numRings); err != nil {
		return nil, err
	}
	rings := make([]LinearRing, numRings)
	for i := uint32(0); i < numRings; i++ {
		if points, err := readLinearRing(r, byteOrder); err != nil {
			return nil, err
		} else {
			rings[i] = points
		}
	}
	return Polygon{rings}, nil
}

func polygonZReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	var numRings uint32
	if err := binary.Read(r, byteOrder, &numRings); err != nil {
		return nil, err
	}
	ringZs := make([]LinearRingZ, numRings)
	for i := uint32(0); i < numRings; i++ {
		if pointZs, err := readLinearRingZ(r, byteOrder); err != nil {
			return nil, err
		} else {
			ringZs[i] = pointZs
		}
	}
	return PolygonZ{ringZs}, nil
}

func polygonMReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	var numRings uint32
	if err := binary.Read(r, byteOrder, &numRings); err != nil {
		return nil, err
	}
	ringMs := make([]LinearRingM, numRings)
	for i := uint32(0); i < numRings; i++ {
		if pointMs, err := readLinearRingM(r, byteOrder); err != nil {
			return nil, err
		} else {
			ringMs[i] = pointMs
		}
	}
	return PolygonM{ringMs}, nil
}

func polygonZMReader(r io.Reader, byteOrder binary.ByteOrder) (Geom, error) {
	var numRings uint32
	if err := binary.Read(r, byteOrder, &numRings); err != nil {
		return nil, err
	}
	ringZMs := make([]LinearRingZM, numRings)
	for i := uint32(0); i < numRings; i++ {
		if pointZMs, err := readLinearRingZM(r, byteOrder); err != nil {
			return nil, err
		} else {
			ringZMs[i] = pointZMs
		}
	}
	return PolygonZM{ringZMs}, nil
}

func Read(r io.Reader) (Geom, error) {

	var wkbByteOrder uint8
	if err := binary.Read(r, binary.LittleEndian, &wkbByteOrder); err != nil {
		return nil, err
	}
	var byteOrder binary.ByteOrder
	switch wkbByteOrder {
	case wkbXDR:
		byteOrder = binary.BigEndian
	case wkbNDR:
		byteOrder = binary.LittleEndian
	default:
		return nil, fmt.Errorf("invalid byte order %u", wkbByteOrder)
	}

	var wkbGeometryType uint32
	if err := binary.Read(r, byteOrder, &wkbGeometryType); err != nil {
		return nil, err
	}

	if reader, ok := wkbReaders[wkbGeometryType]; ok {
		return reader(r, byteOrder)
	} else {
		return nil, fmt.Errorf("unsupported geometry type %u", wkbGeometryType)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strconv"

	"github.com/Sirupsen/logrus"
	"github.com/gorilla/context"

	pg "github.com/mozilla/tls-observatory/database"
	"github.com/mozilla/tls-observatory/logger"
)

func ScanHandler(w http.ResponseWriter, r *http.Request) {
	var (
		status int
		err    error
	)

	defer func() {
		if nil != err {
			http.Error(w, err.Error(), status)
		}
	}()

	log := logger.GetLogger()
	status = http.StatusInternalServerError

	log.WithFields(logrus.Fields{
		"form values": r.Form,
		"headers":     r.Header,
	}).Debug("Scan endpoint received request")

	val, ok := context.GetOk(r, dbKey)
	if !ok {
		log.Error("Could not find db in request context")
		err = errors.New("Could not access database.")
		return
	}

	db := val.(*pg.DB)

	domain := r.FormValue("target")

	if validateDomain(domain) {

		scan, err := db.NewScan(domain, -1) //no replay
		if err != nil {
			log.WithFields(logrus.Fields{
				"domain": domain,
				"error":  err.Error(),
			}).Error("Could not create new scan")
			err = errors.New("Could not create new scan")
			return
		}

		resp := fmt.Sprintf(`{"scan_id":"%d"}`, scan.ID)
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, resp)
	} else {
		w.WriteHeader(http.StatusBadRequest)
	}
}

func ResultHandler(w http.ResponseWriter, r *http.Request) {

	var (
		status int
		err    error
	)

	defer func() {
		if nil != err {
			http.Error(w, err.Error(), status)
		}
	}()

	log := logger.GetLogger()
	status = http.StatusInternalServerError

	log.WithFields(logrus.Fields{
		"form values": r.Form,
		"headers":     r.Header,
	}).Debug("Results endpoint received request")

	val, ok := context.GetOk(r, dbKey)
	if !ok {
		log.Error("Could not find db in request context")
		err = errors.New("Could not access database.")
		return
	}

	db := val.(*pg.DB)

	idStr := r.FormValue("id")

	id, err := strconv.ParseInt(idStr, 10, 64)
	if err != nil {
		log.WithFields(logrus.Fields{
			"scan_id": idStr,
			"error":   err.Error(),
		}).Error("Could not parse scanid")
		err = errors.New("Could not parse provided scan id")
		status = http.StatusBadRequest
		return
	}

	scan, err := db.GetScanByID(id)
	if err != nil {
		log.WithFields(logrus.Fields{
			"scan_id": id,
			"error":   err.Error(),
		}).Error("Could not get scan from database")
		err = errors.New("Could not access database to get requested scan.")
		return
	}

	if scan.ID == -1 {
		log.WithFields(logrus.Fields{
			"scan_id": id,
		}).Debug("Did not find scan in database")

		err = errors.New("Could not find a scan with the id you provided.")
		status = http.StatusNotFound
		return
	}

	jsScan, err := json.Marshal(scan)
	if err != nil {
		log.WithFields(logrus.Fields{
			"scan_id": id,
			"error":   err.Error(),
		}).Error("Could not Marshal scan")

		err = errors.New("Could not process the requested scan")
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	fmt.Fprint(w, string(jsScan))
}

func CertificateHandler(w http.ResponseWriter, r *http.Request) {

	var (
		status int
		err    error
	)

	defer func() {
		if nil != err {
			http.Error(w, err.Error(), status)
		}
	}()

	log := logger.GetLogger()
	status = http.StatusInternalServerError

	log.WithFields(logrus.Fields{
		"form values": r.Form.Encode(),
		"headers":     r.Header,
	}).Debug("Certificate Endpoint received request")

	val, ok := context.GetOk(r, dbKey)
	if !ok {
		log.Error("Could not find db in request context")
		err = errors.New("Could not access database.")
		return
	}

	db := val.(*pg.DB)

	idStr := r.FormValue("id")

	id, err := strconv.ParseInt(idStr, 10, 64)
	if err != nil {
		log.WithFields(logrus.Fields{
			"cert_id": id,
			"error":   err.Error(),
		}).Error("Could not parse certificate id")

		status = http.StatusBadRequest
		err = errors.New("Could not parse provided certificate id")
		return
	}

	cert, err := db.GetCertByID(id)
	if err != nil {
		log.WithFields(logrus.Fields{
			"cert_id": id,
			"error":   err.Error(),
		}).Error("Could not get cert from database")

		err = errors.New("Could not access database to get requested certificate")
		return
	}

	jsScan, err := json.Marshal(cert)
	if err != nil {
		log.WithFields(logrus.Fields{
			"cert_id": id,
			"error":   err.Error(),
		}).Error("Could not Marshal cert")

		err = errors.New("Could not process requested certificate")
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	fmt.Fprint(w, string(jsScan))
}

func validateDomain(domain string) bool {

	// TODO
	// Need to validate the domain, in a way,
	// before passing it to the retriever queue

	return true
}
<commit_msg>Basic patch to avoid scanning empty targets<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strconv"

	"github.com/Sirupsen/logrus"
	"github.com/gorilla/context"

	pg "github.com/mozilla/tls-observatory/database"
	"github.com/mozilla/tls-observatory/logger"
)

func ScanHandler(w http.ResponseWriter, r *http.Request) {
	var (
		status int
		err    error
	)

	defer func() {
		if nil != err {
			http.Error(w, err.Error(), status)
		}
	}()

	log := logger.GetLogger()
	status = http.StatusInternalServerError

	log.WithFields(logrus.Fields{
		"form values": r.Form,
		"headers":     r.Header,
	}).Debug("Scan endpoint received request")

	val, ok := context.GetOk(r, dbKey)
	if !ok {
		log.Error("Could not find db in request context")
		err = errors.New("Could not access database.")
		return
	}

	db := val.(*pg.DB)

	domain := r.FormValue("target")

	if validateDomain(domain) {

		scan, err := db.NewScan(domain, -1) //no replay
		if err != nil {
			log.WithFields(logrus.Fields{
				"domain": domain,
				"error":  err.Error(),
			}).Error("Could not create new scan")
			err = errors.New("Could not create new scan")
			return
		}

		resp := fmt.Sprintf(`{"scan_id":"%d"}`, scan.ID)
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, resp)
	} else {
		w.WriteHeader(http.StatusBadRequest)
	}
}

func ResultHandler(w http.ResponseWriter, r *http.Request) {

	var (
		status int
		err    error
	)

	defer func() {
		if nil != err {
			http.Error(w, err.Error(), status)
		}
	}()

	log := logger.GetLogger()
	status = http.StatusInternalServerError

	log.WithFields(logrus.Fields{
		"form values": r.Form,
		"headers":     r.Header,
	}).Debug("Results endpoint received request")

	val, ok := context.GetOk(r, dbKey)
	if !ok {
		log.Error("Could not find db in request context")
		err = errors.New("Could not access database.")
		return
	}

	db := val.(*pg.DB)

	idStr := r.FormValue("id")

	id, err := strconv.ParseInt(idStr, 10, 64)
	if err != nil {
		log.WithFields(logrus.Fields{
			"scan_id": idStr,
			"error":   err.Error(),
		}).Error("Could not parse scanid")
		err = errors.New("Could not parse provided scan id")
		status = http.StatusBadRequest
		return
	}

	scan, err := db.GetScanByID(id)
	if err != nil {
		log.WithFields(logrus.Fields{
			"scan_id": id,
			"error":   err.Error(),
		}).Error("Could not get scan from database")
		err = errors.New("Could not access database to get requested scan.")
		return
	}

	if scan.ID == -1 {
		log.WithFields(logrus.Fields{
			"scan_id": id,
		}).Debug("Did not find scan in database")

		err = errors.New("Could not find a scan with the id you provided.")
		status = http.StatusNotFound
		return
	}

	jsScan, err := json.Marshal(scan)
	if err != nil {
		log.WithFields(logrus.Fields{
			"scan_id": id,
			"error":   err.Error(),
		}).Error("Could not Marshal scan")

		err = errors.New("Could not process the requested scan")
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	fmt.Fprint(w, string(jsScan))
}

func CertificateHandler(w http.ResponseWriter, r *http.Request) {

	var (
		status int
		err    error
	)

	defer func() {
		if nil != err {
			http.Error(w, err.Error(), status)
		}
	}()

	log := logger.GetLogger()
	status = http.StatusInternalServerError

	log.WithFields(logrus.Fields{
		"form values": r.Form.Encode(),
		"headers":     r.Header,
	}).Debug("Certificate Endpoint received request")

	val, ok := context.GetOk(r, dbKey)
	if !ok {
		log.Error("Could not find db in request context")
		err = errors.New("Could not access database.")
		return
	}

	db := val.(*pg.DB)

	idStr := r.FormValue("id")

	id, err := strconv.ParseInt(idStr, 10, 64)
	if err != nil {
		log.WithFields(logrus.Fields{
			"cert_id": id,
			"error":   err.Error(),
		}).Error("Could not parse certificate id")

		status = http.StatusBadRequest
		err = errors.New("Could not parse provided certificate id")
		return
	}

	cert, err := db.GetCertByID(id)
	if err != nil {
		log.WithFields(logrus.Fields{
			"cert_id": id,
			"error":   err.Error(),
		}).Error("Could not get cert from database")

		err = errors.New("Could not access database to get requested certificate")
		return
	}

	jsScan, err := json.Marshal(cert)
	if err != nil {
		log.WithFields(logrus.Fields{
			"cert_id": id,
			"error":   err.Error(),
		}).Error("Could not Marshal cert")

		err = errors.New("Could not process requested certificate")
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	fmt.Fprint(w, string(jsScan))
}

func validateDomain(domain string) bool {
	if domain == "" {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>// Copyright 2013 The Prometheus Authors
// 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 rules

import (
	"context"
	"fmt"
	"math"
	"sort"
	"testing"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/prometheus/common/model"

	"github.com/prometheus/prometheus/pkg/labels"
	"github.com/prometheus/prometheus/pkg/timestamp"
	"github.com/prometheus/prometheus/pkg/value"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/util/testutil"
)

func TestAlertingRule(t *testing.T) {
	suite, err := promql.NewTest(t, `
		load 5m
			http_requests{job="app-server", instance="0", group="canary", severity="overwrite-me"}	75 85  95 105 105  95  85
			http_requests{job="app-server", instance="1", group="canary", severity="overwrite-me"}	80 90 100 110 120 130 140
	`)
	testutil.Ok(t, err)
	defer suite.Close()

	err = suite.Run()
	testutil.Ok(t, err)

	expr, err := promql.ParseExpr(`http_requests{group="canary", job="app-server"} < 100`)
	testutil.Ok(t, err)

	rule := NewAlertingRule(
		"HTTPRequestRateLow",
		expr,
		time.Minute,
		labels.FromStrings("severity", "{{\"c\"}}ritical"),
		nil, nil,
	)
	result := promql.Vector{
		{
			Metric: labels.FromStrings(
				"__name__", "ALERTS",
				"alertname", "HTTPRequestRateLow",
				"alertstate", "pending",
				"group", "canary",
				"instance", "0",
				"job", "app-server",
				"severity", "critical",
			),
			Point: promql.Point{V: 1},
		},
		{
			Metric: labels.FromStrings(
				"__name__", "ALERTS",
				"alertname", "HTTPRequestRateLow",
				"alertstate", "pending",
				"group", "canary",
				"instance", "1",
				"job", "app-server",
				"severity", "critical",
			),
			Point: promql.Point{V: 1},
		},
		{
			Metric: labels.FromStrings(
				"__name__", "ALERTS",
				"alertname", "HTTPRequestRateLow",
				"alertstate", "firing",
				"group", "canary",
				"instance", "0",
				"job", "app-server",
				"severity", "critical",
			),
			Point: promql.Point{V: 1},
		},
		{
			Metric: labels.FromStrings(
				"__name__", "ALERTS",
				"alertname", "HTTPRequestRateLow",
				"alertstate", "firing",
				"group", "canary",
				"instance", "1",
				"job", "app-server",
				"severity", "critical",
			),
			Point: promql.Point{V: 1},
		},
	}

	baseTime := time.Unix(0, 0)

	var tests = []struct {
		time   time.Duration
		result promql.Vector
	}{
		{
			time:   0,
			result: result[:2],
		}, {
			time:   5 * time.Minute,
			result: result[2:],
		}, {
			time:   10 * time.Minute,
			result: result[2:3],
		},
		{
			time:   15 * time.Minute,
			result: nil,
		},
		{
			time:   20 * time.Minute,
			result: nil,
		},
		{
			time:   25 * time.Minute,
			result: result[:1],
		},
		{
			time:   30 * time.Minute,
			result: result[2:3],
		},
	}

	for i, test := range tests {
		t.Logf("case %d", i)

		evalTime := baseTime.Add(test.time)

		res, err := rule.Eval(suite.Context(), evalTime, EngineQueryFunc(suite.QueryEngine(), suite.Storage()), nil)
		testutil.Ok(t, err)

		for i := range test.result {
			test.result[i].T = timestamp.FromTime(evalTime)
		}
		testutil.Assert(t, len(test.result) == len(res), "%d. Number of samples in expected and actual output don't match (%d vs. %d)", i, len(test.result), len(res))

		sort.Slice(res, func(i, j int) bool {
			return labels.Compare(res[i].Metric, res[j].Metric) < 0
		})
		testutil.Equals(t, test.result, res)

		for _, aa := range rule.ActiveAlerts() {
			testutil.Assert(t, aa.Labels.Get(model.MetricNameLabel) == "", "%s label set on active alert: %s", model.MetricNameLabel, aa.Labels)
		}
	}
}

func annotateWithTime(lines []string, ts time.Time) []string {
	annotatedLines := []string{}
	for _, line := range lines {
		annotatedLines = append(annotatedLines, fmt.Sprintf(line, timestamp.FromTime(ts)))
	}
	return annotatedLines
}

func TestStaleness(t *testing.T) {
	storage := testutil.NewStorage(t)
	defer storage.Close()
	engine := promql.NewEngine(nil, nil, 10, 10*time.Second)
	opts := &ManagerOptions{
		QueryFunc:  EngineQueryFunc(engine, storage),
		Appendable: storage,
		Context:    context.Background(),
		Logger:     log.NewNopLogger(),
	}

	expr, err := promql.ParseExpr("a + 1")
	testutil.Ok(t, err)
	rule := NewRecordingRule("a_plus_one", expr, labels.Labels{})
	group := NewGroup("default", "", time.Second, []Rule{rule}, opts)

	// A time series that has two samples and then goes stale.
	app, _ := storage.Appender()
	app.Add(labels.FromStrings(model.MetricNameLabel, "a"), 0, 1)
	app.Add(labels.FromStrings(model.MetricNameLabel, "a"), 1000, 2)
	app.Add(labels.FromStrings(model.MetricNameLabel, "a"), 2000, math.Float64frombits(value.StaleNaN))

	err = app.Commit()
	testutil.Ok(t, err)

	ctx := context.Background()

	// Execute 3 times, 1 second apart.
	group.Eval(ctx, time.Unix(0, 0))
	group.Eval(ctx, time.Unix(1, 0))
	group.Eval(ctx, time.Unix(2, 0))

	querier, err := storage.Querier(context.Background(), 0, 2000)
	testutil.Ok(t, err)
	defer querier.Close()

	matcher, err := labels.NewMatcher(labels.MatchEqual, model.MetricNameLabel, "a_plus_one")
	testutil.Ok(t, err)

	set, err := querier.Select(nil, matcher)
	testutil.Ok(t, err)

	samples, err := readSeriesSet(set)
	testutil.Ok(t, err)

	metric := labels.FromStrings(model.MetricNameLabel, "a_plus_one").String()
	metricSample, ok := samples[metric]

	testutil.Assert(t, ok, "Series %s not returned.", metric)
	testutil.Assert(t, value.IsStaleNaN(metricSample[2].V), "Appended second sample not as expected. Wanted: stale NaN Got: %x", math.Float64bits(metricSample[2].V))
	metricSample[2].V = 42 // reflect.DeepEqual cannot handle NaN.

	want := map[string][]promql.Point{
		metric: []promql.Point{{0, 2}, {1000, 3}, {2000, 42}},
	}

	testutil.Equals(t, want, samples)
}

// Convert a SeriesSet into a form useable with reflect.DeepEqual.
func readSeriesSet(ss storage.SeriesSet) (map[string][]promql.Point, error) {
	result := map[string][]promql.Point{}

	for ss.Next() {
		series := ss.At()

		points := []promql.Point{}
		it := series.Iterator()
		for it.Next() {
			t, v := it.At()
			points = append(points, promql.Point{T: t, V: v})
		}

		name := series.Labels().String()
		result[name] = points
	}
	return result, ss.Err()
}

func TestCopyState(t *testing.T) {
	oldGroup := &Group{
		rules: []Rule{
			NewAlertingRule("alert", nil, 0, nil, nil, nil),
			NewRecordingRule("rule1", nil, nil),
			NewRecordingRule("rule2", nil, nil),
			NewRecordingRule("rule3", nil, nil),
			NewRecordingRule("rule3", nil, nil),
		},
		seriesInPreviousEval: []map[string]labels.Labels{
			map[string]labels.Labels{"a": nil},
			map[string]labels.Labels{"r1": nil},
			map[string]labels.Labels{"r2": nil},
			map[string]labels.Labels{"r3a": nil},
			map[string]labels.Labels{"r3b": nil},
		},
		evaluationTime: time.Second,
	}
	oldGroup.rules[0].(*AlertingRule).active[42] = nil
	newGroup := &Group{
		rules: []Rule{
			NewRecordingRule("rule3", nil, nil),
			NewRecordingRule("rule3", nil, nil),
			NewRecordingRule("rule3", nil, nil),
			NewAlertingRule("alert", nil, 0, nil, nil, nil),
			NewRecordingRule("rule1", nil, nil),
			NewRecordingRule("rule4", nil, nil),
		},
		seriesInPreviousEval: make([]map[string]labels.Labels, 6),
	}
	newGroup.copyState(oldGroup)

	want := []map[string]labels.Labels{
		map[string]labels.Labels{"r3a": nil},
		map[string]labels.Labels{"r3b": nil},
		nil,
		map[string]labels.Labels{"a": nil},
		map[string]labels.Labels{"r1": nil},
		nil,
	}
	testutil.Equals(t, want, newGroup.seriesInPreviousEval)
	testutil.Equals(t, oldGroup.rules[0], newGroup.rules[3])
	testutil.Equals(t, oldGroup.evaluationTime, newGroup.evaluationTime)
}

func TestUpdate(t *testing.T) {
	expected := map[string]labels.Labels{
		"test": labels.FromStrings("name", "value"),
	}
	ruleManager := NewManager(&ManagerOptions{
		Context: context.Background(),
		Logger:  log.NewNopLogger(),
	})
	ruleManager.Run()

	err := ruleManager.Update(0, nil)
	testutil.Ok(t, err)
	for _, g := range ruleManager.groups {
		g.seriesInPreviousEval = []map[string]labels.Labels{
			expected,
		}
	}

	err = ruleManager.Update(0, nil)
	testutil.Ok(t, err)
	for _, g := range ruleManager.groups {
		for _, actual := range g.seriesInPreviousEval {
			testutil.Equals(t, expected, actual)
		}
	}
}
<commit_msg>Cleans up test by removing unused function (#3969)<commit_after>// Copyright 2013 The Prometheus Authors
// 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 rules

import (
	"context"
	"math"
	"sort"
	"testing"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/prometheus/common/model"

	"github.com/prometheus/prometheus/pkg/labels"
	"github.com/prometheus/prometheus/pkg/timestamp"
	"github.com/prometheus/prometheus/pkg/value"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/storage"
	"github.com/prometheus/prometheus/util/testutil"
)

func TestAlertingRule(t *testing.T) {
	suite, err := promql.NewTest(t, `
		load 5m
			http_requests{job="app-server", instance="0", group="canary", severity="overwrite-me"}	75 85  95 105 105  95  85
			http_requests{job="app-server", instance="1", group="canary", severity="overwrite-me"}	80 90 100 110 120 130 140
	`)
	testutil.Ok(t, err)
	defer suite.Close()

	err = suite.Run()
	testutil.Ok(t, err)

	expr, err := promql.ParseExpr(`http_requests{group="canary", job="app-server"} < 100`)
	testutil.Ok(t, err)

	rule := NewAlertingRule(
		"HTTPRequestRateLow",
		expr,
		time.Minute,
		labels.FromStrings("severity", "{{\"c\"}}ritical"),
		nil, nil,
	)
	result := promql.Vector{
		{
			Metric: labels.FromStrings(
				"__name__", "ALERTS",
				"alertname", "HTTPRequestRateLow",
				"alertstate", "pending",
				"group", "canary",
				"instance", "0",
				"job", "app-server",
				"severity", "critical",
			),
			Point: promql.Point{V: 1},
		},
		{
			Metric: labels.FromStrings(
				"__name__", "ALERTS",
				"alertname", "HTTPRequestRateLow",
				"alertstate", "pending",
				"group", "canary",
				"instance", "1",
				"job", "app-server",
				"severity", "critical",
			),
			Point: promql.Point{V: 1},
		},
		{
			Metric: labels.FromStrings(
				"__name__", "ALERTS",
				"alertname", "HTTPRequestRateLow",
				"alertstate", "firing",
				"group", "canary",
				"instance", "0",
				"job", "app-server",
				"severity", "critical",
			),
			Point: promql.Point{V: 1},
		},
		{
			Metric: labels.FromStrings(
				"__name__", "ALERTS",
				"alertname", "HTTPRequestRateLow",
				"alertstate", "firing",
				"group", "canary",
				"instance", "1",
				"job", "app-server",
				"severity", "critical",
			),
			Point: promql.Point{V: 1},
		},
	}

	baseTime := time.Unix(0, 0)

	var tests = []struct {
		time   time.Duration
		result promql.Vector
	}{
		{
			time:   0,
			result: result[:2],
		}, {
			time:   5 * time.Minute,
			result: result[2:],
		}, {
			time:   10 * time.Minute,
			result: result[2:3],
		},
		{
			time:   15 * time.Minute,
			result: nil,
		},
		{
			time:   20 * time.Minute,
			result: nil,
		},
		{
			time:   25 * time.Minute,
			result: result[:1],
		},
		{
			time:   30 * time.Minute,
			result: result[2:3],
		},
	}

	for i, test := range tests {
		t.Logf("case %d", i)

		evalTime := baseTime.Add(test.time)

		res, err := rule.Eval(suite.Context(), evalTime, EngineQueryFunc(suite.QueryEngine(), suite.Storage()), nil)
		testutil.Ok(t, err)

		for i := range test.result {
			test.result[i].T = timestamp.FromTime(evalTime)
		}
		testutil.Assert(t, len(test.result) == len(res), "%d. Number of samples in expected and actual output don't match (%d vs. %d)", i, len(test.result), len(res))

		sort.Slice(res, func(i, j int) bool {
			return labels.Compare(res[i].Metric, res[j].Metric) < 0
		})
		testutil.Equals(t, test.result, res)

		for _, aa := range rule.ActiveAlerts() {
			testutil.Assert(t, aa.Labels.Get(model.MetricNameLabel) == "", "%s label set on active alert: %s", model.MetricNameLabel, aa.Labels)
		}
	}
}

func TestStaleness(t *testing.T) {
	storage := testutil.NewStorage(t)
	defer storage.Close()
	engine := promql.NewEngine(nil, nil, 10, 10*time.Second)
	opts := &ManagerOptions{
		QueryFunc:  EngineQueryFunc(engine, storage),
		Appendable: storage,
		Context:    context.Background(),
		Logger:     log.NewNopLogger(),
	}

	expr, err := promql.ParseExpr("a + 1")
	testutil.Ok(t, err)
	rule := NewRecordingRule("a_plus_one", expr, labels.Labels{})
	group := NewGroup("default", "", time.Second, []Rule{rule}, opts)

	// A time series that has two samples and then goes stale.
	app, _ := storage.Appender()
	app.Add(labels.FromStrings(model.MetricNameLabel, "a"), 0, 1)
	app.Add(labels.FromStrings(model.MetricNameLabel, "a"), 1000, 2)
	app.Add(labels.FromStrings(model.MetricNameLabel, "a"), 2000, math.Float64frombits(value.StaleNaN))

	err = app.Commit()
	testutil.Ok(t, err)

	ctx := context.Background()

	// Execute 3 times, 1 second apart.
	group.Eval(ctx, time.Unix(0, 0))
	group.Eval(ctx, time.Unix(1, 0))
	group.Eval(ctx, time.Unix(2, 0))

	querier, err := storage.Querier(context.Background(), 0, 2000)
	testutil.Ok(t, err)
	defer querier.Close()

	matcher, err := labels.NewMatcher(labels.MatchEqual, model.MetricNameLabel, "a_plus_one")
	testutil.Ok(t, err)

	set, err := querier.Select(nil, matcher)
	testutil.Ok(t, err)

	samples, err := readSeriesSet(set)
	testutil.Ok(t, err)

	metric := labels.FromStrings(model.MetricNameLabel, "a_plus_one").String()
	metricSample, ok := samples[metric]

	testutil.Assert(t, ok, "Series %s not returned.", metric)
	testutil.Assert(t, value.IsStaleNaN(metricSample[2].V), "Appended second sample not as expected. Wanted: stale NaN Got: %x", math.Float64bits(metricSample[2].V))
	metricSample[2].V = 42 // reflect.DeepEqual cannot handle NaN.

	want := map[string][]promql.Point{
		metric: []promql.Point{{0, 2}, {1000, 3}, {2000, 42}},
	}

	testutil.Equals(t, want, samples)
}

// Convert a SeriesSet into a form useable with reflect.DeepEqual.
func readSeriesSet(ss storage.SeriesSet) (map[string][]promql.Point, error) {
	result := map[string][]promql.Point{}

	for ss.Next() {
		series := ss.At()

		points := []promql.Point{}
		it := series.Iterator()
		for it.Next() {
			t, v := it.At()
			points = append(points, promql.Point{T: t, V: v})
		}

		name := series.Labels().String()
		result[name] = points
	}
	return result, ss.Err()
}

func TestCopyState(t *testing.T) {
	oldGroup := &Group{
		rules: []Rule{
			NewAlertingRule("alert", nil, 0, nil, nil, nil),
			NewRecordingRule("rule1", nil, nil),
			NewRecordingRule("rule2", nil, nil),
			NewRecordingRule("rule3", nil, nil),
			NewRecordingRule("rule3", nil, nil),
		},
		seriesInPreviousEval: []map[string]labels.Labels{
			map[string]labels.Labels{"a": nil},
			map[string]labels.Labels{"r1": nil},
			map[string]labels.Labels{"r2": nil},
			map[string]labels.Labels{"r3a": nil},
			map[string]labels.Labels{"r3b": nil},
		},
		evaluationTime: time.Second,
	}
	oldGroup.rules[0].(*AlertingRule).active[42] = nil
	newGroup := &Group{
		rules: []Rule{
			NewRecordingRule("rule3", nil, nil),
			NewRecordingRule("rule3", nil, nil),
			NewRecordingRule("rule3", nil, nil),
			NewAlertingRule("alert", nil, 0, nil, nil, nil),
			NewRecordingRule("rule1", nil, nil),
			NewRecordingRule("rule4", nil, nil),
		},
		seriesInPreviousEval: make([]map[string]labels.Labels, 6),
	}
	newGroup.copyState(oldGroup)

	want := []map[string]labels.Labels{
		map[string]labels.Labels{"r3a": nil},
		map[string]labels.Labels{"r3b": nil},
		nil,
		map[string]labels.Labels{"a": nil},
		map[string]labels.Labels{"r1": nil},
		nil,
	}
	testutil.Equals(t, want, newGroup.seriesInPreviousEval)
	testutil.Equals(t, oldGroup.rules[0], newGroup.rules[3])
	testutil.Equals(t, oldGroup.evaluationTime, newGroup.evaluationTime)
}

func TestUpdate(t *testing.T) {
	expected := map[string]labels.Labels{
		"test": labels.FromStrings("name", "value"),
	}
	ruleManager := NewManager(&ManagerOptions{
		Context: context.Background(),
		Logger:  log.NewNopLogger(),
	})
	ruleManager.Run()

	err := ruleManager.Update(0, nil)
	testutil.Ok(t, err)
	for _, g := range ruleManager.groups {
		g.seriesInPreviousEval = []map[string]labels.Labels{
			expected,
		}
	}

	err = ruleManager.Update(0, nil)
	testutil.Ok(t, err)
	for _, g := range ruleManager.groups {
		for _, actual := range g.seriesInPreviousEval {
			testutil.Equals(t, expected, actual)
		}
	}
}
<|endoftext|>


<commit_before>package boltdb

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math"
	"sort"

	"github.com/Workiva/go-datastructures/augmentedtree"
	region "github.com/akhenakh/regionagogo"
	"github.com/akhenakh/regionagogo/geostore"
	"github.com/boltdb/bolt"
	"github.com/golang/geo/s2"
	"github.com/golang/protobuf/proto"
	lru "github.com/hashicorp/golang-lru"
)

const (
	loopBucket              = "loop"
	coverBucket             = "cover"
	earthCircumferenceMeter = 40075017
)

var (
	defaultCoverer = s2.RegionCoverer{MinLevel: 1, MaxLevel: 30, MaxCells: 8}
)

// GeoFenceBoltDB provides an in memory index and boltdb query engine for fences lookup
type GeoFenceBoltDB struct {
	augmentedtree.Tree
	*bolt.DB
	cache *lru.Cache
	debug bool
}

// GeoSearchOption used to pass options to NewGeoSearch
type GeoFenceBoltDBOption func(*geoFenceBoltDBOptions)

type geoFenceBoltDBOptions struct {
	maxCachedEntries uint
	debug            bool
	ro               bool
}

// WithCachedEntries enable an LRU cache default is disabled
func WithCachedEntries(maxCachedEntries uint) GeoFenceBoltDBOption {
	return func(o *geoFenceBoltDBOptions) {
		o.maxCachedEntries = maxCachedEntries
	}
}

// WithDebug enable debug
func WithDebug(debug bool) GeoFenceBoltDBOption {
	return func(o *geoFenceBoltDBOptions) {
		o.debug = debug
	}
}

// WithReadOnly enable read only db
func WithReadOnly(ro bool) GeoFenceBoltDBOption {
	return func(o *geoFenceBoltDBOptions) {
		o.ro = ro
	}
}

// NewGeoFenceBoltDB creates a new geo database, needs a writable path for BoltDB
func NewGeoFenceBoltDB(dbpath string, opts ...GeoFenceBoltDBOption) (*GeoFenceBoltDB, error) {
	var geoOpts geoFenceBoltDBOptions

	for _, opt := range opts {
		opt(&geoOpts)
	}

	db, err := bolt.Open(dbpath, 0600, &bolt.Options{ReadOnly: geoOpts.ro})
	if err != nil {
		return nil, err
	}

	// create bucket if we have write permission
	if !geoOpts.ro {
		if errdb := db.Update(func(tx *bolt.Tx) error {
			if _, errtx := tx.CreateBucketIfNotExists([]byte(loopBucket)); errtx != nil {
				return fmt.Errorf("create bucket: %s", errtx)
			}
			if _, errtx := tx.CreateBucketIfNotExists([]byte(coverBucket)); errtx != nil {
				return fmt.Errorf("create bucket: %s", errtx)
			}
			return nil
		}); errdb != nil {
			return nil, errdb
		}
	}

	gs := &GeoFenceBoltDB{
		Tree:  augmentedtree.New(1),
		DB:    db,
		debug: geoOpts.debug,
	}

	if geoOpts.maxCachedEntries != 0 {
		cache, err := lru.New(int(geoOpts.maxCachedEntries))
		if err != nil {
			return nil, err
		}
		gs.cache = cache
	}

	if err := gs.importGeoData(); err != nil {
		return nil, err
	}

	return gs, nil
}

// index indexes each cells of the cover and set its loopID
func (gs *GeoFenceBoltDB) index(fc *geostore.FenceCover, loopID uint64) {
	for _, cell := range fc.Cellunion {
		s2interval := &region.S2Interval{CellID: s2.CellID(cell)}
		intervals := gs.Query(s2interval)
		found := false

		if len(intervals) != 0 {
			for _, existInterval := range intervals {
				if existInterval.LowAtDimension(1) == s2interval.LowAtDimension(1) &&
					existInterval.HighAtDimension(1) == s2interval.HighAtDimension(1) {
					// update existing interval
					existS2Interval := existInterval.(*region.S2Interval)
					if gs.debug {
						log.Printf("added %d to existing interval %s containing %v", loopID, existS2Interval, existS2Interval.LoopIDs)
					}

					existS2Interval.LoopIDs = append(existS2Interval.LoopIDs, loopID)
					found = true
					break
				}
			}
		}

		if !found {
			// create new interval with current loop
			s2interval.LoopIDs = []uint64{loopID}
			gs.Add(s2interval)
			if gs.debug {
				log.Printf("added %v to new interval %s", s2interval.LoopIDs, s2interval)
			}
		}
	}
}

// importGeoData loads all existing cells into the segment tree
func (gs *GeoFenceBoltDB) importGeoData() error {
	var count int
	err := gs.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(coverBucket))
		cur := b.Cursor()

		// load the cell ranges into the tree
		var fc geostore.FenceCover
		for k, v := cur.First(); k != nil; k, v = cur.Next() {
			count++
			err := proto.Unmarshal(v, &fc)
			if err != nil {
				return err
			}
			if gs.debug {
				log.Println("read", fc.Cellunion)
			}

			// read back the loopID from the key
			var loopID uint64
			buf := bytes.NewBuffer(k)
			err = binary.Read(buf, binary.BigEndian, &loopID)
			if err != nil {
				return err
			}

			gs.index(&fc, loopID)
		}

		return nil
	})
	if err != nil {
		return err
	}

	if count != 0 {
		log.Println("loaded", count, "existing fences")
	} else {
		log.Println("initialized database")
	}

	return nil
}

// FenceByID returns a region from DB by its id
func (gs *GeoFenceBoltDB) FenceByID(loopID uint64) *region.Fence {
	// TODO: refactor as Fence ?
	if gs.cache != nil {
		if val, ok := gs.cache.Get(loopID); ok {
			r, _ := val.(*region.Fence)
			return r
		}
	}
	var rs *geostore.FenceStorage
	err := gs.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(loopBucket))
		v := b.Get(itob(loopID))

		if len(v) == 0 {
			return nil
		}

		var frs geostore.FenceStorage
		err := proto.Unmarshal(v, &frs)
		if err != nil {
			return err
		}
		rs = &frs
		return nil
	})
	if err != nil {
		return nil
	}
	r := region.NewFenceFromStorage(rs)
	if gs.cache != nil && r != nil {
		gs.cache.Add(loopID, r)
	}
	return r
}

// StubbingQuery returns the fence for the corresponding lat, lng point
func (gs *GeoFenceBoltDB) StubbingQuery(lat, lng float64, opts ...region.QueryOptionsFunc) (region.Fences, error) {
	// the CellID at L30
	q := s2.CellIDFromLatLng(s2.LatLngFromDegrees(lat, lng))

	i := &region.S2Interval{CellID: q}

	if gs.debug {
		log.Println("lookup", lat, lng, q)
	}
	r := gs.Tree.Query(i)

	var foundFence *region.Fence

	var queryOpts region.QueryOptions
	for _, opt := range opts {
		opt(&queryOpts)
	}

	var res []*region.Fence

	for _, itv := range r {
		sitv := itv.(*region.S2Interval)
		if gs.debug {
			log.Println("found possible solution", sitv, sitv.LoopIDs)
		}

		for _, loopID := range sitv.LoopIDs {
			fence := gs.FenceByID(loopID)
			if fence != nil && fence.Loop.ContainsPoint(q.Point()) {
				res = append(res, fence)
				if foundFence == nil {
					if gs.debug {
						log.Println("found matching fence", sitv, sitv.LoopIDs, fence.Loop.NumEdges())
					}
					foundFence = fence
					continue
				}

				// a fence can include a smaller fence
				// return only the one that is contained in the other if asked
				if !queryOpts.MultipleFences {
					// we take the 1st vertex of the fence.Loop if it is contained in previousLoop
					// region loop  is more precise
					if foundFence.Loop.ContainsPoint(fence.Loop.Vertex(0)) {
						foundFence = fence
					}
				}
			}

		}

	}

	if !queryOpts.MultipleFences && foundFence != nil {
		return []*region.Fence{foundFence}, nil
	}

	// Order fences by their size
	sort.Sort(sort.Reverse(region.BySize(res)))

	return res, nil
}

// RectQuery perform rectangular query ur upper right bl bottom left
func (gs *GeoFenceBoltDB) RectQuery(urlat, urlng, bllat, bllng float64, opts ...region.QueryOptionsFunc) (region.Fences, error) {
	rect := s2.RectFromLatLng(s2.LatLngFromDegrees(bllat, bllng))
	rect = rect.AddPoint(s2.LatLngFromDegrees(urlat, urlng))

	rc := &s2.RegionCoverer{MaxLevel: 30, MaxCells: 1}
	covering := rc.Covering(rect)
	if len(covering) != 1 {
		return nil, errors.New("impossible covering")
	}
	i := &region.S2Interval{CellID: covering[0]}
	r := gs.Tree.Query(i)

	fences := make(map[uint64]*region.Fence)

	for _, itv := range r {
		sitv := itv.(*region.S2Interval)
		for _, loopID := range sitv.LoopIDs {
			var region *region.Fence
			if v, ok := fences[loopID]; ok {
				region = v
			} else {
				region = gs.FenceByID(loopID)
				// testing the found loop is actually inside the rect
				// (since we are using only one large cover it may be outside)
				if rect.Contains(region.Loop.RectBound()) {
					fences[loopID] = region
				}
			}

		}
	}

	var res []*region.Fence
	for _, v := range fences {
		res = append(res, v)
	}
	return region.Fences(res), nil
}

// RadiusQuery is performing a radius query
func (gs *GeoFenceBoltDB) RadiusQuery(lat, lng, radius float64, opts ...region.QueryOptionsFunc) (region.Fences, error) {
	center := s2.PointFromLatLng(s2.LatLngFromDegrees(lat, lng))
	cap := s2.CapFromCenterArea(center, s2RadialAreaMeters(radius))
	covering := defaultCoverer.Covering(cap)

	var res []*region.Fence

	fencesIds := make(map[uint64]struct{})

	for _, cellID := range covering {
		i := &region.S2Interval{CellID: cellID}
		r := gs.Tree.Query(i)

		for _, itv := range r {
			sitv := itv.(*region.S2Interval)
			for _, loopID := range sitv.LoopIDs {
				fencesIds[loopID] = struct{}{}
			}
		}
	}

	for k := range fencesIds {
		fence := gs.FenceByID(k)
		if fence != nil {
			res = append(res, fence)
		}
	}

	return res, nil
}

func s2RadialAreaMeters(radius float64) float64 {
	r := (radius / earthCircumferenceMeter) * math.Pi * 2
	return (math.Pi * r * r)
}

// StoreFence stores a fence into the database and load its index in memory
func (gs *GeoFenceBoltDB) StoreFence(fs *geostore.FenceStorage, cover []uint64) error {
	return gs.Update(func(tx *bolt.Tx) error {
		loopB := tx.Bucket([]byte(loopBucket))
		coverBucket := tx.Bucket([]byte(coverBucket))

		loopID, err := loopB.NextSequence()
		if err != nil {
			return err
		}

		buf, err := proto.Marshal(fs)
		if err != nil {
			return err
		}

		if gs.debug {
			log.Println("inserted", loopID, fs.Data, cover)
		}

		// convert our loopID to bigendian to be used as key
		k := itob(loopID)

		err = loopB.Put(k, buf)
		if err != nil {
			return err
		}

		// inserting into cover index using the same key
		fc := &geostore.FenceCover{Cellunion: cover}
		bufc, err := proto.Marshal(fc)
		if err != nil {
			return err
		}

		// also load into memory
		gs.index(fc, loopID)

		return coverBucket.Put(k, bufc)
	})
}

// itob returns an 8-byte big endian representation of v.
func itob(v uint64) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, v)
	return b
}
<commit_msg>added NewGeoFenceIdx use a passed  *bolt.DB<commit_after>package boltdb

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math"
	"sort"

	"github.com/Workiva/go-datastructures/augmentedtree"
	region "github.com/akhenakh/regionagogo"
	"github.com/akhenakh/regionagogo/geostore"
	"github.com/boltdb/bolt"
	"github.com/golang/geo/s2"
	"github.com/golang/protobuf/proto"
	lru "github.com/hashicorp/golang-lru"
)

const (
	loopBucket              = "loop"
	coverBucket             = "cover"
	earthCircumferenceMeter = 40075017
)

var (
	defaultCoverer = s2.RegionCoverer{MinLevel: 1, MaxLevel: 30, MaxCells: 8}
)

// GeoFenceBoltDB provides an in memory index and boltdb query engine for fences lookup
type GeoFenceBoltDB struct {
	augmentedtree.Tree
	*bolt.DB
	cache *lru.Cache
	debug bool
	ro    bool
}

// GeoSearchOption used to pass options to NewGeoSearch
type GeoFenceBoltDBOption func(*geoFenceBoltDBOptions)

type geoFenceBoltDBOptions struct {
	maxCachedEntries uint
	debug            bool
	ro               bool
}

// WithCachedEntries enable an LRU cache default is disabled
func WithCachedEntries(maxCachedEntries uint) GeoFenceBoltDBOption {
	return func(o *geoFenceBoltDBOptions) {
		o.maxCachedEntries = maxCachedEntries
	}
}

// WithDebug enable debug
func WithDebug(debug bool) GeoFenceBoltDBOption {
	return func(o *geoFenceBoltDBOptions) {
		o.debug = debug
	}
}

// WithReadOnly enable read only db
func WithReadOnly(ro bool) GeoFenceBoltDBOption {
	return func(o *geoFenceBoltDBOptions) {
		o.ro = ro
	}
}

// NewGeoFenceBoltDB creates or reopen a bolt geo database
func NewGeoFenceBoltDB(dbpath string, opts ...GeoFenceBoltDBOption) (*GeoFenceBoltDB, error) {
	var geoOpts geoFenceBoltDBOptions

	for _, opt := range opts {
		opt(&geoOpts)
	}

	db, err := bolt.Open(dbpath, 0600, &bolt.Options{ReadOnly: geoOpts.ro})
	if err != nil {
		return nil, err
	}

	return NewGeoFenceIdx(db, opts...)
}

// NewGeoFenceIdx a geo index over a BoltDB storage
func NewGeoFenceIdx(db *bolt.DB, opts ...GeoFenceBoltDBOption) (*GeoFenceBoltDB, error) {
	var geoOpts geoFenceBoltDBOptions

	for _, opt := range opts {
		opt(&geoOpts)
	}

	// create bucket if we have write permission
	if !geoOpts.ro {
		if errdb := db.Update(func(tx *bolt.Tx) error {
			if _, errtx := tx.CreateBucketIfNotExists([]byte(loopBucket)); errtx != nil {
				return fmt.Errorf("create bucket: %s", errtx)
			}
			if _, errtx := tx.CreateBucketIfNotExists([]byte(coverBucket)); errtx != nil {
				return fmt.Errorf("create bucket: %s", errtx)
			}
			return nil
		}); errdb != nil {
			return nil, errdb
		}
	}

	gs := &GeoFenceBoltDB{
		Tree:  augmentedtree.New(1),
		DB:    db,
		debug: geoOpts.debug,
		ro:    geoOpts.ro,
	}

	if geoOpts.maxCachedEntries != 0 {
		cache, err := lru.New(int(geoOpts.maxCachedEntries))
		if err != nil {
			return nil, err
		}
		gs.cache = cache
	}

	if err := gs.importGeoData(); err != nil {
		return nil, err
	}

	return gs, nil
}

// index indexes each cells of the cover and set its loopID
func (gs *GeoFenceBoltDB) index(fc *geostore.FenceCover, loopID uint64) {
	for _, cell := range fc.Cellunion {
		s2interval := &region.S2Interval{CellID: s2.CellID(cell)}
		intervals := gs.Query(s2interval)
		found := false

		if len(intervals) != 0 {
			for _, existInterval := range intervals {
				if existInterval.LowAtDimension(1) == s2interval.LowAtDimension(1) &&
					existInterval.HighAtDimension(1) == s2interval.HighAtDimension(1) {
					// update existing interval
					existS2Interval := existInterval.(*region.S2Interval)
					if gs.debug {
						log.Printf("added %d to existing interval %s containing %v", loopID, existS2Interval, existS2Interval.LoopIDs)
					}

					existS2Interval.LoopIDs = append(existS2Interval.LoopIDs, loopID)
					found = true
					break
				}
			}
		}

		if !found {
			// create new interval with current loop
			s2interval.LoopIDs = []uint64{loopID}
			gs.Add(s2interval)
			if gs.debug {
				log.Printf("added %v to new interval %s", s2interval.LoopIDs, s2interval)
			}
		}
	}
}

// importGeoData loads all existing cells into the segment tree
func (gs *GeoFenceBoltDB) importGeoData() error {
	var count int
	err := gs.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(coverBucket))
		cur := b.Cursor()

		// load the cell ranges into the tree
		var fc geostore.FenceCover
		for k, v := cur.First(); k != nil; k, v = cur.Next() {
			count++
			err := proto.Unmarshal(v, &fc)
			if err != nil {
				return err
			}
			if gs.debug {
				log.Println("read", fc.Cellunion)
			}

			// read back the loopID from the key
			var loopID uint64
			buf := bytes.NewBuffer(k)
			err = binary.Read(buf, binary.BigEndian, &loopID)
			if err != nil {
				return err
			}

			gs.index(&fc, loopID)
		}

		return nil
	})
	if err != nil {
		return err
	}

	if count != 0 {
		log.Println("loaded", count, "existing fences")
	} else {
		log.Println("initialized database")
	}

	return nil
}

// FenceByID returns a region from DB by its id
func (gs *GeoFenceBoltDB) FenceByID(loopID uint64) *region.Fence {
	// TODO: refactor as Fence ?
	if gs.cache != nil {
		if val, ok := gs.cache.Get(loopID); ok {
			r, _ := val.(*region.Fence)
			return r
		}
	}
	var rs *geostore.FenceStorage
	err := gs.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(loopBucket))
		v := b.Get(itob(loopID))

		if len(v) == 0 {
			return nil
		}

		var frs geostore.FenceStorage
		err := proto.Unmarshal(v, &frs)
		if err != nil {
			return err
		}
		rs = &frs
		return nil
	})
	if err != nil {
		return nil
	}
	r := region.NewFenceFromStorage(rs)
	if gs.cache != nil && r != nil {
		gs.cache.Add(loopID, r)
	}
	return r
}

// StubbingQuery returns the fence for the corresponding lat, lng point
func (gs *GeoFenceBoltDB) StubbingQuery(lat, lng float64, opts ...region.QueryOptionsFunc) (region.Fences, error) {
	// the CellID at L30
	q := s2.CellIDFromLatLng(s2.LatLngFromDegrees(lat, lng))

	i := &region.S2Interval{CellID: q}

	if gs.debug {
		log.Println("lookup", lat, lng, q)
	}
	r := gs.Tree.Query(i)

	var foundFence *region.Fence

	var queryOpts region.QueryOptions
	for _, opt := range opts {
		opt(&queryOpts)
	}

	var res []*region.Fence

	for _, itv := range r {
		sitv := itv.(*region.S2Interval)
		if gs.debug {
			log.Println("found possible solution", sitv, sitv.LoopIDs)
		}

		for _, loopID := range sitv.LoopIDs {
			fence := gs.FenceByID(loopID)
			if fence != nil && fence.Loop.ContainsPoint(q.Point()) {
				res = append(res, fence)
				if foundFence == nil {
					if gs.debug {
						log.Println("found matching fence", sitv, sitv.LoopIDs, fence.Loop.NumEdges())
					}
					foundFence = fence
					continue
				}

				// a fence can include a smaller fence
				// return only the one that is contained in the other if asked
				if !queryOpts.MultipleFences {
					// we take the 1st vertex of the fence.Loop if it is contained in previousLoop
					// region loop  is more precise
					if foundFence.Loop.ContainsPoint(fence.Loop.Vertex(0)) {
						foundFence = fence
					}
				}
			}

		}

	}

	if !queryOpts.MultipleFences && foundFence != nil {
		return []*region.Fence{foundFence}, nil
	}

	// Order fences by their size
	sort.Sort(sort.Reverse(region.BySize(res)))

	return res, nil
}

// RectQuery perform rectangular query ur upper right bl bottom left
func (gs *GeoFenceBoltDB) RectQuery(urlat, urlng, bllat, bllng float64, opts ...region.QueryOptionsFunc) (region.Fences, error) {
	rect := s2.RectFromLatLng(s2.LatLngFromDegrees(bllat, bllng))
	rect = rect.AddPoint(s2.LatLngFromDegrees(urlat, urlng))

	rc := &s2.RegionCoverer{MaxLevel: 30, MaxCells: 1}
	covering := rc.Covering(rect)
	if len(covering) != 1 {
		return nil, errors.New("impossible covering")
	}
	i := &region.S2Interval{CellID: covering[0]}
	r := gs.Tree.Query(i)

	fences := make(map[uint64]*region.Fence)

	for _, itv := range r {
		sitv := itv.(*region.S2Interval)
		for _, loopID := range sitv.LoopIDs {
			var region *region.Fence
			if v, ok := fences[loopID]; ok {
				region = v
			} else {
				region = gs.FenceByID(loopID)
				// testing the found loop is actually inside the rect
				// (since we are using only one large cover it may be outside)
				if rect.Contains(region.Loop.RectBound()) {
					fences[loopID] = region
				}
			}

		}
	}

	var res []*region.Fence
	for _, v := range fences {
		res = append(res, v)
	}
	return region.Fences(res), nil
}

// RadiusQuery is performing a radius query
func (gs *GeoFenceBoltDB) RadiusQuery(lat, lng, radius float64, opts ...region.QueryOptionsFunc) (region.Fences, error) {
	center := s2.PointFromLatLng(s2.LatLngFromDegrees(lat, lng))
	cap := s2.CapFromCenterArea(center, s2RadialAreaMeters(radius))
	covering := defaultCoverer.Covering(cap)

	var res []*region.Fence

	fencesIds := make(map[uint64]struct{})

	for _, cellID := range covering {
		i := &region.S2Interval{CellID: cellID}
		r := gs.Tree.Query(i)

		for _, itv := range r {
			sitv := itv.(*region.S2Interval)
			for _, loopID := range sitv.LoopIDs {
				fencesIds[loopID] = struct{}{}
			}
		}
	}

	for k := range fencesIds {
		fence := gs.FenceByID(k)
		if fence != nil {
			res = append(res, fence)
		}
	}

	return res, nil
}

func s2RadialAreaMeters(radius float64) float64 {
	r := (radius / earthCircumferenceMeter) * math.Pi * 2
	return (math.Pi * r * r)
}

// StoreFence stores a fence into the database and load its index in memory
func (gs *GeoFenceBoltDB) StoreFence(fs *geostore.FenceStorage, cover []uint64) error {
	if gs.ro {
		return errors.New("db is in read only mode")
	}
	return gs.Update(func(tx *bolt.Tx) error {
		loopB := tx.Bucket([]byte(loopBucket))
		coverBucket := tx.Bucket([]byte(coverBucket))

		loopID, err := loopB.NextSequence()
		if err != nil {
			return err
		}

		buf, err := proto.Marshal(fs)
		if err != nil {
			return err
		}

		if gs.debug {
			log.Println("inserted", loopID, fs.Data, cover)
		}

		// convert our loopID to bigendian to be used as key
		k := itob(loopID)

		err = loopB.Put(k, buf)
		if err != nil {
			return err
		}

		// inserting into cover index using the same key
		fc := &geostore.FenceCover{Cellunion: cover}
		bufc, err := proto.Marshal(fc)
		if err != nil {
			return err
		}

		// also load into memory
		gs.index(fc, loopID)

		return coverBucket.Put(k, bufc)
	})
}

// itob returns an 8-byte big endian representation of v.
func itob(v uint64) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, v)
	return b
}
<|endoftext|>


<commit_before>package wavreader

import (
	"fmt"
	"io"
)

func Convert(r *Reader, channels, rate, bits int) (*Reader, error) {
	// Fast path: don't convert anything if not absolutely necessary
	if r.Channels == channels && r.SampleRate == rate && r.BitsPerSample == bits {
		//return r, nil
	}

	if channels < 1 {
		return nil, fmt.Errorf("need at least 1 output channel")
	}

	rv, wri := Pipe()
	rv.Channels = channels
	rv.SampleRate = rate
	rv.BitsPerSample = bits

	go doConversion(wri, r, channels, rate, bits)

	return rv, nil
}

func doConversion(wri *io.PipeWriter, r *Reader, channels, rate, bits int) {
	// Samples at a time. We're reading the input stream in chunks of, say, 10ms.
	saatIn, saatOut := (r.SampleRate+99)/100, 1+(rate+99)/100

	// Bytes per sample
	Bin, Bout := (r.BitsPerSample+7)/8, (bits+7)/8

	// Create buffers
	bufIn := make([]byte, r.Channels*saatIn*Bin)
	bufChan := make([][]byte, channels)
	bufRate := make([]*rateConverter, channels)
	bufBits := make([][]byte, channels)
	bufOut := make([]byte, r.Channels*saatOut*Bout)

	for i, _ := range bufChan {
		bufChan[i] = make([]byte, saatIn*Bin)

		if r.SampleRate == rate {
			// Optimization: since we're not converting, just reuse the last output buffer
			bufRate[i] = &rateConverter{Output: bufChan[i]}
		} else {
			bufRate[i] = &rateConverter{Output: make([]byte, saatOut*Bin)}
		}

		if r.BitsPerSample == bits {
			// Optimization: since we're not converting, just reuse the last output buffer
			bufBits[i] = bufRate[i].Output
		} else {
			bufBits[i] = make([]byte, saatOut*Bout)
		}
	}

	var err error
	var nMono, nRate, nBits int

	for {
		n, errRead := io.ReadFull(r, bufIn)
		if n == 0 && errRead != nil {
			wri.CloseWithError(errRead)
			return
		}

		nMono, err = monoChannels(bufChan, bufIn[:n], r, Bin)
		if err != nil {
			wri.CloseWithError(err)
			return
		}

		for i, rc := range bufRate {
			nRate, err = rc.convert(bufChan[i][:nMono], r, Bin, rate)
			if err != nil {
				wri.CloseWithError(err)
				return
			}
		}

		for i, ch := range bufRate {
			nBits, err = convertBits(bufBits[i], ch.Output[:nRate], r, Bin, Bout, bits)
			if err != nil {
				wri.CloseWithError(err)
				return
			}
		}

		n, err = interleave(bufOut, bufBits, nBits, Bout)
		if err != nil {
			wri.CloseWithError(err)
			return
		}

		_, errWrite := wri.Write(bufOut[:n])
		if errWrite != nil {
			wri.CloseWithError(errWrite)
			return
		}
		if errRead != nil {
			wri.CloseWithError(errRead)
			return
		}
	}
}

func monoChannels(out [][]byte, in []byte, r *Reader, Bin int) (int, error) {
	if r.Channels == 1 {
		// Copy mono input data to all output channels
		for _, b := range out {
			copy(b, in)
		}
		return len(in), nil
	} else if r.Channels == len(out) {
		// Spice interleaved sample data into per-channel buffers
		j := 0
		for j*r.Channels < len(in) {
			for i, ch := range out {
				ioff := j*r.Channels + i*Bin
				copy(ch[j:], in[ioff:ioff+Bin])
			}
			j += Bin
		}
		return j, nil
	} else {
		return 0, fmt.Errorf("n:m channel mapping is not implemented")
	}
}

type rateConverter struct {
	Output  []byte
	skipped int
}

func (rc *rateConverter) convert(in []byte, r *Reader, Bin, rate int) (int, error) {
	if r.SampleRate == rate {
		// We've caught this case by reusing buffers
		return len(in), nil
	}

	if r.SampleRate > rate && (r.SampleRate%rate) == 0 {
		// Fast path: the source sample rate is a multiple of the target rate

		// Output a sample every c samples
		c := r.SampleRate / rate

		i, n := 0, 0

		for i < len(in) {
			if rc.skipped == 0 {
				copy(rc.Output[n:], in[i:i+Bin])
				n += Bin
			}

			rc.skipped = (rc.skipped + 1) % c
			i += Bin
		}

		return n, nil
	}

	return 0, fmt.Errorf("sample rate conversion is not implemented")
}

func convertBits(out []byte, in []byte, r *Reader, Bin, Bout, bits int) (int, error) {
	if r.BitsPerSample == bits {
		// We've caught this case by reusing buffers
		return len(in), nil
	}

	// FIXME: handle bit lengths that aren't multiples of 8 bits. (Do those exist?)

	if r.BitsPerSample > bits {
		// Discard input bits in the output stream
		i, n := 0, 0

		for i < len(in) {
			// HACK: Handle conversion from signed 16-bit to unsigned 8-bit
			if bits == 8 {
				out[n] = 128 + in[i+Bin-1]
			} else {
				copy(out[n:], in[i+Bin-Bout:i+Bin])
			}

			i += Bin
			n += Bout
		}

		return n, nil
	} else {
		// Pad the input stream to fill the output stream
		i, n := 0, 0

		// HACK: Handle conversion from unsigned 8-bit to signed 16-bit
		var offset uint8 = 0
		if Bin == 1 {
			offset = 128
		}

		for i < len(in) {
			for j := 0; j < Bout; j++ {
				out[n+Bout-j-1] = in[i+(Bin-(j%Bin)-1)] + offset
				// TODO: Add dithering, or some other method of hiding rounding errors
			}

			i += Bin
			n += Bout
		}

		return n, nil
	}
}

func interleave(out []byte, in [][]byte, length int, Bout int) (int, error) {
	c := len(in)
	j := 0
	for j < length {
		for i, ch := range in {
			ooff := j*c + i*Bout
			copy(out[ooff:], ch[j:j+Bout])
		}
		j += Bout
	}
	return j * c, nil
}
<commit_msg>Add sample rate upscaling with squarewaves<commit_after>package wavreader

import (
	"fmt"
	"io"
)

func Convert(r *Reader, channels, rate, bits int) (*Reader, error) {
	// Fast path: don't convert anything if not absolutely necessary
	if r.Channels == channels && r.SampleRate == rate && r.BitsPerSample == bits {
		//return r, nil
	}

	if channels < 1 {
		return nil, fmt.Errorf("need at least 1 output channel")
	}

	rv, wri := Pipe()
	rv.Channels = channels
	rv.SampleRate = rate
	rv.BitsPerSample = bits

	go doConversion(wri, r, channels, rate, bits)

	return rv, nil
}

func doConversion(wri *io.PipeWriter, r *Reader, channels, rate, bits int) {
	// Samples at a time. We're reading the input stream in chunks of, say, 10ms.
	saatIn, saatOut := (r.SampleRate+99)/100, 1+(rate+99)/100

	// Bytes per sample
	Bin, Bout := (r.BitsPerSample+7)/8, (bits+7)/8

	// Create buffers
	bufIn := make([]byte, r.Channels*saatIn*Bin)
	bufChan := make([][]byte, channels)
	bufRate := make([]*rateConverter, channels)
	bufBits := make([][]byte, channels)
	bufOut := make([]byte, r.Channels*saatOut*Bout)

	for i, _ := range bufChan {
		bufChan[i] = make([]byte, saatIn*Bin)

		if r.SampleRate == rate {
			// Optimization: since we're not converting, just reuse the last output buffer
			bufRate[i] = &rateConverter{Output: bufChan[i]}
		} else {
			bufRate[i] = &rateConverter{Output: make([]byte, saatOut*Bin)}
		}

		if r.BitsPerSample == bits {
			// Optimization: since we're not converting, just reuse the last output buffer
			bufBits[i] = bufRate[i].Output
		} else {
			bufBits[i] = make([]byte, saatOut*Bout)
		}
	}

	var err error
	var nMono, nRate, nBits int

	for {
		n, errRead := io.ReadFull(r, bufIn)
		if n == 0 && errRead != nil {
			wri.CloseWithError(errRead)
			return
		}

		nMono, err = monoChannels(bufChan, bufIn[:n], r, Bin)
		if err != nil {
			wri.CloseWithError(err)
			return
		}

		for i, rc := range bufRate {
			nRate, err = rc.convert(bufChan[i][:nMono], r, Bin, rate)
			if err != nil {
				wri.CloseWithError(err)
				return
			}
		}

		for i, ch := range bufRate {
			nBits, err = convertBits(bufBits[i], ch.Output[:nRate], r, Bin, Bout, bits)
			if err != nil {
				wri.CloseWithError(err)
				return
			}
		}

		n, err = interleave(bufOut, bufBits, nBits, Bout)
		if err != nil {
			wri.CloseWithError(err)
			return
		}

		_, errWrite := wri.Write(bufOut[:n])
		if errWrite != nil {
			wri.CloseWithError(errWrite)
			return
		}
		if errRead != nil {
			wri.CloseWithError(errRead)
			return
		}
	}
}

func monoChannels(out [][]byte, in []byte, r *Reader, Bin int) (int, error) {
	if r.Channels == 1 {
		// Copy mono input data to all output channels
		for _, b := range out {
			copy(b, in)
		}
		return len(in), nil
	} else if r.Channels == len(out) {
		// Spice interleaved sample data into per-channel buffers
		j := 0
		for j*r.Channels < len(in) {
			for i, ch := range out {
				ioff := j*r.Channels + i*Bin
				copy(ch[j:], in[ioff:ioff+Bin])
			}
			j += Bin
		}
		return j, nil
	} else {
		return 0, fmt.Errorf("n:m channel mapping is not implemented")
	}
}

type rateConverter struct {
	Output  []byte
	skipped int
}

func (rc *rateConverter) convert(in []byte, r *Reader, Bin, rate int) (int, error) {
	if r.SampleRate == rate {
		// We've caught this case by reusing buffers
		return len(in), nil
	}

	if r.SampleRate > rate && (r.SampleRate%rate) == 0 {
		// Fast path: the source sample rate is a multiple of the target rate

		// Output a sample every c samples
		c := r.SampleRate / rate

		i, n := 0, 0

		for i < len(in) {
			if rc.skipped == 0 {
				copy(rc.Output[n:], in[i:i+Bin])
				n += Bin
			}

			rc.skipped = (rc.skipped + 1) % c
			i += Bin
		}

		return n, nil
	} else if r.SampleRate < rate && (rate%r.SampleRate) == 0 {
		// Another fast path: approximate the upscaled version with squarewaves

		// Repeat each sample c times
		c := rate / r.SampleRate

		i := 0
		for i < len(in) {
			for j := 0; j < c; j++ {
				copy(rc.Output[i*c+j*Bin:], in[i:i+Bin])
			}
			i += Bin
		}

		return len(in) * c, nil
	}

	return 0, fmt.Errorf("sample rate conversion is not implemented")
}

func convertBits(out []byte, in []byte, r *Reader, Bin, Bout, bits int) (int, error) {
	if r.BitsPerSample == bits {
		// We've caught this case by reusing buffers
		return len(in), nil
	}

	// FIXME: handle bit lengths that aren't multiples of 8 bits. (Do those exist?)

	if r.BitsPerSample > bits {
		// Discard input bits in the output stream
		i, n := 0, 0

		for i < len(in) {
			// HACK: Handle conversion from signed 16-bit to unsigned 8-bit
			if bits == 8 {
				out[n] = 128 + in[i+Bin-1]
			} else {
				copy(out[n:], in[i+Bin-Bout:i+Bin])
			}

			i += Bin
			n += Bout
		}

		return n, nil
	} else {
		// Pad the input stream to fill the output stream
		i, n := 0, 0

		// HACK: Handle conversion from unsigned 8-bit to signed 16-bit
		var offset uint8 = 0
		if Bin == 1 {
			offset = 128
		}

		for i < len(in) {
			for j := 0; j < Bout; j++ {
				out[n+Bout-j-1] = in[i+(Bin-(j%Bin)-1)] + offset
				// TODO: Add dithering, or some other method of hiding rounding errors
			}

			i += Bin
			n += Bout
		}

		return n, nil
	}
}

func interleave(out []byte, in [][]byte, length int, Bout int) (int, error) {
	c := len(in)
	j := 0
	for j < length {
		for i, ch := range in {
			ooff := j*c + i*Bout
			copy(out[ooff:], ch[j:j+Bout])
		}
		j += Bout
	}
	return j * c, nil
}
<|endoftext|>


<commit_before>package libbusyna

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

// ExtractShellCreate creates the shell help script that is executed as the
// shell itself.
func ExtractShellCreate(rcfile *os.File) *os.File {
	shfile, err := ioutil.TempFile("", "busyna-shell-")
	if err != nil {
		log.Fatal(err)
	}
	shfile.Chmod(0755)
	for _, l := range []string{
		`#!/bin/sh`,
		``,
		`shift # get rid of -c`,
		fmt.Sprintf(`SUBMAKE=$(echo "$@" | sed 's@%s@make SHELL=%s MAKE=%s@')`, rcfile.Name(), shfile.Name(), rcfile.Name()),
		`if [ "$SUBMAKE" != "$(echo "$@")" ]; then # recursive make call`,
		"\t/bin/sh -c \"$SUBMAKE\"",
		"\texit $?",
		`fi`,
		`(`,
		`# get rid of standard environment`,
		`unset MAKE`,
		`unset MAKEOVERRIDES`,
		`unset MAKEFLAGS`,
		`unset MFLAGS`,
		`unset MAKELEVEL`,
	} {
		shfile.WriteString(l)
		shfile.WriteString("\n")
	}
	for _, e := range os.Environ() {
		ev := strings.Split(e, "=")
		shfile.WriteString(fmt.Sprintf("unset %s\n", ev[0]))
	}
	shfile.WriteString(fmt.Sprintf("/usr/bin/env | /bin/sed 's@\\(^[^=]\\+\\)=\\(.*\\)$@\\1='\"'\"'\\2'\"'\"'@' >> %s\n", rcfile.Name()))
	shfile.WriteString(")\n")
	shfile.WriteString(fmt.Sprintf("echo cd \"$PWD\" >> %s\n\n", rcfile.Name()))
	shfile.WriteString(fmt.Sprintf("echo -n '(' >> %s\n\n", rcfile.Name()))
	shfile.WriteString(fmt.Sprintf("echo \"$@\" | sed 's@\\\\$@@' | tr -d '\\n' >> %s\n\n", rcfile.Name()))
	shfile.WriteString(fmt.Sprintf("echo ')' >> %s\n\n", rcfile.Name()))
	//shfile.WriteString("echo exec /bin/sh -c \"$@\" >&2\n\n")
	shfile.WriteString("exec /bin/sh -c \"$@\"\n\n")
	shfile.Close() // Must close to avoid "text file busy"
	return shfile
}

// ExtractMake creates the shell script and the Makefile that are used to
// create a busyna.rc from an existing Makefile
func ExtractMake(outputfile string) {
	outputfileabs, err := filepath.Abs(outputfile)
	if err != nil {
		log.Fatal(err)
	}
	rcfile, err := ioutil.TempFile(filepath.Dir(outputfileabs), "busyna.rc-")
	if err != nil {
		log.Fatal(err)
	}
	defer os.Remove(rcfile.Name())
	rcfile.Close()

	shfile := ExtractShellCreate(rcfile)
	defer os.Remove(shfile.Name())

	cmd := exec.Command("make", "-B", "-j1", fmt.Sprintf("SHELL=%s", shfile.Name()), fmt.Sprintf("MAKE=%s", rcfile.Name()))
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		log.Fatal(fmt.Sprintf("%s while running make with shell %s to fill %s - keeping files", err, shfile.Name(), rcfile.Name()))
	}

	DeployRcCmd(RcParse(rcfile.Name(), ChanFromFile(rcfile.Name())), outputfile)
}
<commit_msg>Remove new env vars from new gnu make<commit_after>package libbusyna

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

// ExtractShellCreate creates the shell help script that is executed as the
// shell itself.
func ExtractShellCreate(rcfile *os.File) *os.File {
	shfile, err := ioutil.TempFile("", "busyna-shell-")
	if err != nil {
		log.Fatal(err)
	}
	shfile.Chmod(0755)
	for _, l := range []string{
		`#!/bin/sh`,
		``,
		`shift # get rid of -c`,
		fmt.Sprintf(`SUBMAKE=$(echo "$@" | sed 's@%s@make SHELL=%s MAKE=%s@')`, rcfile.Name(), shfile.Name(), rcfile.Name()),
		`if [ "$SUBMAKE" != "$(echo "$@")" ]; then # recursive make call`,
		"\t/bin/sh -c \"$SUBMAKE\"",
		"\texit $?",
		`fi`,
		`(`,
		`# get rid of standard environment`,
		`unset MAKE`,
		`unset MAKEOVERRIDES`,
		`unset MAKEFLAGS`,
		`unset MFLAGS`,
		`unset MAKELEVEL`,
		`unset MAKE_TERMOUT`,
		`unset MAKE_TERMERR`,
	} {
		shfile.WriteString(l)
		shfile.WriteString("\n")
	}
	for _, e := range os.Environ() {
		ev := strings.Split(e, "=")
		shfile.WriteString(fmt.Sprintf("unset %s\n", ev[0]))
	}
	shfile.WriteString(fmt.Sprintf("/usr/bin/env | /bin/sed 's@\\(^[^=]\\+\\)=\\(.*\\)$@\\1='\"'\"'\\2'\"'\"'@' >> %s\n", rcfile.Name()))
	shfile.WriteString(")\n")
	shfile.WriteString(fmt.Sprintf("echo cd \"$PWD\" >> %s\n\n", rcfile.Name()))
	shfile.WriteString(fmt.Sprintf("echo -n '(' >> %s\n\n", rcfile.Name()))
	shfile.WriteString(fmt.Sprintf("echo \"$@\" | sed 's@\\\\$@@' | tr -d '\\n' >> %s\n\n", rcfile.Name()))
	shfile.WriteString(fmt.Sprintf("echo ')' >> %s\n\n", rcfile.Name()))
	//shfile.WriteString("echo exec /bin/sh -c \"$@\" >&2\n\n")
	shfile.WriteString("exec /bin/sh -c \"$@\"\n\n")
	shfile.Close() // Must close to avoid "text file busy"
	return shfile
}

// ExtractMake creates the shell script and the Makefile that are used to
// create a busyna.rc from an existing Makefile
func ExtractMake(outputfile string) {
	outputfileabs, err := filepath.Abs(outputfile)
	if err != nil {
		log.Fatal(err)
	}
	rcfile, err := ioutil.TempFile(filepath.Dir(outputfileabs), "busyna.rc-")
	if err != nil {
		log.Fatal(err)
	}
	defer os.Remove(rcfile.Name())
	rcfile.Close()

	shfile := ExtractShellCreate(rcfile)
	defer os.Remove(shfile.Name())

	cmd := exec.Command("make", "-B", "-j1", fmt.Sprintf("SHELL=%s", shfile.Name()), fmt.Sprintf("MAKE=%s", rcfile.Name()))
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		log.Fatal(fmt.Sprintf("%s while running make with shell %s to fill %s - keeping files", err, shfile.Name(), rcfile.Name()))
	}

	DeployRcCmd(RcParse(rcfile.Name(), ChanFromFile(rcfile.Name())), outputfile)
}
<|endoftext|>


<commit_before>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"errors"
	"fmt"
	"sync"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
)

type extendedIdentify struct {
	behavior keybase1.TLFIdentifyBehavior

	// lock guards userBreaks and tlfBreaks
	lock       sync.Mutex
	userBreaks chan keybase1.TLFIdentifyFailure
	tlfBreaks  *keybase1.TLFBreak
}

func (ei *extendedIdentify) userBreak(username libkb.NormalizedUsername, uid keybase1.UID, breaks *keybase1.IdentifyTrackBreaks) {
	if ei.userBreaks == nil {
		return
	}

	ei.userBreaks <- keybase1.TLFIdentifyFailure{
		Breaks: breaks,
		User: keybase1.User{
			Uid:      uid,
			Username: string(username),
		},
	}
}

func (ei *extendedIdentify) makeTlfBreaksIfNeeded(
	ctx context.Context, numUserInTlf int) error {
	if ei.userBreaks == nil {
		return nil
	}

	ei.lock.Lock()
	defer ei.lock.Unlock()

	b := &keybase1.TLFBreak{}
	for i := 0; i < numUserInTlf; i++ {
		select {
		case ub, ok := <-ei.userBreaks:
			if !ok {
				return errors.New("makeTlfBreaksIfNeeded called on extendedIdentify" +
					" with closed userBreaks channel.")
			}
			if ub.Breaks != nil {
				b.Breaks = append(b.Breaks, ub)
			}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	ei.tlfBreaks = b

	return nil
}

// getTlfBreakOrBust returns a keybase1.TLFBreak. This should only be called
// for behavior.WarningInsteadOfErrorOnBrokenTracks() == true, and after
// makeTlfBreaksIfNeeded is called, to make sure user proof breaks get
// populated in GUI mode.
//
// If called otherwise, we don't panic here anymore, since we can't panic on
// nil ei.tlfBreaks. The reason is if a previous successful identify has
// already happened recently, it could cause this identify to be skipped, which
// means ei.tlfBreaks is never populated. In this case, it's safe to return an
// empty keybase1.TLFBreak.
func (ei *extendedIdentify) getTlfBreakAndClose() keybase1.TLFBreak {
	ei.lock.Lock()
	defer ei.lock.Unlock()

	if ei.userBreaks != nil {
		close(ei.userBreaks)
		ei.userBreaks = nil
	}
	if ei.tlfBreaks != nil {
		return *ei.tlfBreaks
	}

	return keybase1.TLFBreak{}
}

// ctxExtendedIdentifyKeyType is a type for the context key for using
// extendedIdentify
type ctxExtendedIdentifyKeyType int

const (
	// ctxExtendedIdentifyKeyType is a context key for using extendedIdentify
	ctxExtendedIdentifyKey ctxExtendedIdentifyKeyType = iota
)

// ExtendedIdentifyAlreadyExists is returned when makeExtendedIdentify is
// called on a context already with extendedIdentify.
type ExtendedIdentifyAlreadyExists struct{}

func (e ExtendedIdentifyAlreadyExists) Error() string {
	return "extendedIdentify already exists"
}

func makeExtendedIdentify(ctx context.Context,
	behavior keybase1.TLFIdentifyBehavior) (context.Context, error) {
	if _, ok := ctx.Value(ctxExtendedIdentifyKey).(*extendedIdentify); ok {
		return nil, ExtendedIdentifyAlreadyExists{}
	}

	if !behavior.WarningInsteadOfErrorOnBrokenTracks() {
		return NewContextReplayable(ctx, func(ctx context.Context) context.Context {
			return context.WithValue(ctx, ctxExtendedIdentifyKey, &extendedIdentify{
				behavior: behavior,
			})
		}), nil
	}

	ch := make(chan keybase1.TLFIdentifyFailure)
	return NewContextReplayable(ctx, func(ctx context.Context) context.Context {
		return context.WithValue(ctx, ctxExtendedIdentifyKey, &extendedIdentify{
			behavior:   behavior,
			userBreaks: ch,
		})
	}), nil
}

func getExtendedIdentify(ctx context.Context) (ei *extendedIdentify) {
	if ei, ok := ctx.Value(ctxExtendedIdentifyKey).(*extendedIdentify); ok {
		return ei
	}
	return &extendedIdentify{
		behavior: keybase1.TLFIdentifyBehavior_DEFAULT_KBFS,
	}
}

// identifyUID performs identify based only on UID. It should be
// used only if the username is not known - as e.g. when rekeying.
func identifyUID(ctx context.Context, nug normalizedUsernameGetter,
	identifier identifier, uid keybase1.UID, isPublic bool) error {
	username, err := nug.GetNormalizedUsername(ctx, uid)
	if err != nil {
		return err
	}
	return identifyUser(ctx, nug, identifier, username, uid, isPublic)
}

// identifyUser is the preferred way to run identifies.
func identifyUser(ctx context.Context, nug normalizedUsernameGetter,
	identifier identifier, username libkb.NormalizedUsername,
	uid keybase1.UID, isPublic bool) error {
	var reason string
	if isPublic {
		reason = "You accessed a public folder."
	} else {
		reason = fmt.Sprintf("You accessed a private folder with %s.", username.String())
	}
	userInfo, err := identifier.Identify(ctx, username.String(), reason)
	if err != nil {
		// Convert libkb.NoSigChainError into one we can report.  (See
		// KBFS-1252).
		if _, ok := err.(libkb.NoSigChainError); ok {
			return NoSigChainError{username}
		}
		return err
	}
	if userInfo.Name != username {
		return fmt.Errorf("Identify returned name=%s, expected %s", userInfo.Name, username)
	}
	if userInfo.UID != uid {
		return fmt.Errorf("Identify returned uid=%s, expected %s", userInfo.UID, uid)
	}
	return nil
}

// identifyUserToChan calls identifyUser and plugs the result into the error channnel.
func identifyUserToChan(ctx context.Context, nug normalizedUsernameGetter,
	identifier identifier, name libkb.NormalizedUsername, uid keybase1.UID,
	isPublic bool, errChan chan error) {
	errChan <- identifyUser(ctx, nug, identifier, name, uid, isPublic)
}

// identifyUsers identifies the users in the given maps.
func identifyUsers(ctx context.Context, nug normalizedUsernameGetter,
	identifier identifier, users map[keybase1.UID]libkb.NormalizedUsername,
	public bool) error {
	eg, ctx := errgroup.WithContext(ctx)

	// TODO: limit the number of concurrent identifies?
	// TODO: implement a version of errgroup with limited concurrency.
	for uid, name := range users {
		uid, name := uid, name
		eg.Go(func() error {
			return identifyUser(ctx, nug, identifier, name, uid, public)
		})
	}

	return eg.Wait()
}

// identifyUserList identifies the users in the given list.
func identifyUserList(ctx context.Context, nug normalizedUsernameGetter, identifier identifier, uids []keybase1.UID, public bool) error {
	eg, ctx := errgroup.WithContext(ctx)

	// TODO: limit the number of concurrent identifies?
	// TODO: implement concurrency limited version of errgroup.
	for _, uid := range uids {
		uid := uid
		eg.Go(func() error {
			return identifyUID(ctx, nug, identifier, uid, public)
		})
	}

	return eg.Wait()
}

// identifyUsersForTLF is a helper for identifyHandle for easier testing.
func identifyUsersForTLF(ctx context.Context, nug normalizedUsernameGetter,
	identifier identifier, users map[keybase1.UID]libkb.NormalizedUsername,
	public bool) error {
	eg, ctx := errgroup.WithContext(ctx)

	eg.Go(func() error {
		ei := getExtendedIdentify(ctx)
		return ei.makeTlfBreaksIfNeeded(ctx, len(users))
	})

	eg.Go(func() error {
		return identifyUsers(ctx, nug, identifier, users, public)
	})

	return eg.Wait()
}

// identifyHandle identifies the canonical names in the given handle.
func identifyHandle(ctx context.Context, nug normalizedUsernameGetter, identifier identifier, h *TlfHandle) error {
	return identifyUsersForTLF(ctx, nug, identifier,
		h.ResolvedUsersMap(), h.IsPublic())
}
<commit_msg>libkbfs: Add more comments<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"errors"
	"fmt"
	"sync"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
)

type extendedIdentify struct {
	behavior keybase1.TLFIdentifyBehavior

	// lock guards userBreaks and tlfBreaks
	lock       sync.Mutex
	userBreaks chan keybase1.TLFIdentifyFailure
	tlfBreaks  *keybase1.TLFBreak
}

func (ei *extendedIdentify) userBreak(username libkb.NormalizedUsername, uid keybase1.UID, breaks *keybase1.IdentifyTrackBreaks) {
	if ei.userBreaks == nil {
		return
	}

	ei.userBreaks <- keybase1.TLFIdentifyFailure{
		Breaks: breaks,
		User: keybase1.User{
			Uid:      uid,
			Username: string(username),
		},
	}
}

func (ei *extendedIdentify) makeTlfBreaksIfNeeded(
	ctx context.Context, numUserInTlf int) error {
	if ei.userBreaks == nil {
		return nil
	}

	ei.lock.Lock()
	defer ei.lock.Unlock()

	b := &keybase1.TLFBreak{}
	for i := 0; i < numUserInTlf; i++ {
		select {
		case ub, ok := <-ei.userBreaks:
			if !ok {
				return errors.New("makeTlfBreaksIfNeeded called on extendedIdentify" +
					" with closed userBreaks channel.")
			}
			if ub.Breaks != nil {
				b.Breaks = append(b.Breaks, ub)
			}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	ei.tlfBreaks = b

	return nil
}

// getTlfBreakOrBust returns a keybase1.TLFBreak. This should only be called
// for behavior.WarningInsteadOfErrorOnBrokenTracks() == true, and after
// makeTlfBreaksIfNeeded is called, to make sure user proof breaks get
// populated in GUI mode.
//
// If called otherwise, we don't panic here anymore, since we can't panic on
// nil ei.tlfBreaks. The reason is if a previous successful identify has
// already happened recently, it could cause this identify to be skipped, which
// means ei.tlfBreaks is never populated. In this case, it's safe to return an
// empty keybase1.TLFBreak.
func (ei *extendedIdentify) getTlfBreakAndClose() keybase1.TLFBreak {
	ei.lock.Lock()
	defer ei.lock.Unlock()

	if ei.userBreaks != nil {
		close(ei.userBreaks)
		ei.userBreaks = nil
	}
	if ei.tlfBreaks != nil {
		return *ei.tlfBreaks
	}

	return keybase1.TLFBreak{}
}

// ctxExtendedIdentifyKeyType is a type for the context key for using
// extendedIdentify
type ctxExtendedIdentifyKeyType int

const (
	// ctxExtendedIdentifyKeyType is a context key for using extendedIdentify
	ctxExtendedIdentifyKey ctxExtendedIdentifyKeyType = iota
)

// ExtendedIdentifyAlreadyExists is returned when makeExtendedIdentify is
// called on a context already with extendedIdentify.
type ExtendedIdentifyAlreadyExists struct{}

func (e ExtendedIdentifyAlreadyExists) Error() string {
	return "extendedIdentify already exists"
}

func makeExtendedIdentify(ctx context.Context,
	behavior keybase1.TLFIdentifyBehavior) (context.Context, error) {
	if _, ok := ctx.Value(ctxExtendedIdentifyKey).(*extendedIdentify); ok {
		return nil, ExtendedIdentifyAlreadyExists{}
	}

	if !behavior.WarningInsteadOfErrorOnBrokenTracks() {
		return NewContextReplayable(ctx, func(ctx context.Context) context.Context {
			return context.WithValue(ctx, ctxExtendedIdentifyKey, &extendedIdentify{
				behavior: behavior,
			})
		}), nil
	}

	ch := make(chan keybase1.TLFIdentifyFailure)
	return NewContextReplayable(ctx, func(ctx context.Context) context.Context {
		return context.WithValue(ctx, ctxExtendedIdentifyKey, &extendedIdentify{
			behavior:   behavior,
			userBreaks: ch,
		})
	}), nil
}

func getExtendedIdentify(ctx context.Context) (ei *extendedIdentify) {
	if ei, ok := ctx.Value(ctxExtendedIdentifyKey).(*extendedIdentify); ok {
		return ei
	}
	return &extendedIdentify{
		behavior: keybase1.TLFIdentifyBehavior_DEFAULT_KBFS,
	}
}

// identifyUID performs identify based only on UID. It should be
// used only if the username is not known - as e.g. when rekeying.
func identifyUID(ctx context.Context, nug normalizedUsernameGetter,
	identifier identifier, uid keybase1.UID, isPublic bool) error {
	username, err := nug.GetNormalizedUsername(ctx, uid)
	if err != nil {
		return err
	}
	return identifyUser(ctx, nug, identifier, username, uid, isPublic)
}

// identifyUser is the preferred way to run identifies.
func identifyUser(ctx context.Context, nug normalizedUsernameGetter,
	identifier identifier, username libkb.NormalizedUsername,
	uid keybase1.UID, isPublic bool) error {
	var reason string
	if isPublic {
		reason = "You accessed a public folder."
	} else {
		reason = fmt.Sprintf("You accessed a private folder with %s.", username.String())
	}
	userInfo, err := identifier.Identify(ctx, username.String(), reason)
	if err != nil {
		// Convert libkb.NoSigChainError into one we can report.  (See
		// KBFS-1252).
		if _, ok := err.(libkb.NoSigChainError); ok {
			return NoSigChainError{username}
		}
		return err
	}
	if userInfo.Name != username {
		return fmt.Errorf("Identify returned name=%s, expected %s", userInfo.Name, username)
	}
	if userInfo.UID != uid {
		return fmt.Errorf("Identify returned uid=%s, expected %s", userInfo.UID, uid)
	}
	return nil
}

// identifyUserToChan calls identifyUser and plugs the result into the error channnel.
func identifyUserToChan(ctx context.Context, nug normalizedUsernameGetter,
	identifier identifier, name libkb.NormalizedUsername, uid keybase1.UID,
	isPublic bool, errChan chan error) {
	errChan <- identifyUser(ctx, nug, identifier, name, uid, isPublic)
}

// identifyUsers identifies the users in the given maps.
func identifyUsers(ctx context.Context, nug normalizedUsernameGetter,
	identifier identifier, users map[keybase1.UID]libkb.NormalizedUsername,
	public bool) error {
	eg, ctx := errgroup.WithContext(ctx)

	// TODO: limit the number of concurrent identifies?
	// TODO: implement a version of errgroup with limited concurrency.
	for uid, name := range users {
		// Capture range variables.
		uid, name := uid, name
		eg.Go(func() error {
			return identifyUser(ctx, nug, identifier, name, uid, public)
		})
	}

	return eg.Wait()
}

// identifyUserList identifies the users in the given list.
// Only use this when the usernames are not known - like when rekeying.
func identifyUserList(ctx context.Context, nug normalizedUsernameGetter, identifier identifier, uids []keybase1.UID, public bool) error {
	eg, ctx := errgroup.WithContext(ctx)

	// TODO: limit the number of concurrent identifies?
	// TODO: implement concurrency limited version of errgroup.
	for _, uid := range uids {
		// Capture range variable.
		uid := uid
		eg.Go(func() error {
			return identifyUID(ctx, nug, identifier, uid, public)
		})
	}

	return eg.Wait()
}

// identifyUsersForTLF is a helper for identifyHandle for easier testing.
func identifyUsersForTLF(ctx context.Context, nug normalizedUsernameGetter,
	identifier identifier, users map[keybase1.UID]libkb.NormalizedUsername,
	public bool) error {
	eg, ctx := errgroup.WithContext(ctx)

	eg.Go(func() error {
		ei := getExtendedIdentify(ctx)
		return ei.makeTlfBreaksIfNeeded(ctx, len(users))
	})

	eg.Go(func() error {
		return identifyUsers(ctx, nug, identifier, users, public)
	})

	return eg.Wait()
}

// identifyHandle identifies the canonical names in the given handle.
func identifyHandle(ctx context.Context, nug normalizedUsernameGetter, identifier identifier, h *TlfHandle) error {
	return identifyUsersForTLF(ctx, nug, identifier,
		h.ResolvedUsersMap(), h.IsPublic())
}
<|endoftext|>


<commit_before>package main

import (
	"image"
	"image/color"
	"image/draw"

	"github.com/MJKWoolnough/minecraft"
	"github.com/MJKWoolnough/minecraft/nbt"
	"github.com/MJKWoolnough/ora"
)

type terrain struct {
	Base, Top minecraft.Block
	TopLevel  uint8
}

var (
	terrainColours = color.Palette{
		color.RGBA{},
		color.RGBA{255, 255, 0, 255},   // Yellow - Sand
		color.RGBA{0, 255, 0, 255},     // Green - Grass
		color.RGBA{87, 59, 12, 255},    // Brown - Dirt
		color.RGBA{255, 128, 0, 255},   // Orange - Farm
		color.RGBA{128, 128, 128, 255}, // Grey - Stone
		color.RGBA{255, 255, 255, 255}, // White - Snow
	}
	terrainBlocks = []terrain{
		{},
		{minecraft.Block{ID: 24, Data: 2}, minecraft.Block{ID: 12}, 5}, // Sandstone - Sand
		{minecraft.Block{ID: 3}, minecraft.Block{ID: 2}, 1},            // Dirt - Grass
		{minecraft.Block{ID: 3}, minecraft.Block{ID: 3}, 0},            // Dirt - Dirt
		{minecraft.Block{ID: 3}, minecraft.Block{ID: 60, Data: 7}, 1},  // Dirt - Farmland
		{minecraft.Block{ID: 1}, minecraft.Block{ID: 1}, 0},            // Stone - Stone
		{minecraft.Block{ID: 1}, minecraft.Block{ID: 50}, 3},           // Stone - Snow
	}
)

func modeTerrain(p *image.Paletted) uint8 {
	b := p.Bounds()
	var modeMap [7]uint16
	most := uint16(0)
	mode := uint8(0)
	for i := b.Min.X; i < b.Max.X; i++ {
		for j := b.Min.Y; j < b.Max.Y; j++ {
			pos := p.ColorIndexAt(i, j)
			modeMap[pos]++
			if m := modeMap[pos]; m > most {
				most = m
				mode = pos
			}
		}
	}
	return mode
}

func meanHeight(g *image.Gray) uint8 {
	b := g.Bounds()
	var total int64
	for i := b.Min.X; i < b.Max.X; i++ {
		for j := b.Min.Y; j < b.Max.Y; j++ {
			total += int64(g.GrayAt(i, j).Y)
		}
	}
	return uint8(total / int64((b.Dx() * b.Dy())))
}

type chunkCache struct {
	mem   *minecraft.MemPath
	level *minecraft.Level
	clear nbt.Tag
	cache map[uint16]nbt.Tag
}

func newCache() *chunkCache {
	mem := minecraft.NewMemPath()
	l, _ := minecraft.NewLevel(mem)

	bedrock := minecraft.Block{ID: 7}

	l.SetBlock(0, 0, 0, minecraft.Block{})
	l.Save()
	l.Close()
	clearChunk, _ := mem.GetChunk(0, 0)

	for i := int32(-16); i < 16; i++ {
		for j := int32(0); j < 255; j++ {
			for k := int32(-16); j < 16; j++ {
				l.SetBlock(i, j, k, bedrock)
			}
		}
	}
	l.Save()
	l.Close()
	mem.SetChunk(clearChunk)
	return &chunkCache{
		mem,
		l,
		clearChunk,
		make(map[uint16]nbt.Tag),
	}
}

func (c *chunkCache) getFromCache(x, z int32, terrain uint8, height int32) nbt.Tag {
	cacheID := uint16(terrain)<<8 | uint16(height)
	chunk, ok := c.cache[cacheID]
	if !ok {
		b := terrainBlocks[terrain].Base
		for j := int32(0); j < height; j++ {
			for i := int32(0); i < 16; i++ {
				for k := int32(0); k < 16; k++ {
					c.level.SetBlock(i, j, k, b)
				}
			}
		}
		c.level.Save()
		c.level.Close()
		chunk, _ = c.mem.GetChunk(0, 0)
		c.mem.SetChunk(c.clear)
		c.cache[cacheID] = chunk
	}
	ld := chunk.Data().(nbt.Compound).Get("Level").Data().(nbt.Compound)
	ld.Set(nbt.NewTag("xPos", nbt.Int(x)))
	ld.Set(nbt.NewTag("zPos", nbt.Int(z)))
	return chunk
}

func buildMap(o *ora.ORA, c chan paint, m chan string, e chan error) {
	defer close(e)
	terrain := o.Layer("terrain")
	height := o.Layer("height")
	sTerrain := image.NewPaletted(o.Bounds(), terrainColours)
	terrainI, err := terrain.Image()
	if err != nil {
		e <- err
		return
	}
	draw.Draw(sTerrain, image.Rect(terrain.X, terrain.Y, sTerrain.Bounds().Max.X, sTerrain.Bounds().Max.Y), terrainI, image.Point{}, draw.Src)
	terrainI = nil
	sHeight := image.NewGray(o.Bounds())
	heightI, err := height.Image()
	if err != nil {
		e <- err
		return
	}
	draw.Draw(sHeight, image.Rect(height.X, height.Y, sTerrain.Bounds().Max.X, sTerrain.Bounds().Max.Y), heightI, image.Point{}, draw.Src)
	heightI = nil
	p, err := minecraft.NewFilePath("./test/")
	if err != nil {
		e <- err
		return
	}

	m <- "Building Terrain"
	buildTerrain(p, sTerrain, sHeight, c)
	m <- "Building Height Map"
	level, err := minecraft.NewLevel(p)
	if err != nil {
		e <- err
		return
	}
	level.Save()
}

func buildTerrain(mpath minecraft.Path, terrain *image.Paletted, height *image.Gray, c chan paint) {
	cc := newCache()
	b := terrain.Bounds()
	for j := 0; j < b.Max.Y; j += 16 {
		chunkZ := int32(j >> 4)
		for i := 0; i < b.Max.X; i += 16 {
			chunkX := int32(i >> 4)
			p := terrain.SubImage(image.Rect(i, j, i+16, j+16)).(*image.Paletted)
			g := height.SubImage(image.Rect(i, j, i+16, j+16)).(*image.Gray)
			t := modeTerrain(p)
			h := int32(meanHeight(g))
			mpath.SetChunk(cc.getFromCache(chunkX, chunkZ, t, h))
			c <- paint{
				terrainColours[t],
				chunkX, chunkZ,
				nil,
			}
		}
	}
}

func buildHeight(level *minecraft.Level, terrain *image.Paletted, height *image.Gray, c chan paint) {

}
<commit_msg>Set Minecraft Options<commit_after>package main

import (
	"image"
	"image/color"
	"image/draw"

	"github.com/MJKWoolnough/minecraft"
	"github.com/MJKWoolnough/minecraft/nbt"
	"github.com/MJKWoolnough/ora"
)

type terrain struct {
	Base, Top minecraft.Block
	TopLevel  uint8
}

var (
	terrainColours = color.Palette{
		color.RGBA{},
		color.RGBA{255, 255, 0, 255},   // Yellow - Sand
		color.RGBA{0, 255, 0, 255},     // Green - Grass
		color.RGBA{87, 59, 12, 255},    // Brown - Dirt
		color.RGBA{255, 128, 0, 255},   // Orange - Farm
		color.RGBA{128, 128, 128, 255}, // Grey - Stone
		color.RGBA{255, 255, 255, 255}, // White - Snow
	}
	terrainBlocks = []terrain{
		{},
		{minecraft.Block{ID: 24, Data: 2}, minecraft.Block{ID: 12}, 5}, // Sandstone - Sand
		{minecraft.Block{ID: 3}, minecraft.Block{ID: 2}, 1},            // Dirt - Grass
		{minecraft.Block{ID: 3}, minecraft.Block{ID: 3}, 0},            // Dirt - Dirt
		{minecraft.Block{ID: 3}, minecraft.Block{ID: 60, Data: 7}, 1},  // Dirt - Farmland
		{minecraft.Block{ID: 1}, minecraft.Block{ID: 1}, 0},            // Stone - Stone
		{minecraft.Block{ID: 1}, minecraft.Block{ID: 50}, 3},           // Stone - Snow
	}
)

func modeTerrain(p *image.Paletted) uint8 {
	b := p.Bounds()
	var modeMap [7]uint16
	most := uint16(0)
	mode := uint8(0)
	for i := b.Min.X; i < b.Max.X; i++ {
		for j := b.Min.Y; j < b.Max.Y; j++ {
			pos := p.ColorIndexAt(i, j)
			modeMap[pos]++
			if m := modeMap[pos]; m > most {
				most = m
				mode = pos
			}
		}
	}
	return mode
}

func meanHeight(g *image.Gray) uint8 {
	b := g.Bounds()
	var total int64
	for i := b.Min.X; i < b.Max.X; i++ {
		for j := b.Min.Y; j < b.Max.Y; j++ {
			total += int64(g.GrayAt(i, j).Y)
		}
	}
	return uint8(total / int64((b.Dx() * b.Dy())))
}

type chunkCache struct {
	mem   *minecraft.MemPath
	level *minecraft.Level
	clear nbt.Tag
	cache map[uint16]nbt.Tag
}

func newCache() *chunkCache {
	mem := minecraft.NewMemPath()
	l, _ := minecraft.NewLevel(mem)

	bedrock := minecraft.Block{ID: 7}

	l.SetBlock(0, 0, 0, minecraft.Block{})
	l.Save()
	l.Close()
	clearChunk, _ := mem.GetChunk(0, 0)

	for i := int32(-16); i < 16; i++ {
		for j := int32(0); j < 255; j++ {
			for k := int32(-16); j < 16; j++ {
				l.SetBlock(i, j, k, bedrock)
			}
		}
	}
	l.Save()
	l.Close()
	mem.SetChunk(clearChunk)
	return &chunkCache{
		mem,
		l,
		clearChunk,
		make(map[uint16]nbt.Tag),
	}
}

func (c *chunkCache) getFromCache(x, z int32, terrain uint8, height int32) nbt.Tag {
	cacheID := uint16(terrain)<<8 | uint16(height)
	chunk, ok := c.cache[cacheID]
	if !ok {
		b := terrainBlocks[terrain].Base
		for j := int32(0); j < height; j++ {
			for i := int32(0); i < 16; i++ {
				for k := int32(0); k < 16; k++ {
					c.level.SetBlock(i, j, k, b)
				}
			}
		}
		c.level.Save()
		c.level.Close()
		chunk, _ = c.mem.GetChunk(0, 0)
		c.mem.SetChunk(c.clear)
		c.cache[cacheID] = chunk
	}
	ld := chunk.Data().(nbt.Compound).Get("Level").Data().(nbt.Compound)
	ld.Set(nbt.NewTag("xPos", nbt.Int(x)))
	ld.Set(nbt.NewTag("zPos", nbt.Int(z)))
	return chunk
}

func buildMap(o *ora.ORA, c chan paint, m chan string, e chan error) {
	defer close(e)
	terrain := o.Layer("terrain")
	height := o.Layer("height")
	sTerrain := image.NewPaletted(o.Bounds(), terrainColours)
	terrainI, err := terrain.Image()
	if err != nil {
		e <- err
		return
	}
	draw.Draw(sTerrain, image.Rect(terrain.X, terrain.Y, sTerrain.Bounds().Max.X, sTerrain.Bounds().Max.Y), terrainI, image.Point{}, draw.Src)
	terrainI = nil
	sHeight := image.NewGray(o.Bounds())
	heightI, err := height.Image()
	if err != nil {
		e <- err
		return
	}
	draw.Draw(sHeight, image.Rect(height.X, height.Y, sTerrain.Bounds().Max.X, sTerrain.Bounds().Max.Y), heightI, image.Point{}, draw.Src)
	heightI = nil
	p, err := minecraft.NewFilePath("./test/")
	if err != nil {
		e <- err
		return
	}

	m <- "Building Terrain"
	buildTerrain(p, sTerrain, sHeight, c)
	m <- "Building Height Map"
	level, err := minecraft.NewLevel(p)
	if err != nil {
		e <- err
		return
	}
	level.LevelName("Test Generation")
	level.MobSpawning(false)
	level.KeepInventory(true)
	level.FireTick(false)
	level.DayLightCycle(false)
	level.MobGriefing(false)
	level.Spawn(10, 250, 10)
	level.Generator(minecraft.FlatGenerator)
	level.GeneratorOptions("0")
	level.GameMode(minecraft.Creative)
	level.AllowCommands(true)
	level.Save()
	level.Close()
}

func buildTerrain(mpath minecraft.Path, terrain *image.Paletted, height *image.Gray, c chan paint) {
	cc := newCache()
	b := terrain.Bounds()
	for j := 0; j < b.Max.Y; j += 16 {
		chunkZ := int32(j >> 4)
		for i := 0; i < b.Max.X; i += 16 {
			chunkX := int32(i >> 4)
			p := terrain.SubImage(image.Rect(i, j, i+16, j+16)).(*image.Paletted)
			g := height.SubImage(image.Rect(i, j, i+16, j+16)).(*image.Gray)
			t := modeTerrain(p)
			h := int32(meanHeight(g))
			mpath.SetChunk(cc.getFromCache(chunkX, chunkZ, t, h))
			c <- paint{
				terrainColours[t],
				chunkX, chunkZ,
				nil,
			}
		}
	}
}

func buildHeight(level *minecraft.Level, terrain *image.Paletted, height *image.Gray, c chan paint) {

}
<|endoftext|>