<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/ChristopherRabotin/smd"
)

const (
	pps1350 thrusterType = iota + 1
	pps5000
	bht1500
	bht8000
	hermes
	vx200
)

var (
	wg             sync.WaitGroup
	numCPUs        int
	opti           bool
	departEarth    bool
	interplanetary bool
	coarse         bool
	timeStep       time.Duration
	cpuChan        chan (bool)
)

func init() {
	flag.BoolVar(&departEarth, "fromEarth", true, "set to true to leave Earth")
	flag.BoolVar(&interplanetary, "interp", false, "set to true for the interplanetary missions")
	flag.BoolVar(&opti, "opti", false, "set to true to use Naasz laws")
	flag.BoolVar(&coarse, "coarse", false, "set to true to perform only a coarse simulation")
	flag.IntVar(&numCPUs, "cpus", 0, "number of CPUs to use for after first finding (set to 0 for max CPUs)")
}

type thrusterType uint8

func (tt thrusterType) Type() smd.EPThruster {
	switch tt {
	case pps1350:
		return new(smd.PPS1350)
	case pps5000:
		return new(smd.PPS5000)
	case bht1500:
		return new(smd.BHT1500)
	case bht8000:
		return new(smd.BHT8000)
	case hermes:
		return new(smd.HERMeS)
	case vx200:
		return new(smd.VX200)
	default:
		panic("unknown thruster")
	}
}

func (tt thrusterType) String() string {
	switch tt {
	case pps1350:
		return "PPS1350"
	case pps5000:
		return "PPS5000"
	case bht1500:
		return "BHT1500"
	case bht8000:
		return "BHT8000"
	case hermes:
		return "HERMeS"
	case vx200:
		return "VX200"
	default:
		panic("unknown thruster")
	}
}

func createSpacecraft(thruster thrusterType, numThrusters int, dist float64, further bool) (*smd.Spacecraft, float64) {
	/* Building spacecraft */
	thrusters := make([]smd.EPThruster, numThrusters)
	thrust := 0.0
	for i := 0; i < numThrusters; i++ {
		thrusters[i] = thruster.Type()
		voltage, power := thruster.Type().Max()
		thisThrust, _ := thruster.Type().Thrust(voltage, power)
		thrust += thisThrust
	}
	dryMass := 1.0
	fuelMass := 5e3
	name := fmt.Sprintf("%dx%s", numThrusters, thruster)
	fmt.Printf("\n===== %s ======\n", name)
	waypoints := []smd.Waypoint{smd.NewReachDistance(dist, further, nil)}
	if opti {
		if interplanetary {
			if departEarth {
				waypoints = []smd.Waypoint{smd.NewOrbitTarget(smd.Mars.HelioOrbit(time.Now()), nil, smd.Naasz, smd.OptiΔaCL, smd.OptiΔiCL), smd.NewCruiseToDistance(dist, further, nil)}
			} else {
				waypoints = []smd.Waypoint{smd.NewOrbitTarget(smd.Earth.HelioOrbit(time.Now()), nil, smd.Naasz, smd.OptiΔaCL, smd.OptiΔiCL, smd.OptiΔeCL), smd.NewCruiseToDistance(dist, further, nil)}
			}
		} else {
			if departEarth {
				// Create virtual orbit
				tgt := smd.NewOrbitFromOE(smd.Earth.SOI, 0.75, 0, 0, 230, 0, smd.Earth)
				waypoints = []smd.Waypoint{smd.NewOrbitTarget(*tgt, nil, smd.Naasz, smd.OptiΔaCL, smd.OptiΔeCL), smd.NewCruiseToDistance(dist, further, nil)}
			} else {
				tgt := smd.NewOrbitFromOE(smd.Mars.SOI, 0.85, 0, 0, 230, 0, smd.Mars)
				waypoints = []smd.Waypoint{smd.NewOrbitTarget(*tgt, nil, smd.Naasz, smd.OptiΔaCL, smd.OptiΔeCL), smd.NewCruiseToDistance(dist, further, nil)}
			}
		}
	}
	return smd.NewSpacecraft(name, dryMass, fuelMass, smd.NewUnlimitedEPS(), thrusters, false, []*smd.Cargo{}, waypoints), thrust
}

func main() {
	flag.Parse()
	availableCPUs := runtime.NumCPU()
	if numCPUs <= 0 || numCPUs > availableCPUs {
		numCPUs = availableCPUs
	}
	runtime.GOMAXPROCS(numCPUs)
	cpuChan = make(chan (bool), numCPUs)
	fmt.Printf("Running on %d CPUs\n", numCPUs)
	if coarse {
		timeStep = 5 * time.Minute
		fmt.Println("=== WARNING ===\n Using a COARSE time step -- Results may be incorrect")
	} else if interplanetary {
		timeStep = 30 * time.Second
	} else {
		timeStep = 10 * time.Second
	}
	// Go through all combinations
	combinations := []struct {
		missionNo, numThrusters int
	}{{1, 1}, {1, 2}, {2, 1}, {2, 2}, {3, 8}, {3, 12}}
	for _, combin := range combinations {
		cpuChan <- true
		go run(combin.missionNo, combin.numThrusters)
	}
	wg.Wait()
}

func run(missionNo, numThrusters int) {
	fmt.Printf("\n\n====== MISSION %d -- intp: %v -- depEarth: %v ======\n\n", missionNo, interplanetary, departEarth)
	var fn string
	if departEarth {
		if interplanetary {
			fn = "earth2mars-tof"
		} else {
			fn = "gto2escape-tof"
		}
	} else {
		if interplanetary {
			fn = "mars2earth-tof"
		} else {
			fn = "mro2escape-tof"
		}
	}
	if missionNo < 3 {
		if numThrusters == 1 {
			fn += fmt.Sprintf("-%da", missionNo)
		} else {
			fn += fmt.Sprintf("-%db", missionNo)
		}
	} else if numThrusters == 8 {
		fn += "-3a"
	} else {
		fn += "-3b"
	}
	if opti {
		fn += "-opti-we"
	} else {
		fn += "-notopti"
	}
	rslts := make(chan string, 10)
	wg.Add(1)
	go streamResults(fn, rslts)

	aGTO, eGTO := smd.Radii2ae(39300+smd.Earth.Radius, 290+smd.Earth.Radius)
	aMRO, eMRO := smd.Radii2ae(44500+smd.Mars.Radius, 426+smd.Mars.Radius)
	startDT := time.Now()
	earthOrbit := smd.Earth.HelioOrbit(startDT)
	marsOrbit := smd.Mars.HelioOrbit(startDT)

	combinations := []thrusterType{pps1350, pps5000, bht1500, bht8000, hermes, vx200}
	if missionNo == 3 {
		combinations = []thrusterType{pps5000, bht8000, hermes, vx200}
	}

	for _, thruster := range combinations {
		var bestCSV string
		var bestTOF = 1e9
		for ω := 0.0; ω < 360; ω += 10.0 {
			fmt.Printf("\n##### %.1f deg #####\n", ω)
			var initOrbit *smd.Orbit
			var distance float64
			var further bool
			if departEarth {
				if interplanetary {
					ugh := smd.Earth.HelioOrbit(startDT)
					initOrbit = &ugh
					further = true
					distance = marsOrbit.RNorm()
				} else {
					initOrbit = smd.NewOrbitFromOE(aGTO, eGTO, 0.0, 0, ω, 210, smd.Earth)
					distance = smd.Earth.SOI
					further = true
				}
			} else {
				if interplanetary {
					ugh := smd.Mars.HelioOrbit(startDT)
					initOrbit = &ugh
					further = false
					distance = earthOrbit.RNorm()
				} else {
					initOrbit = smd.NewOrbitFromOE(aMRO, eMRO, 0.0, 0, ω, 180, smd.Mars)
					distance = smd.Mars.SOI
					further = true
				}
			}
			sc, maxThrust := createSpacecraft(thruster, numThrusters, distance, further)
			if missionNo == 2 {
				sc.FuelMass = 3e3
				if departEarth {
					sc.DryMass = 900 + 577 + 1e3 + 1e3 // Add MRO, Curiosity and Schiaparelli, and suppose 1 ton bus.
				} else {
					// Suppose less return mass
					sc.DryMass = 500 + 1e3 // Add Schiaparelli return, and suppose 1 ton bus.
				}
			} else if missionNo == 3 {
				sc.DryMass = 52e3
				sc.FuelMass = 24e3
			}
			initV := initOrbit.VNorm()
			initFuel := sc.FuelMass
			// Propagate
			export := smd.ExportConfig{Filename: fn + sc.Name, AsCSV: true, Cosmo: true, Timestamp: false}
			endDT := startDT.Add(-1)
			astro := smd.NewPreciseMission(sc, initOrbit, startDT, endDT, smd.Perturbations{}, timeStep, false, export)
			astro.Propagate()
			// Compute data.
			tof := astro.CurrentDT.Sub(startDT).Hours() / 24
			deltaV := astro.Orbit.VNorm() - initV
			deltaFuel := sc.FuelMass - initFuel
			var vInf float64
			if departEarth {
				if interplanetary { // Arriving at Mars, check how fast we're going compared to some standard velocity
					vInf = astro.Orbit.VNorm() - marsOrbit.VNorm()
				} else {
					astro.Orbit.ToXCentric(smd.Sun, astro.CurrentDT)
					vInf = astro.Orbit.VNorm() - smd.Earth.HelioOrbit(astro.CurrentDT).VNorm()
				}
			} else {
				if interplanetary {
					vInf = astro.Orbit.VNorm() - earthOrbit.VNorm()
				} else {
					astro.Orbit.ToXCentric(smd.Sun, astro.CurrentDT)
					vInf = astro.Orbit.VNorm() - smd.Mars.HelioOrbit(astro.CurrentDT).VNorm()
				}
			}

			csv := fmt.Sprintf("%.3f,%s,%.3f,%.3f,%.3f,%.3f,%.3f\n", ω, sc.Name, maxThrust, tof, deltaV, deltaFuel, vInf)
			// Check if best
			if tof < bestTOF {
				bestTOF = tof
				bestCSV = csv
			}
			if true { // XXX: Switch to false for arg of periapsis finding.
				rslts <- csv
				break
			}
		}

		if !interplanetary && false { // XXX: Idem
			rslts <- bestCSV
		}
	}
	close(rslts)
	<-cpuChan
}

func streamResults(fn string, rslts <-chan string) {
	// Write CSV file.
	f, err := os.Create(fmt.Sprintf("./%s.csv", fn))
	if err != nil {
		panic(err)
	}
	defer f.Close()
	// Header
	f.WriteString("arg. periapsis (deg), name, thrust (N), T.O.F. (days), \\Delta V (km/s), fuel (kf), \\V_{\\infty} (km/s)\n")
	for {
		rslt, more := <-rslts
		if more {
			if _, err := f.WriteString(rslt); err != nil {
				panic(err)
			}
		} else {
			break
		}
	}
	f.Close()
	wg.Done()
}
<commit_msg>Add initial and final fuel in CSV output<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/ChristopherRabotin/smd"
)

const (
	pps1350 thrusterType = iota + 1
	pps5000
	bht1500
	bht8000
	hermes
	vx200
)

var (
	wg             sync.WaitGroup
	numCPUs        int
	opti           bool
	departEarth    bool
	interplanetary bool
	coarse         bool
	argPeri        bool
	timeStep       time.Duration
	cpuChan        chan (bool)
)

func init() {
	flag.BoolVar(&departEarth, "fromEarth", true, "set to true to leave Earth")
	flag.BoolVar(&interplanetary, "interp", false, "set to true for the interplanetary missions")
	flag.BoolVar(&opti, "opti", false, "set to true to use Naasz laws")
	flag.BoolVar(&coarse, "coarse", false, "set to true to perform only a coarse simulation")
	flag.IntVar(&numCPUs, "cpus", 0, "number of CPUs to use for after first finding (set to 0 for max CPUs)")
	flag.BoolVar(&argPeri, "peri", false, "set to true to search for periapsis argument (requires a lot of disk space)")
}

type thrusterType uint8

func (tt thrusterType) Type() smd.EPThruster {
	switch tt {
	case pps1350:
		return new(smd.PPS1350)
	case pps5000:
		return new(smd.PPS5000)
	case bht1500:
		return new(smd.BHT1500)
	case bht8000:
		return new(smd.BHT8000)
	case hermes:
		return new(smd.HERMeS)
	case vx200:
		return new(smd.VX200)
	default:
		panic("unknown thruster")
	}
}

func (tt thrusterType) String() string {
	switch tt {
	case pps1350:
		return "PPS1350"
	case pps5000:
		return "PPS5000"
	case bht1500:
		return "BHT1500"
	case bht8000:
		return "BHT8000"
	case hermes:
		return "HERMeS"
	case vx200:
		return "VX200"
	default:
		panic("unknown thruster")
	}
}

func createSpacecraft(thruster thrusterType, numThrusters int, dist float64, further bool) (*smd.Spacecraft, float64) {
	/* Building spacecraft */
	thrusters := make([]smd.EPThruster, numThrusters)
	thrust := 0.0
	for i := 0; i < numThrusters; i++ {
		thrusters[i] = thruster.Type()
		voltage, power := thruster.Type().Max()
		thisThrust, _ := thruster.Type().Thrust(voltage, power)
		thrust += thisThrust
	}
	dryMass := 1.0
	fuelMass := 5e3
	name := fmt.Sprintf("%dx%s", numThrusters, thruster)
	fmt.Printf("\n===== %s ======\n", name)
	waypoints := []smd.Waypoint{smd.NewReachDistance(dist, further, nil)}
	if opti {
		if interplanetary {
			if departEarth {
				waypoints = []smd.Waypoint{smd.NewOrbitTarget(smd.Mars.HelioOrbit(time.Now()), nil, smd.Naasz, smd.OptiΔaCL, smd.OptiΔiCL), smd.NewCruiseToDistance(dist, further, nil)}
			} else {
				waypoints = []smd.Waypoint{smd.NewOrbitTarget(smd.Earth.HelioOrbit(time.Now()), nil, smd.Naasz, smd.OptiΔaCL, smd.OptiΔiCL, smd.OptiΔeCL), smd.NewCruiseToDistance(dist, further, nil)}
			}
		} else {
			if departEarth {
				// Create virtual orbit
				tgt := smd.NewOrbitFromOE(smd.Earth.SOI, 0.75, 0, 0, 230, 0, smd.Earth)
				waypoints = []smd.Waypoint{smd.NewOrbitTarget(*tgt, nil, smd.Naasz, smd.OptiΔaCL, smd.OptiΔeCL), smd.NewCruiseToDistance(dist, further, nil)}
			} else {
				tgt := smd.NewOrbitFromOE(smd.Mars.SOI, 0.85, 0, 0, 230, 0, smd.Mars)
				waypoints = []smd.Waypoint{smd.NewOrbitTarget(*tgt, nil, smd.Naasz, smd.OptiΔaCL, smd.OptiΔeCL), smd.NewCruiseToDistance(dist, further, nil)}
			}
		}
	}
	return smd.NewSpacecraft(name, dryMass, fuelMass, smd.NewUnlimitedEPS(), thrusters, false, []*smd.Cargo{}, waypoints), thrust
}

func main() {
	flag.Parse()
	availableCPUs := runtime.NumCPU()
	if numCPUs <= 0 || numCPUs > availableCPUs {
		numCPUs = availableCPUs
	}
	runtime.GOMAXPROCS(numCPUs)
	cpuChan = make(chan (bool), numCPUs)
	fmt.Printf("Running on %d CPUs\n", numCPUs)
	if coarse {
		timeStep = 5 * time.Minute
		fmt.Println("=== WARNING ===\n Using a COARSE time step -- Results may be incorrect")
	} else if interplanetary {
		timeStep = 30 * time.Second
	} else {
		timeStep = 10 * time.Second
	}
	// Go through all combinations
	combinations := []struct {
		missionNo, numThrusters int
	}{{1, 1}, {1, 2}, {2, 1}, {2, 2}, {3, 8}, {3, 12}}
	for _, combin := range combinations {
		cpuChan <- true
		go run(combin.missionNo, combin.numThrusters)
	}
	wg.Wait()
}

func run(missionNo, numThrusters int) {
	fmt.Printf("\n\n====== MISSION %d -- intp: %v -- depEarth: %v ======\n\n", missionNo, interplanetary, departEarth)
	var fn string
	if departEarth {
		if interplanetary {
			fn = "earth2mars-tof"
		} else {
			fn = "gto2escape-tof"
		}
	} else {
		if interplanetary {
			fn = "mars2earth-tof"
		} else {
			fn = "mro2escape-tof"
		}
	}
	if missionNo < 3 {
		if numThrusters == 1 {
			fn += fmt.Sprintf("-%da", missionNo)
		} else {
			fn += fmt.Sprintf("-%db", missionNo)
		}
	} else if numThrusters == 8 {
		fn += "-3a"
	} else {
		fn += "-3b"
	}
	if opti {
		fn += "-opti-we"
	} else {
		fn += "-notopti"
	}
	rslts := make(chan string, 10)
	wg.Add(1)
	go streamResults(fn, rslts)

	aGTO, eGTO := smd.Radii2ae(39300+smd.Earth.Radius, 290+smd.Earth.Radius)
	aMRO, eMRO := smd.Radii2ae(44500+smd.Mars.Radius, 426+smd.Mars.Radius)
	startDT := time.Now()
	earthOrbit := smd.Earth.HelioOrbit(startDT)
	marsOrbit := smd.Mars.HelioOrbit(startDT)

	combinations := []thrusterType{pps1350, pps5000, bht1500, bht8000, hermes, vx200}
	if missionNo == 3 {
		combinations = []thrusterType{pps5000, bht8000, hermes, vx200}
	}

	for _, thruster := range combinations {
		var bestCSV string
		var bestTOF = 1e9
		for ω := 0.0; ω < 360; ω += 10.0 {
			fmt.Printf("\n##### %.1f deg #####\n", ω)
			var initOrbit *smd.Orbit
			var distance float64
			var further bool
			if departEarth {
				if interplanetary {
					ugh := smd.Earth.HelioOrbit(startDT)
					initOrbit = &ugh
					further = true
					distance = marsOrbit.RNorm()
				} else {
					initOrbit = smd.NewOrbitFromOE(aGTO, eGTO, 0.0, 0, ω, 210, smd.Earth)
					distance = smd.Earth.SOI
					further = true
				}
			} else {
				if interplanetary {
					ugh := smd.Mars.HelioOrbit(startDT)
					initOrbit = &ugh
					further = false
					distance = earthOrbit.RNorm()
				} else {
					initOrbit = smd.NewOrbitFromOE(aMRO, eMRO, 0.0, 0, ω, 180, smd.Mars)
					distance = smd.Mars.SOI
					further = true
				}
			}
			sc, maxThrust := createSpacecraft(thruster, numThrusters, distance, further)
			if missionNo == 2 {
				sc.FuelMass = 3e3
				if departEarth {
					sc.DryMass = 900 + 577 + 1e3 + 1e3 // Add MRO, Curiosity and Schiaparelli, and suppose 1 ton bus.
				} else {
					// Suppose less return mass
					sc.DryMass = 500 + 1e3 // Add Schiaparelli return, and suppose 1 ton bus.
				}
			} else if missionNo == 3 {
				sc.DryMass = 52e3
				sc.FuelMass = 24e3
			}
			initV := initOrbit.VNorm()
			initFuel := sc.FuelMass
			// Propagate
			export := smd.ExportConfig{Filename: fn + sc.Name, AsCSV: argPeri, Cosmo: false, Timestamp: false}
			endDT := startDT.Add(-1)
			astro := smd.NewPreciseMission(sc, initOrbit, startDT, endDT, smd.Perturbations{}, timeStep, false, export)
			astro.Propagate()
			// Compute data.
			tof := astro.CurrentDT.Sub(startDT).Hours() / 24
			deltaV := astro.Orbit.VNorm() - initV
			var vInf float64
			if departEarth {
				if interplanetary { // Arriving at Mars, check how fast we're going compared to some standard velocity
					vInf = astro.Orbit.VNorm() - marsOrbit.VNorm()
				} else {
					astro.Orbit.ToXCentric(smd.Sun, astro.CurrentDT)
					vInf = astro.Orbit.VNorm() - smd.Earth.HelioOrbit(astro.CurrentDT).VNorm()
				}
			} else {
				if interplanetary {
					vInf = astro.Orbit.VNorm() - earthOrbit.VNorm()
				} else {
					astro.Orbit.ToXCentric(smd.Sun, astro.CurrentDT)
					vInf = astro.Orbit.VNorm() - smd.Mars.HelioOrbit(astro.CurrentDT).VNorm()
				}
			}

			csv := fmt.Sprintf("%.3f,%s,%.3f,%.3f,%.3f,%.3f,%.3f,%3f\n", ω, sc.Name, maxThrust, tof, deltaV, initFuel, sc.FuelMass, vInf)
			// Check if best
			if tof < bestTOF {
				bestTOF = tof
				bestCSV = csv
			}
			if !argPeri {
				rslts <- csv
				break
			}
		}

		if !interplanetary && argPeri {
			rslts <- bestCSV
		}
	}
	close(rslts)
	<-cpuChan
}

func streamResults(fn string, rslts <-chan string) {
	// Write CSV file.
	f, err := os.Create(fmt.Sprintf("./%s.csv", fn))
	if err != nil {
		panic(err)
	}
	defer f.Close()
	// Header
	f.WriteString("arg. periapsis (deg), name, thrust (N), T.O.F. (days), \\Delta V (km/s), init fuel (kg), final fuel (kg), \\V_{\\infty} (km/s)\n")
	for {
		rslt, more := <-rslts
		if more {
			if _, err := f.WriteString(rslt); err != nil {
				panic(err)
			}
		} else {
			break
		}
	}
	f.Close()
	wg.Done()
}
<|endoftext|>


<commit_before>package postgres

import (
	"bytes"
	"fmt"
	"github.com/jasonish/evebox/core"
	"github.com/jasonish/evebox/elasticsearch"
	"github.com/jasonish/evebox/log"
	"github.com/satori/go.uuid"
	"gopkg.in/square/go-jose.v1/json"
)

type EventService struct {
	core.NotImplementedEventService

	pg *Service
}

func NewEventService(pg *Service) *EventService {
	service := EventService{}
	service.pg = pg
	return &service
}

func (s *EventService) FindNetflow(options core.EventQueryOptions, sortBy string, order string) (interface{}, error) {

	size := int64(10)

	sql := `SELECT
                  uuid,
                  source
                FROM events_master
                WHERE source ->> 'event_type' = 'netflow'
`

	if order == "" {
		order = "desc"
	}

	if options.Size > 0 {
		size = options.Size
	}

	if options.TimeRange != "" {
		sql += fmt.Sprintf(" AND timestamp >= NOW() - '%s'::interval", options.TimeRange)
	}

	if sortBy == "netflow.pkts" {
		sql += fmt.Sprintf(" ORDER BY source -> 'netflow' ->> 'pkts' %s", order)
	} else if sortBy == "netflow.bytes" {
		sql += fmt.Sprintf(" ORDER BY source -> 'netflow' ->> 'bytes' %s", order)
	} else {
		return nil, fmt.Errorf("Unknown sort field: %s", sortBy)
	}

	sql += fmt.Sprintf(" LIMIT %d", size)

	log.Println(sql)

	rows, err := s.pg.db.Query(sql)
	if err != nil {
		return nil, err
	}

	data := []map[string]interface{}{}

	for rows.Next() {
		var eventId uuid.UUID
		var rawSource []byte
		var source map[string]interface{}

		err = rows.Scan(&eventId, &rawSource)
		if err != nil {
			return nil, err
		}

		decoder := json.NewDecoder(bytes.NewReader(rawSource))
		decoder.UseNumber()
		err = decoder.Decode(&source)
		if err != nil {
			return nil, err
		}

		log.Println(elasticsearch.ToJson(source))

		data = append(data, map[string]interface{}{
			"_id":     eventId.String(),
			"_source": source,
		})
	}

	return map[string]interface{}{
		"data": data,
	}, nil
}
<commit_msg>fix json import<commit_after>package postgres

import (
	"bytes"
	"fmt"
	"github.com/jasonish/evebox/core"
	"github.com/jasonish/evebox/elasticsearch"
	"github.com/jasonish/evebox/log"
	"github.com/satori/go.uuid"
	"encoding/json"
)

type EventService struct {
	core.NotImplementedEventService

	pg *Service
}

func NewEventService(pg *Service) *EventService {
	service := EventService{}
	service.pg = pg
	return &service
}

func (s *EventService) FindNetflow(options core.EventQueryOptions, sortBy string, order string) (interface{}, error) {

	size := int64(10)

	sql := `SELECT
                  uuid,
                  source
                FROM events_master
                WHERE source ->> 'event_type' = 'netflow'
`

	if order == "" {
		order = "desc"
	}

	if options.Size > 0 {
		size = options.Size
	}

	if options.TimeRange != "" {
		sql += fmt.Sprintf(" AND timestamp >= NOW() - '%s'::interval", options.TimeRange)
	}

	if sortBy == "netflow.pkts" {
		sql += fmt.Sprintf(" ORDER BY source -> 'netflow' ->> 'pkts' %s", order)
	} else if sortBy == "netflow.bytes" {
		sql += fmt.Sprintf(" ORDER BY source -> 'netflow' ->> 'bytes' %s", order)
	} else {
		return nil, fmt.Errorf("Unknown sort field: %s", sortBy)
	}

	sql += fmt.Sprintf(" LIMIT %d", size)

	log.Println(sql)

	rows, err := s.pg.db.Query(sql)
	if err != nil {
		return nil, err
	}

	data := []map[string]interface{}{}

	for rows.Next() {
		var eventId uuid.UUID
		var rawSource []byte
		var source map[string]interface{}

		err = rows.Scan(&eventId, &rawSource)
		if err != nil {
			return nil, err
		}

		decoder := json.NewDecoder(bytes.NewReader(rawSource))
		decoder.UseNumber()
		err = decoder.Decode(&source)
		if err != nil {
			return nil, err
		}

		log.Println(elasticsearch.ToJson(source))

		data = append(data, map[string]interface{}{
			"_id":     eventId.String(),
			"_source": source,
		})
	}

	return map[string]interface{}{
		"data": data,
	}, nil
}
<|endoftext|>


<commit_before>package convert

import (
	"crypto/md5"
	"fmt"
	"io"
	"strconv"
)

func Int(val interface{}) int64 {
	if val != nil {
		switch val.(type) {
		case float64:
			return int64(val.(float64))
		case int64:
			return val.(int64)
		case string:
			ret, err := strconv.Atoi(val.(string))
			if err != nil {
				return 0
			} else {
				return int64(ret)
			}
		}
	}
	return 0
}

func IntWith(val interface{}, defaultValue int64) int64 {
	if val != nil {
		switch val.(type) {
		case float64:
			return int64(val.(float64))
		case int64:
			return val.(int64)
		case string:
			ret, err := strconv.Atoi(val.(string))
			if err == nil {
				return int64(ret)
			}
		}
	}
	return defaultValue
}

func String(val interface{}) string {
	switch val.(type) {
	case string:
		return val.(string)
	case []byte:
		return string(val.([]byte))
	case nil:
		return ""
	default:
		return fmt.Sprintf("%v", val)
	}
}

func MD5(src string) string {
	h := md5.New()
	io.WriteString(h, src)
	return h.Sum(nil)
}
<commit_msg>convert MD5 error fix<commit_after>package convert

import (
	"crypto/md5"
	"fmt"
	"io"
	"strconv"
)

func Int(val interface{}) int64 {
	if val != nil {
		switch val.(type) {
		case float64:
			return int64(val.(float64))
		case int64:
			return val.(int64)
		case string:
			ret, err := strconv.Atoi(val.(string))
			if err != nil {
				return 0
			} else {
				return int64(ret)
			}
		}
	}
	return 0
}

func IntWith(val interface{}, defaultValue int64) int64 {
	if val != nil {
		switch val.(type) {
		case float64:
			return int64(val.(float64))
		case int64:
			return val.(int64)
		case string:
			ret, err := strconv.Atoi(val.(string))
			if err == nil {
				return int64(ret)
			}
		}
	}
	return defaultValue
}

func String(val interface{}) string {
	switch val.(type) {
	case string:
		return val.(string)
	case []byte:
		return string(val.([]byte))
	case nil:
		return ""
	default:
		return fmt.Sprintf("%v", val)
	}
}

func MD5(src string) string {
	h := md5.New()
	io.WriteString(h, src)
	return string(h.Sum(nil))
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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

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

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

package main

import (
	"flag"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"strings"
	"time"

	"github.com/golang/glog"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	v3 "google.golang.org/api/monitoring/v3"

	"github.com/GoogleCloudPlatform/k8s-stackdriver/prometheus-to-sd/config"
	"github.com/GoogleCloudPlatform/k8s-stackdriver/prometheus-to-sd/flags"
	"github.com/GoogleCloudPlatform/k8s-stackdriver/prometheus-to-sd/translator"
)

var (
	host       = flag.String("target-host", "localhost", "The monitored component's hostname. DEPRECATED: Use --source instead.")
	port       = flag.Uint("target-port", 80, "The monitored component's port. DEPRECATED: Use --source instead.")
	component  = flag.String("component", "", "The monitored target's name. DEPRECATED: Use --source instead.")
	resolution = flag.Duration("metrics-resolution", 60*time.Second,
		"The resolution at which prometheus-to-sd will scrape the component for metrics.")
	metricsPrefix = flag.String("stackdriver-prefix", "container.googleapis.com/master",
		"Prefix that is appended to every metric.")
	whitelisted = flag.String("whitelisted-metrics", "",
		"Comma-separated list of whitelisted metrics. If empty all metrics will be exported. DEPRECATED: Use --source instead.")
	autoWhitelistMetrics = flag.Bool("auto-whitelist-metrics", false,
		"If component has no whitelisted metrics, prometheus-to-sd will fetch them from Stackdriver.")
	metricDescriptorsResolution = flag.Duration("metric-descriptors-resolution", 10*time.Minute,
		"The resolution at which prometheus-to-sd will scrape metric descriptors from Stackdriver.")
	apioverride = flag.String("api-override", "",
		"The stackdriver API endpoint to override the default one used (which is prod).")
	source = flags.Uris{}
	podId  = flag.String("pod-id", "machine",
		"Name of the pod in which monitored component is running.")
	namespaceId = flag.String("namespace-id", "",
		"Namespace name of the pod in which monitored component is running.")
	omitComponentName = flag.Bool("omit-component-name", true,
		"If metric name starts with the component name then this substring is removed to keep metric name shorter.")
	debugPort = flag.Uint("port", 6060, "Port on which debug information is exposed.")

	customMetricsPrefix = "custom.googleapis.com"
)

func main() {
	flag.Set("logtostderr", "true")
	flag.Var(&source, "source", "source(s) to watch in [component-name]:http://host:port?whitelisted=a,b,c format")

	defer glog.Flush()
	flag.Parse()

	sourceConfigs := config.SourceConfigsFromFlags(source, component, host, port, whitelisted)

	gceConf, err := config.GetGceConfig(*metricsPrefix)
	podConfig := &config.PodConfig{
		PodId:       *podId,
		NamespaceId: *namespaceId,
	}
	if err != nil {
		glog.Fatalf("Failed to get GCE config: %v", err)
	}
	glog.Infof("GCE config: %+v", gceConf)

	go func() {
		glog.Error(http.ListenAndServe(fmt.Sprintf(":%d", *debugPort), nil))
	}()

	client := oauth2.NewClient(context.Background(), google.ComputeTokenSource(""))
	stackdriverService, err := v3.New(client)
	if *apioverride != "" {
		stackdriverService.BasePath = *apioverride
	}
	if err != nil {
		glog.Fatalf("Failed to create Stackdriver client: %v", err)
	}
	glog.V(4).Infof("Successfully created Stackdriver client")

	if len(sourceConfigs) == 0 {
		glog.Fatalf("No sources defined. Please specify at least one --source flag.")
	}

	for _, sourceConfig := range sourceConfigs {
		glog.V(4).Infof("Starting goroutine for %+v", sourceConfig)

		// Pass sourceConfig as a parameter to avoid using the last sourceConfig by all goroutines.
		go readAndPushDataToStackdriver(stackdriverService, gceConf, podConfig, sourceConfig)
	}

	// As worker goroutines work forever, block main thread as well.
	<-make(chan int)
}

func readAndPushDataToStackdriver(stackdriverService *v3.Service, gceConf *config.GceConfig, podConfig *config.PodConfig, sourceConfig config.SourceConfig) {
	glog.Infof("Running prometheus-to-sd, monitored target is %s %v:%v", sourceConfig.Component, sourceConfig.Host, sourceConfig.Port)
	commonConfig := &config.CommonConfig{
		GceConfig:     gceConf,
		PodConfig:     podConfig,
		ComponentName: sourceConfig.Component,
	}
	metricDescriptorCache := translator.NewMetricDescriptorCache(stackdriverService, commonConfig, sourceConfig.Component)
	signal := time.After(0)
	useWhitelistedMetricsAutodiscovery := *autoWhitelistMetrics && len(sourceConfig.Whitelisted) == 0

	for range time.Tick(*resolution) {
		// Mark cache at the beginning of each iteration as stale. Cache is considered refreshed only if during
		// current iteration there was successful call to Refresh function.
		metricDescriptorCache.MarkStale()
		glog.V(4).Infof("Scraping metrics of component %v", sourceConfig.Component)
		select {
		case <-signal:
			glog.V(4).Infof("Updating metrics cache for component %v", sourceConfig.Component)
			metricDescriptorCache.Refresh()
			if useWhitelistedMetricsAutodiscovery {
				sourceConfig.UpdateWhitelistedMetrics(metricDescriptorCache.GetMetricNames())
				glog.V(2).Infof("Autodiscovered whitelisted metrics for component %v: %v", commonConfig.ComponentName, sourceConfig.Whitelisted)
			}
			signal = time.After(*metricDescriptorsResolution)
		default:
		}
		if useWhitelistedMetricsAutodiscovery && len(sourceConfig.Whitelisted) == 0 {
			glog.V(4).Infof("Skipping %v component as there are no metric to expose.", sourceConfig.Component)
			continue
		}
		metrics, err := translator.GetPrometheusMetrics(sourceConfig.Host, sourceConfig.Port)
		if err != nil {
			glog.V(2).Infof("Error while getting Prometheus metrics %v for component %v", err, sourceConfig.Component)
			continue
		}
		if *omitComponentName {
			metrics = translator.OmitComponentName(metrics, sourceConfig.Component)
		}
		if strings.HasPrefix(commonConfig.GceConfig.MetricsPrefix, customMetricsPrefix) {
			metricDescriptorCache.UpdateMetricDescriptors(metrics, sourceConfig.Whitelisted)
		}
		ts := translator.TranslatePrometheusToStackdriver(commonConfig, sourceConfig.Whitelisted, metrics, metricDescriptorCache)
		translator.SendToStackdriver(stackdriverService, commonConfig, ts)
	}
}
<commit_msg>Update default debug port used by prometheus-to-sd.<commit_after>/*
Copyright 2017 Google Inc.

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

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

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

package main

import (
	"flag"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"strings"
	"time"

	"github.com/golang/glog"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	v3 "google.golang.org/api/monitoring/v3"

	"github.com/GoogleCloudPlatform/k8s-stackdriver/prometheus-to-sd/config"
	"github.com/GoogleCloudPlatform/k8s-stackdriver/prometheus-to-sd/flags"
	"github.com/GoogleCloudPlatform/k8s-stackdriver/prometheus-to-sd/translator"
)

var (
	host       = flag.String("target-host", "localhost", "The monitored component's hostname. DEPRECATED: Use --source instead.")
	port       = flag.Uint("target-port", 80, "The monitored component's port. DEPRECATED: Use --source instead.")
	component  = flag.String("component", "", "The monitored target's name. DEPRECATED: Use --source instead.")
	resolution = flag.Duration("metrics-resolution", 60*time.Second,
		"The resolution at which prometheus-to-sd will scrape the component for metrics.")
	metricsPrefix = flag.String("stackdriver-prefix", "container.googleapis.com/master",
		"Prefix that is appended to every metric.")
	whitelisted = flag.String("whitelisted-metrics", "",
		"Comma-separated list of whitelisted metrics. If empty all metrics will be exported. DEPRECATED: Use --source instead.")
	autoWhitelistMetrics = flag.Bool("auto-whitelist-metrics", false,
		"If component has no whitelisted metrics, prometheus-to-sd will fetch them from Stackdriver.")
	metricDescriptorsResolution = flag.Duration("metric-descriptors-resolution", 10*time.Minute,
		"The resolution at which prometheus-to-sd will scrape metric descriptors from Stackdriver.")
	apioverride = flag.String("api-override", "",
		"The stackdriver API endpoint to override the default one used (which is prod).")
	source = flags.Uris{}
	podId  = flag.String("pod-id", "machine",
		"Name of the pod in which monitored component is running.")
	namespaceId = flag.String("namespace-id", "",
		"Namespace name of the pod in which monitored component is running.")
	omitComponentName = flag.Bool("omit-component-name", true,
		"If metric name starts with the component name then this substring is removed to keep metric name shorter.")
	debugPort = flag.Uint("port", 6061, "Port on which debug information is exposed.")

	customMetricsPrefix = "custom.googleapis.com"
)

func main() {
	flag.Set("logtostderr", "true")
	flag.Var(&source, "source", "source(s) to watch in [component-name]:http://host:port?whitelisted=a,b,c format")

	defer glog.Flush()
	flag.Parse()

	sourceConfigs := config.SourceConfigsFromFlags(source, component, host, port, whitelisted)

	gceConf, err := config.GetGceConfig(*metricsPrefix)
	podConfig := &config.PodConfig{
		PodId:       *podId,
		NamespaceId: *namespaceId,
	}
	if err != nil {
		glog.Fatalf("Failed to get GCE config: %v", err)
	}
	glog.Infof("GCE config: %+v", gceConf)

	go func() {
		glog.Error(http.ListenAndServe(fmt.Sprintf(":%d", *debugPort), nil))
	}()

	client := oauth2.NewClient(context.Background(), google.ComputeTokenSource(""))
	stackdriverService, err := v3.New(client)
	if *apioverride != "" {
		stackdriverService.BasePath = *apioverride
	}
	if err != nil {
		glog.Fatalf("Failed to create Stackdriver client: %v", err)
	}
	glog.V(4).Infof("Successfully created Stackdriver client")

	if len(sourceConfigs) == 0 {
		glog.Fatalf("No sources defined. Please specify at least one --source flag.")
	}

	for _, sourceConfig := range sourceConfigs {
		glog.V(4).Infof("Starting goroutine for %+v", sourceConfig)

		// Pass sourceConfig as a parameter to avoid using the last sourceConfig by all goroutines.
		go readAndPushDataToStackdriver(stackdriverService, gceConf, podConfig, sourceConfig)
	}

	// As worker goroutines work forever, block main thread as well.
	<-make(chan int)
}

func readAndPushDataToStackdriver(stackdriverService *v3.Service, gceConf *config.GceConfig, podConfig *config.PodConfig, sourceConfig config.SourceConfig) {
	glog.Infof("Running prometheus-to-sd, monitored target is %s %v:%v", sourceConfig.Component, sourceConfig.Host, sourceConfig.Port)
	commonConfig := &config.CommonConfig{
		GceConfig:     gceConf,
		PodConfig:     podConfig,
		ComponentName: sourceConfig.Component,
	}
	metricDescriptorCache := translator.NewMetricDescriptorCache(stackdriverService, commonConfig, sourceConfig.Component)
	signal := time.After(0)
	useWhitelistedMetricsAutodiscovery := *autoWhitelistMetrics && len(sourceConfig.Whitelisted) == 0

	for range time.Tick(*resolution) {
		// Mark cache at the beginning of each iteration as stale. Cache is considered refreshed only if during
		// current iteration there was successful call to Refresh function.
		metricDescriptorCache.MarkStale()
		glog.V(4).Infof("Scraping metrics of component %v", sourceConfig.Component)
		select {
		case <-signal:
			glog.V(4).Infof("Updating metrics cache for component %v", sourceConfig.Component)
			metricDescriptorCache.Refresh()
			if useWhitelistedMetricsAutodiscovery {
				sourceConfig.UpdateWhitelistedMetrics(metricDescriptorCache.GetMetricNames())
				glog.V(2).Infof("Autodiscovered whitelisted metrics for component %v: %v", commonConfig.ComponentName, sourceConfig.Whitelisted)
			}
			signal = time.After(*metricDescriptorsResolution)
		default:
		}
		if useWhitelistedMetricsAutodiscovery && len(sourceConfig.Whitelisted) == 0 {
			glog.V(4).Infof("Skipping %v component as there are no metric to expose.", sourceConfig.Component)
			continue
		}
		metrics, err := translator.GetPrometheusMetrics(sourceConfig.Host, sourceConfig.Port)
		if err != nil {
			glog.V(2).Infof("Error while getting Prometheus metrics %v for component %v", err, sourceConfig.Component)
			continue
		}
		if *omitComponentName {
			metrics = translator.OmitComponentName(metrics, sourceConfig.Component)
		}
		if strings.HasPrefix(commonConfig.GceConfig.MetricsPrefix, customMetricsPrefix) {
			metricDescriptorCache.UpdateMetricDescriptors(metrics, sourceConfig.Whitelisted)
		}
		ts := translator.TranslatePrometheusToStackdriver(commonConfig, sourceConfig.Whitelisted, metrics, metricDescriptorCache)
		translator.SendToStackdriver(stackdriverService, commonConfig, ts)
	}
}
<|endoftext|>


<commit_before>// +build go1.9

package prometheus

import (
	"fmt"
	"log"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/armon/go-metrics"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/push"
)

var (
	// DefaultPrometheusOpts is the default set of options used when creating a
	// PrometheusSink.
	DefaultPrometheusOpts = PrometheusOpts{
		Expiration: 60 * time.Second,
	}
)

// PrometheusOpts is used to configure the Prometheus Sink
type PrometheusOpts struct {
	// Expiration is the duration a metric is valid for, after which it will be
	// untracked. If the value is zero, a metric is never expired.
	Expiration time.Duration
	Registerer prometheus.Registerer
}

type PrometheusSink struct {
	// If these will ever be copied, they should be converted to *sync.Map values and initialized appropriately
	gauges     sync.Map
	summaries  sync.Map
	counters   sync.Map
	expiration time.Duration
}

type PrometheusGauge struct {
	prometheus.Gauge
	updatedAt time.Time
}

type PrometheusSummary struct {
	prometheus.Summary
	updatedAt time.Time
}

type PrometheusCounter struct {
	prometheus.Counter
	updatedAt time.Time
}

// NewPrometheusSink creates a new PrometheusSink using the default options.
func NewPrometheusSink() (*PrometheusSink, error) {
	return NewPrometheusSinkFrom(DefaultPrometheusOpts)
}

// NewPrometheusSinkFrom creates a new PrometheusSink using the passed options.
func NewPrometheusSinkFrom(opts PrometheusOpts) (*PrometheusSink, error) {
	sink := &PrometheusSink{
		gauges:     sync.Map{},
		summaries:  sync.Map{},
		counters:   sync.Map{},
		expiration: opts.Expiration,
	}

	reg := opts.Registerer
	if reg == nil {
		reg = prometheus.DefaultRegisterer
	}

	return sink, reg.Register(sink)
}

// Describe is needed to meet the Collector interface.
func (p *PrometheusSink) Describe(c chan<- *prometheus.Desc) {
	// We must emit some description otherwise an error is returned. This
	// description isn't shown to the user!
	prometheus.NewGauge(prometheus.GaugeOpts{Name: "Dummy", Help: "Dummy"}).Describe(c)
}

// Collect meets the collection interface and allows us to enforce our expiration
// logic to clean up ephemeral metrics if their value haven't been set for a
// duration exceeding our allowed expiration time.
func (p *PrometheusSink) Collect(c chan<- prometheus.Metric) {
	expire := p.expiration != 0
	now := time.Now()
	p.gauges.Range(func(k, v interface{}) bool {
		if v != nil {
			lastUpdate := v.(*PrometheusGauge).updatedAt
			if expire && lastUpdate.Add(p.expiration).Before(now) {
				p.gauges.Delete(k)
			} else {
				v.(*PrometheusGauge).Collect(c)
			}
		}
		return true
	})
	p.summaries.Range(func(k, v interface{}) bool {
		if v != nil {
			lastUpdate := v.(*PrometheusSummary).updatedAt
			if expire && lastUpdate.Add(p.expiration).Before(now) {
				p.summaries.Delete(k)
			} else {
				v.(*PrometheusSummary).Collect(c)
			}
		}
		return true
	})
	p.counters.Range(func(k, v interface{}) bool {
		if v != nil {
			lastUpdate := v.(*PrometheusCounter).updatedAt
			if expire && lastUpdate.Add(p.expiration).Before(now) {
				p.counters.Delete(k)
			} else {
				v.(*PrometheusCounter).Collect(c)
			}
		}
		return true
	})
}

var forbiddenChars = regexp.MustCompile("[ .=\\-/]")

func (p *PrometheusSink) flattenKey(parts []string, labels []metrics.Label) (string, string) {
	key := strings.Join(parts, "_")
	key = forbiddenChars.ReplaceAllString(key, "_")

	hash := key
	for _, label := range labels {
		hash += fmt.Sprintf(";%s=%s", label.Name, label.Value)
	}

	return key, hash
}

func prometheusLabels(labels []metrics.Label) prometheus.Labels {
	l := make(prometheus.Labels)
	for _, label := range labels {
		l[label.Name] = label.Value
	}
	return l
}

func (p *PrometheusSink) SetGauge(parts []string, val float32) {
	p.SetGaugeWithLabels(parts, val, nil)
}

func (p *PrometheusSink) SetGaugeWithLabels(parts []string, val float32, labels []metrics.Label) {
	key, hash := p.flattenKey(parts, labels)
	pg, ok := p.gauges.Load(hash)

	// The sync.Map underlying gauges stores pointers to our structs. If we need to make updates,
	// rather than modifying the underlying value directly, which would be racy, we make a local
	// copy by dereferencing the pointer we get back, making the appropriate changes, and then
	// storing a pointer to our local copy. The underlying Prometheus types are threadsafe,
	// so there's no issues there. It's possible for racy updates to occur to the updatedAt
	// value, but since we're always setting it to time.Now(), it doesn't really matter.
	if ok {
		localGauge := *pg.(*PrometheusGauge)
		localGauge.Set(float64(val))
		localGauge.updatedAt = time.Now()
		p.gauges.Store(hash, &localGauge)
	} else {
		g := prometheus.NewGauge(prometheus.GaugeOpts{
			Name:        key,
			Help:        key,
			ConstLabels: prometheusLabels(labels),
		})
		g.Set(float64(val))
		pg = &PrometheusGauge{
			g, time.Now(),
		}
		p.gauges.Store(hash, pg)
	}
}

func (p *PrometheusSink) AddSample(parts []string, val float32) {
	p.AddSampleWithLabels(parts, val, nil)
}

func (p *PrometheusSink) AddSampleWithLabels(parts []string, val float32, labels []metrics.Label) {
	key, hash := p.flattenKey(parts, labels)
	ps, ok := p.summaries.Load(hash)

	if ok {
		localSummary := *ps.(*PrometheusSummary)
		localSummary.Observe(float64(val))
		localSummary.updatedAt = time.Now()
		p.summaries.Store(hash, &localSummary)
	} else {
		s := prometheus.NewSummary(prometheus.SummaryOpts{
			Name:        key,
			Help:        key,
			MaxAge:      10 * time.Second,
			ConstLabels: prometheusLabels(labels),
			Objectives:  map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001},
		})
		s.Observe(float64(val))
		ps = &PrometheusSummary{
			s, time.Now(),
		}
		p.summaries.Store(hash, ps)
	}
}

// EmitKey is not implemented. Prometheus doesn’t offer a type for which an
// arbitrary number of values is retained, as Prometheus works with a pull
// model, rather than a push model.
func (p *PrometheusSink) EmitKey(key []string, val float32) {
}

func (p *PrometheusSink) IncrCounter(parts []string, val float32) {
	p.IncrCounterWithLabels(parts, val, nil)
}

func (p *PrometheusSink) IncrCounterWithLabels(parts []string, val float32, labels []metrics.Label) {
	key, hash := p.flattenKey(parts, labels)
	pc, ok := p.counters.Load(hash)

	if ok {
		localCounter := *pc.(*PrometheusCounter)
		localCounter.Add(float64(val))
		localCounter.updatedAt = time.Now()
		p.counters.Store(hash, &localCounter)
	} else {
		c := prometheus.NewCounter(prometheus.CounterOpts{
			Name:        key,
			Help:        key,
			ConstLabels: prometheusLabels(labels),
		})
		c.Add(float64(val))
		pc = &PrometheusCounter{
			c, time.Now(),
		}
		p.counters.Store(hash, pc)
	}
}

type PrometheusPushSink struct {
	*PrometheusSink
	pusher       *push.Pusher
	address      string
	pushInterval time.Duration
	stopChan     chan struct{}
}

func NewPrometheusPushSink(address string, pushIterval time.Duration, name string) (*PrometheusPushSink, error) {
	promSink := &PrometheusSink{
		gauges:     sync.Map{},
		summaries:  sync.Map{},
		counters:   sync.Map{},
		expiration: 60 * time.Second,
	}

	pusher := push.New(address, name).Collector(promSink)

	sink := &PrometheusPushSink{
		promSink,
		pusher,
		address,
		pushIterval,
		make(chan struct{}),
	}

	sink.flushMetrics()
	return sink, nil
}

func (s *PrometheusPushSink) flushMetrics() {
	ticker := time.NewTicker(s.pushInterval)

	go func() {
		for {
			select {
			case <-ticker.C:
				err := s.pusher.Push()
				if err != nil {
					log.Printf("[ERR] Error pushing to Prometheus! Err: %s", err)
				}
			case <-s.stopChan:
				ticker.Stop()
				return
			}
		}
	}()
}

func (s *PrometheusPushSink) Shutdown() {
	close(s.stopChan)
}
<commit_msg>add static metrics to PrometheusSink<commit_after>// +build go1.9

package prometheus

import (
	"fmt"
	"log"
	"math"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/armon/go-metrics"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/push"
)

var (
	// DefaultPrometheusOpts is the default set of options used when creating a
	// PrometheusSink.
	DefaultPrometheusOpts = PrometheusOpts{
		Expiration: 60 * time.Second,
	}
)

// PrometheusOpts is used to configure the Prometheus Sink
type PrometheusOpts struct {
	// Expiration is the duration a metric is valid for, after which it will be
	// untracked. If the value is zero, a metric is never expired.
	Expiration time.Duration
	Registerer prometheus.Registerer

	// Gauges, Summaries, and Counters allow us to pre-declare metrics by giving their Name and ConstLabels to the
	// PrometheusSink when it is created. Metrics declared in this way will be initialized at zero and will not be
	// deleted when their expiry is reached.
	// - Gauges and Summaries will be set to NaN when they expire.
	// - Counters continue to Collect their last known value.
	// Ex:
	// PrometheusOpts{
	//     Expiration: 10 * time.Second,
	//     Gauges: []PrometheusGauge{
	//         {
	//	         Name: []string{ "application", "component", "measurement"},
	//           ConstLabels: []metrics.Label{ { Name: "datacenter", Value: "dc1" }, },
	//         },
	//     },
	// }
	Gauges     []PrometheusGauge
	Summaries  []PrometheusSummary
	Counters   []PrometheusCounter
}

type PrometheusSink struct {
	// If these will ever be copied, they should be converted to *sync.Map values and initialized appropriately
	gauges     sync.Map
	summaries  sync.Map
	counters   sync.Map
	expiration time.Duration
}

type PrometheusGauge struct {
	Name []string
	ConstLabels []metrics.Label
	prometheus.Gauge
	updatedAt time.Time
	// canDelete is set if the metric is created during runtime so we know it's ephemeral and can delete it on expiry.
	canDelete bool
}

type PrometheusSummary struct {
	Name []string
	ConstLabels []metrics.Label
	prometheus.Summary
	updatedAt time.Time
	canDelete bool
}

type PrometheusCounter struct {
	Name []string
	ConstLabels []metrics.Label
	prometheus.Counter
	updatedAt time.Time
	canDelete bool
}

// NewPrometheusSink creates a new PrometheusSink using the default options.
func NewPrometheusSink() (*PrometheusSink, error) {
	return NewPrometheusSinkFrom(DefaultPrometheusOpts)
}

// NewPrometheusSinkFrom creates a new PrometheusSink using the passed options.
func NewPrometheusSinkFrom(opts PrometheusOpts) (*PrometheusSink, error) {
	sink := &PrometheusSink{
		gauges:     sync.Map{},
		summaries:  sync.Map{},
		counters:   sync.Map{},
		expiration: opts.Expiration,
	}

	initGauges(&sink.gauges, opts.Gauges)
	initSummaries(&sink.summaries, opts.Summaries)
	initCounters(&sink.counters, opts.Counters)

	reg := opts.Registerer
	if reg == nil {
		reg = prometheus.DefaultRegisterer
	}

	return sink, reg.Register(sink)
}

// Describe is needed to meet the Collector interface.
func (p *PrometheusSink) Describe(c chan<- *prometheus.Desc) {
	// We must emit some description otherwise an error is returned. This
	// description isn't shown to the user!
	prometheus.NewGauge(prometheus.GaugeOpts{Name: "Dummy", Help: "Dummy"}).Describe(c)
}

// Collect meets the collection interface and allows us to enforce our expiration
// logic to clean up ephemeral metrics if their value haven't been set for a
// duration exceeding our allowed expiration time.
func (p *PrometheusSink) Collect(c chan<- prometheus.Metric) {
	expire := p.expiration != 0
	now := time.Now()
	p.gauges.Range(func(k, v interface{}) bool {
		if v != nil {
			lastUpdate := v.(*PrometheusGauge).updatedAt
			if expire && lastUpdate.Add(p.expiration).Before(now) {
				if v.(*PrometheusGauge).canDelete {
					p.gauges.Delete(k)
					return true
				}
				// We have not observed the gauge this interval so we don't know its value.
				v.(*PrometheusGauge).Set(math.NaN())
			}
			v.(*PrometheusGauge).Collect(c)
		}
		return true
	})
	p.summaries.Range(func(k, v interface{}) bool {
		if v != nil {
			lastUpdate := v.(*PrometheusSummary).updatedAt
			if expire && lastUpdate.Add(p.expiration).Before(now) {
				if v.(*PrometheusSummary).canDelete {
					p.summaries.Delete(k)
					return true
				}
				// We have observed nothing in this interval.
				v.(*PrometheusSummary).Observe(math.NaN())
			}
			v.(*PrometheusSummary).Collect(c)
		}
		return true
	})
	p.counters.Range(func(k, v interface{}) bool {
		if v != nil {
			lastUpdate := v.(*PrometheusCounter).updatedAt
			if expire && lastUpdate.Add(p.expiration).Before(now) {
				if v.(PrometheusCounter).canDelete {
					p.counters.Delete(k)
					return true
				}
				// Counters remain at their previous value when not observed, so we do not set it to NaN.
			}
			v.(*PrometheusCounter).Collect(c)
		}
		return true
	})
}

func initGauges(m *sync.Map, gauges []PrometheusGauge) {
	for _, gauge := range gauges {
		key, hash := flattenKey(gauge.Name, gauge.ConstLabels)
		g := prometheus.NewGauge(prometheus.GaugeOpts{
			Name:        key,
			Help:        key,
			ConstLabels: prometheusLabels(gauge.ConstLabels),
		})
		g.Set(float64(0)) // Initialize at zero
		gauge.Gauge = g
		m.Store(hash, &gauge)
	}
	return
}

func initSummaries(m *sync.Map, summaries []PrometheusSummary) {
	for _, summary := range summaries {
		key, hash := flattenKey(summary.Name, summary.ConstLabels)
		s := prometheus.NewSummary(prometheus.SummaryOpts{
			Name:        key,
			Help:        key,
			ConstLabels: prometheusLabels(summary.ConstLabels),
		})
		s.Observe(float64(0)) // Initialize at zero
		summary.Summary = s
		m.Store(hash, &summary)
	}
	return
}

func initCounters(m *sync.Map, counters []PrometheusCounter) {
	for _, counter := range counters {
		key, hash := flattenKey(counter.Name, counter.ConstLabels)
		c := prometheus.NewCounter(prometheus.CounterOpts{
			Name:        key,
			Help:        key,
			ConstLabels: prometheusLabels(counter.ConstLabels),
		})
		c.Add(float64(0)) // Initialize at zero
		counter.Counter = c
		m.Store(hash, &counter)
	}
	return
}

var forbiddenChars = regexp.MustCompile("[ .=\\-/]")

func flattenKey(parts []string, labels []metrics.Label) (string, string) {
	key := strings.Join(parts, "_")
	key = forbiddenChars.ReplaceAllString(key, "_")

	hash := key
	for _, label := range labels {
		hash += fmt.Sprintf(";%s=%s", label.Name, label.Value)
	}

	return key, hash
}

func prometheusLabels(labels []metrics.Label) prometheus.Labels {
	l := make(prometheus.Labels)
	for _, label := range labels {
		l[label.Name] = label.Value
	}
	return l
}

func (p *PrometheusSink) SetGauge(parts []string, val float32) {
	p.SetGaugeWithLabels(parts, val, nil)
}

func (p *PrometheusSink) SetGaugeWithLabels(parts []string, val float32, labels []metrics.Label) {
	key, hash := flattenKey(parts, labels)
	pg, ok := p.gauges.Load(hash)

	// The sync.Map underlying gauges stores pointers to our structs. If we need to make updates,
	// rather than modifying the underlying value directly, which would be racy, we make a local
	// copy by dereferencing the pointer we get back, making the appropriate changes, and then
	// storing a pointer to our local copy. The underlying Prometheus types are threadsafe,
	// so there's no issues there. It's possible for racy updates to occur to the updatedAt
	// value, but since we're always setting it to time.Now(), it doesn't really matter.
	if ok {
		localGauge := *pg.(*PrometheusGauge)
		localGauge.Set(float64(val))
		localGauge.updatedAt = time.Now()
		p.gauges.Store(hash, &localGauge)

	// The gauge does not exist, create the gauge and allow it to be deleted
	} else {
		g := prometheus.NewGauge(prometheus.GaugeOpts{
			Name:        key,
			Help:        key,
			ConstLabels: prometheusLabels(labels),
		})
		g.Set(float64(val))
		pg = &PrometheusGauge{
			Gauge: g,
			updatedAt: time.Now(),
			canDelete: true,
		}
		p.gauges.Store(hash, pg)
	}
}

func (p *PrometheusSink) AddSample(parts []string, val float32) {
	p.AddSampleWithLabels(parts, val, nil)
}

func (p *PrometheusSink) AddSampleWithLabels(parts []string, val float32, labels []metrics.Label) {
	key, hash := flattenKey(parts, labels)
	ps, ok := p.summaries.Load(hash)

	// Does the summary already exist for this sample type?
	if ok {
		localSummary := *ps.(*PrometheusSummary)
		localSummary.Observe(float64(val))
		localSummary.updatedAt = time.Now()
		p.summaries.Store(hash, &localSummary)

	// The summary does not exist, create the Summary and allow it to be deleted
	} else {
		s := prometheus.NewSummary(prometheus.SummaryOpts{
			Name:        key,
			Help:        key,
			MaxAge:      10 * time.Second,
			ConstLabels: prometheusLabels(labels),
			Objectives:  map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001},
		})
		s.Observe(float64(val))
		ps = &PrometheusSummary{
			Summary: s,
			updatedAt: time.Now(),
			canDelete: true,
		}
		p.summaries.Store(hash, ps)
	}
}

// EmitKey is not implemented. Prometheus doesn’t offer a type for which an
// arbitrary number of values is retained, as Prometheus works with a pull
// model, rather than a push model.
func (p *PrometheusSink) EmitKey(key []string, val float32) {
}

func (p *PrometheusSink) IncrCounter(parts []string, val float32) {
	p.IncrCounterWithLabels(parts, val, nil)
}

func (p *PrometheusSink) IncrCounterWithLabels(parts []string, val float32, labels []metrics.Label) {
	key, hash := flattenKey(parts, labels)
	pc, ok := p.counters.Load(hash)

	// Does the counter exist?
	if ok {
		localCounter := *pc.(*PrometheusCounter)
		localCounter.Add(float64(val))
		localCounter.updatedAt = time.Now()
		p.counters.Store(hash, &localCounter)

	// The counter does not exist yet, create it and allow it to be deleted
	} else {
		c := prometheus.NewCounter(prometheus.CounterOpts{
			Name:        key,
			Help:        key,
			ConstLabels: prometheusLabels(labels),
		})
		c.Add(float64(val))
		pc = &PrometheusCounter{
			Counter: c,
			updatedAt: time.Now(),
			canDelete: true,
		}
		p.counters.Store(hash, pc)
	}
}

// PrometheusPushSink wraps a normal prometheus sink and provides an address and facilities to export it to an address
// on an interval.
type PrometheusPushSink struct {
	*PrometheusSink
	pusher       *push.Pusher
	address      string
	pushInterval time.Duration
	stopChan     chan struct{}
}

// NewPrometheusPushSink creates a PrometheusPushSink by taking an address, interval, and destination name.
func NewPrometheusPushSink(address string, pushInterval time.Duration, name string) (*PrometheusPushSink, error) {
	promSink := &PrometheusSink{
		gauges:     sync.Map{},
		summaries:  sync.Map{},
		counters:   sync.Map{},
		expiration: 60 * time.Second,
	}

	pusher := push.New(address, name).Collector(promSink)

	sink := &PrometheusPushSink{
		promSink,
		pusher,
		address,
		pushInterval,
		make(chan struct{}),
	}

	sink.flushMetrics()
	return sink, nil
}

func (s *PrometheusPushSink) flushMetrics() {
	ticker := time.NewTicker(s.pushInterval)

	go func() {
		for {
			select {
			case <-ticker.C:
				err := s.pusher.Push()
				if err != nil {
					log.Printf("[ERR] Error pushing to Prometheus! Err: %s", err)
				}
			case <-s.stopChan:
				ticker.Stop()
				return
			}
		}
	}()
}

func (s *PrometheusPushSink) Shutdown() {
	close(s.stopChan)
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package gce

import (
	"github.com/juju/errors"

	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/config"
)

type environProvider struct{}

var providerInstance = environProvider{}
var _ environs.EnvironProvider = providerInstance

func init() {
	// This will only happen in binaries that actually import this provider
	// somewhere. To enable a provider, import it in the "providers/all"
	// package; please do *not* import individual providers anywhere else,
	// except in direct tests for that provider.
	environs.RegisterProvider("gce", providerInstance)
}

func (environProvider) Open(cfg *config.Config) (environs.Environ, error) {
	// You should probably not change this method; prefer to cause SetConfig
	// to completely configure an environment, regardless of the initial state.
	// TODO(ericsnow) This leaves ecfg and gce nil, which might
	// result in nil pointer panics.  Either use new(environ) here
	// or handle nil for those two fields wherever we use them.
	env := &environ{name: cfg.Name()}
	if err := env.SetConfig(cfg); err != nil {
		return nil, errors.Trace(err)
	}
	return env, nil
}

func (environProvider) Prepare(ctx environs.BootstrapContext, cfg *config.Config) (environs.Environ, error) {
	env, err := providerInstance.Open(cfg)
	if err != nil {
		return nil, errors.Trace(err)
	}

	if ctx.ShouldVerifyCredentials() {
		if err := env.(*environ).gce.verifyCredentials(); err != nil {
			return nil, err
		}
	}
	return env, nil
}

func (environProvider) Validate(cfg, old *config.Config) (valid *config.Config, err error) {
	// You should almost certainly not change this method; if you need to change
	// how configs are validated, you should edit validateConfig itself, to ensure
	// that your checks are always applied.
	newEcfg, err := validateConfig(cfg, nil)
	if err != nil {
		return nil, errors.Annotate(err, "invalid config")
	}
	if old != nil {
		oldEcfg, err := validateConfig(old, nil)
		if err != nil {
			return nil, errors.Annotate(err, "invalid base config")
		}
		if newEcfg, err = validateConfig(cfg, oldEcfg.Config); err != nil {
			return nil, errors.Annotate(err, "invalid config change")
		}
	}
	return cfg.Apply(newEcfg.attrs)
}

func (environProvider) SecretAttrs(cfg *config.Config) (map[string]string, error) {
	ecfg, err := validateConfig(cfg, nil)
	if err != nil {
		return nil, errors.Trace(err)
	}
	secretAttrs := make(map[string]string, len(configSecretFields))
	for _, key := range configSecretFields {
		secretAttrs[key] = ecfg.attrs[key].(string)
	}
	return secretAttrs, nil
}

func (environProvider) BoilerplateConfig() string {
	// boilerplateConfig is kept in config.go, in the hope that people editing
	// config will keep it up to date.
	return boilerplateConfig
}
<commit_msg>provider/gce: Drop an old TODO.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package gce

import (
	"github.com/juju/errors"

	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/config"
)

type environProvider struct{}

var providerInstance = environProvider{}
var _ environs.EnvironProvider = providerInstance

func init() {
	// This will only happen in binaries that actually import this provider
	// somewhere. To enable a provider, import it in the "providers/all"
	// package; please do *not* import individual providers anywhere else,
	// except in direct tests for that provider.
	environs.RegisterProvider("gce", providerInstance)
}

func (environProvider) Open(cfg *config.Config) (environs.Environ, error) {
	// You should probably not change this method; prefer to cause SetConfig
	// to completely configure an environment, regardless of the initial state.
	env := &environ{name: cfg.Name()}
	if err := env.SetConfig(cfg); err != nil {
		return nil, errors.Trace(err)
	}
	return env, nil
}

func (environProvider) Prepare(ctx environs.BootstrapContext, cfg *config.Config) (environs.Environ, error) {
	env, err := providerInstance.Open(cfg)
	if err != nil {
		return nil, errors.Trace(err)
	}

	if ctx.ShouldVerifyCredentials() {
		if err := env.(*environ).gce.verifyCredentials(); err != nil {
			return nil, err
		}
	}
	return env, nil
}

func (environProvider) Validate(cfg, old *config.Config) (valid *config.Config, err error) {
	// You should almost certainly not change this method; if you need to change
	// how configs are validated, you should edit validateConfig itself, to ensure
	// that your checks are always applied.
	newEcfg, err := validateConfig(cfg, nil)
	if err != nil {
		return nil, errors.Annotate(err, "invalid config")
	}
	if old != nil {
		oldEcfg, err := validateConfig(old, nil)
		if err != nil {
			return nil, errors.Annotate(err, "invalid base config")
		}
		if newEcfg, err = validateConfig(cfg, oldEcfg.Config); err != nil {
			return nil, errors.Annotate(err, "invalid config change")
		}
	}
	return cfg.Apply(newEcfg.attrs)
}

func (environProvider) SecretAttrs(cfg *config.Config) (map[string]string, error) {
	ecfg, err := validateConfig(cfg, nil)
	if err != nil {
		return nil, errors.Trace(err)
	}
	secretAttrs := make(map[string]string, len(configSecretFields))
	for _, key := range configSecretFields {
		secretAttrs[key] = ecfg.attrs[key].(string)
	}
	return secretAttrs, nil
}

func (environProvider) BoilerplateConfig() string {
	// boilerplateConfig is kept in config.go, in the hope that people editing
	// config will keep it up to date.
	return boilerplateConfig
}
<|endoftext|>


<commit_before>package freedom

import (
	"net"
	"sync"

	"github.com/v2ray/v2ray-core"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
)

type FreedomConnection struct {
	packet v2net.Packet
}

func NewFreedomConnection(firstPacket v2net.Packet) *FreedomConnection {
	return &FreedomConnection{
		packet: firstPacket,
	}
}

func (vconn *FreedomConnection) Start(ray core.OutboundRay) error {
	conn, err := net.Dial(vconn.packet.Destination().Network(), vconn.packet.Destination().Address().String())
	log.Info("Freedom: Opening connection to %s", vconn.packet.Destination().String())
	if err != nil {
		if ray != nil {
			close(ray.OutboundOutput())
		}
		return log.Error("Freedom: Failed to open connection: %s : %v", vconn.packet.Destination().String(), err)
	}

	input := ray.OutboundInput()
	output := ray.OutboundOutput()
	var readMutex, writeMutex sync.Mutex
	readMutex.Lock()
	writeMutex.Lock()

	if chunk := vconn.packet.Chunk(); chunk != nil {
		conn.Write(chunk)
	}

	if !vconn.packet.MoreChunks() {
		if ray != nil {
			close(ray.OutboundOutput())
		}
		writeMutex.Unlock()
	} else {
		go dumpInput(conn, input, &writeMutex)
	}

	go dumpOutput(conn, output, &readMutex)

	go func() {
		writeMutex.Lock()
		if tcpConn, ok := conn.(*net.TCPConn); ok {
			tcpConn.CloseWrite()
		}
		readMutex.Lock()
		conn.Close()
	}()

	return nil
}

func dumpInput(conn net.Conn, input <-chan []byte, finish *sync.Mutex) {
	v2net.ChanToWriter(conn, input)
	finish.Unlock()
}

func dumpOutput(conn net.Conn, output chan<- []byte, finish *sync.Mutex) {
	v2net.ReaderToChan(output, conn)
	finish.Unlock()
	close(output)
}
<commit_msg>Update freedom connection for UDP<commit_after>package freedom

import (
	"net"
	"sync"

	"github.com/v2ray/v2ray-core"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
)

type FreedomConnection struct {
	packet v2net.Packet
}

func NewFreedomConnection(firstPacket v2net.Packet) *FreedomConnection {
	return &FreedomConnection{
		packet: firstPacket,
	}
}

func (vconn *FreedomConnection) Start(ray core.OutboundRay) error {
	conn, err := net.Dial(vconn.packet.Destination().Network(), vconn.packet.Destination().Address().String())
	log.Info("Freedom: Opening connection to %s", vconn.packet.Destination().String())
	if err != nil {
		if ray != nil {
			close(ray.OutboundOutput())
		}
		return log.Error("Freedom: Failed to open connection: %s : %v", vconn.packet.Destination().String(), err)
	}

	input := ray.OutboundInput()
	output := ray.OutboundOutput()
	var readMutex, writeMutex sync.Mutex
	readMutex.Lock()
	writeMutex.Lock()

	if chunk := vconn.packet.Chunk(); chunk != nil {
		conn.Write(chunk)
	}

	if !vconn.packet.MoreChunks() {
		if ray != nil {
			close(ray.OutboundOutput())
		}
		writeMutex.Unlock()
	} else {
		go dumpInput(conn, input, &writeMutex)
	}

	go dumpOutput(conn, output, &readMutex, vconn.packet.Destination().IsUDP())

	go func() {
		writeMutex.Lock()
		if tcpConn, ok := conn.(*net.TCPConn); ok {
			tcpConn.CloseWrite()
		}
		readMutex.Lock()
		conn.Close()
	}()

	return nil
}

func dumpInput(conn net.Conn, input <-chan []byte, finish *sync.Mutex) {
	v2net.ChanToWriter(conn, input)
	finish.Unlock()
}

func dumpOutput(conn net.Conn, output chan<- []byte, finish *sync.Mutex, udp bool) {
	defer finish.Unlock()
	defer close(output)

	response, err := v2net.ReadFrom(conn)
	if len(response) > 0 {
		output <- response
	}
	if err != nil {
		return
	}
	if udp {
		return
	}

	v2net.ReaderToChan(output, conn)
}
<|endoftext|>


<commit_before>package secio

import (
	"crypto/cipher"
	"errors"
	"fmt"
	"io"

	"crypto/hmac"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	proto "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/goprotobuf/proto"
	msgio "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-msgio"
	mpool "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-msgio/mpool"
)

// ErrMACInvalid signals that a MAC verification failed
var ErrMACInvalid = errors.New("MAC verification failed")

// bufPool is a ByteSlicePool for messages. we need buffers because (sadly)
// we cannot encrypt in place-- the user needs their buffer back.
var bufPool = mpool.ByteSlicePool

type etmWriter struct {
	// params
	pool mpool.Pool        // for the buffers with encrypted data
	msg  msgio.WriteCloser // msgio for knowing where boundaries lie
	str  cipher.Stream     // the stream cipher to encrypt with
	mac  HMAC              // the mac to authenticate data with
}

// NewETMWriter Encrypt-Then-MAC
func NewETMWriter(w io.Writer, s cipher.Stream, mac HMAC) msgio.WriteCloser {
	return &etmWriter{msg: msgio.NewWriter(w), str: s, mac: mac, pool: bufPool}
}

// Write writes passed in buffer as a single message.
func (w *etmWriter) Write(b []byte) (int, error) {
	if err := w.WriteMsg(b); err != nil {
		return 0, err
	}
	return len(b), nil
}

// WriteMsg writes the msg in the passed in buffer.
func (w *etmWriter) WriteMsg(b []byte) error {

	// encrypt.
	data := w.pool.Get(uint32(len(b))).([]byte)
	data = data[:len(b)] // the pool's buffer may be larger
	w.str.XORKeyStream(data, b)

	// log.Debugf("ENC plaintext (%d): %s %v", len(b), b, b)
	// log.Debugf("ENC ciphertext (%d): %s %v", len(data), data, data)

	// then, mac.
	if _, err := w.mac.Write(data); err != nil {
		return err
	}

	// Sum appends.
	data = w.mac.Sum(data)
	w.mac.Reset()
	// it's sad to append here. our buffers are -- hopefully -- coming from
	// a shared buffer pool, so the append may not actually cause allocation
	// one can only hope. i guess we'll see.

	return w.msg.WriteMsg(data)
}

func (w *etmWriter) Close() error {
	return w.msg.Close()
}

type etmReader struct {
	msgio.Reader
	io.Closer

	// buffer
	buf []byte

	// params
	msg msgio.ReadCloser // msgio for knowing where boundaries lie
	str cipher.Stream    // the stream cipher to encrypt with
	mac HMAC             // the mac to authenticate data with
}

// NewETMReader Encrypt-Then-MAC
func NewETMReader(r io.Reader, s cipher.Stream, mac HMAC) msgio.ReadCloser {
	return &etmReader{msg: msgio.NewReader(r), str: s, mac: mac}
}

func (r *etmReader) NextMsgLen() (int, error) {
	return r.msg.NextMsgLen()
}

func (r *etmReader) drainBuf(buf []byte) int {
	if r.buf == nil {
		return 0
	}

	n := copy(buf, r.buf)
	r.buf = r.buf[n:]
	return n
}

func (r *etmReader) Read(buf []byte) (int, error) {
	// first, check if we have anything in the buffer
	copied := r.drainBuf(buf)
	buf = buf[copied:]
	if copied > 0 {
		return copied, nil
		// return here to avoid complicating the rest...
		// user can call io.ReadFull.
	}

	// check the buffer has enough space for the next msg
	fullLen, err := r.msg.NextMsgLen()
	if err != nil {
		return 0, err
	}

	buf2 := buf
	changed := false
	// if not enough space, allocate a new buffer.
	if cap(buf) < fullLen {
		buf2 = make([]byte, fullLen)
		changed = true
	}
	buf2 = buf2[:fullLen]

	n, err := io.ReadFull(r.msg, buf2)
	if err != nil {
		return n, err
	}

	m, err := r.macCheckThenDecrypt(buf2)
	if err != nil {
		return 0, err
	}
	buf2 = buf2[:m]
	if !changed {
		return m, nil
	}

	n = copy(buf, buf2)
	if len(buf2) > len(buf) {
		r.buf = buf2[len(buf):] // had some left over? save it.
	}
	return n, nil
}

func (r *etmReader) ReadMsg() ([]byte, error) {
	msg, err := r.msg.ReadMsg()
	if err != nil {
		return nil, err
	}

	n, err := r.macCheckThenDecrypt(msg)
	if err != nil {
		return nil, err
	}
	return msg[:n], nil
}

func (r *etmReader) macCheckThenDecrypt(m []byte) (int, error) {
	l := len(m)
	if l < r.mac.size {
		return 0, fmt.Errorf("buffer (%d) shorter than MAC size (%d)", l, r.mac.size)
	}

	mark := l - r.mac.size
	data := m[:mark]
	macd := m[mark:]

	r.mac.Write(data)
	expected := r.mac.Sum(nil)
	r.mac.Reset()

	// check mac. if failed, return error.
	if !hmac.Equal(macd, expected) {
		log.Error("MAC Invalid:", expected, "!=", macd)
		return 0, ErrMACInvalid
	}

	// ok seems good. decrypt. (can decrypt in place, yay!)
	// log.Debugf("DEC ciphertext (%d): %s %v", len(data), data, data)
	r.str.XORKeyStream(data, data)
	// log.Debugf("DEC plaintext (%d): %s %v", len(data), data, data)

	return mark, nil
}

func (w *etmReader) Close() error {
	return w.msg.Close()
}

// ReleaseMsg signals a buffer can be reused.
func (r *etmReader) ReleaseMsg(b []byte) {
	r.msg.ReleaseMsg(b)
}

// writeMsgCtx is used by the
func writeMsgCtx(ctx context.Context, w msgio.Writer, msg proto.Message) ([]byte, error) {
	enc, err := proto.Marshal(msg)
	if err != nil {
		return nil, err
	}

	// write in a goroutine so we can exit when our context is cancelled.
	done := make(chan error)
	go func(m []byte) {
		err := w.WriteMsg(m)
		done <- err
	}(enc)

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case e := <-done:
		return enc, e
	}
}

func readMsgCtx(ctx context.Context, r msgio.Reader, p proto.Message) ([]byte, error) {
	var msg []byte

	// read in a goroutine so we can exit when our context is cancelled.
	done := make(chan error)
	go func() {
		var err error
		msg, err = r.ReadMsg()
		done <- err
	}()

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case e := <-done:
		if e != nil {
			return nil, e
		}
	}

	return msg, proto.Unmarshal(msg, p)
}
<commit_msg>secio: threadsafe<commit_after>package secio

import (
	"crypto/cipher"
	"errors"
	"fmt"
	"io"
	"sync"

	"crypto/hmac"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	proto "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/goprotobuf/proto"
	msgio "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-msgio"
	mpool "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-msgio/mpool"
)

// ErrMACInvalid signals that a MAC verification failed
var ErrMACInvalid = errors.New("MAC verification failed")

// bufPool is a ByteSlicePool for messages. we need buffers because (sadly)
// we cannot encrypt in place-- the user needs their buffer back.
var bufPool = mpool.ByteSlicePool

type etmWriter struct {
	// params
	pool mpool.Pool        // for the buffers with encrypted data
	msg  msgio.WriteCloser // msgio for knowing where boundaries lie
	str  cipher.Stream     // the stream cipher to encrypt with
	mac  HMAC              // the mac to authenticate data with

	sync.Mutex
}

// NewETMWriter Encrypt-Then-MAC
func NewETMWriter(w io.Writer, s cipher.Stream, mac HMAC) msgio.WriteCloser {
	return &etmWriter{msg: msgio.NewWriter(w), str: s, mac: mac, pool: bufPool}
}

// Write writes passed in buffer as a single message.
func (w *etmWriter) Write(b []byte) (int, error) {
	if err := w.WriteMsg(b); err != nil {
		return 0, err
	}
	return len(b), nil
}

// WriteMsg writes the msg in the passed in buffer.
func (w *etmWriter) WriteMsg(b []byte) error {
	w.Lock()
	defer w.Unlock()

	// encrypt.
	data := w.pool.Get(uint32(len(b))).([]byte)
	data = data[:len(b)] // the pool's buffer may be larger
	w.str.XORKeyStream(data, b)

	// log.Debugf("ENC plaintext (%d): %s %v", len(b), b, b)
	// log.Debugf("ENC ciphertext (%d): %s %v", len(data), data, data)

	// then, mac.
	if _, err := w.mac.Write(data); err != nil {
		return err
	}

	// Sum appends.
	data = w.mac.Sum(data)
	w.mac.Reset()
	// it's sad to append here. our buffers are -- hopefully -- coming from
	// a shared buffer pool, so the append may not actually cause allocation
	// one can only hope. i guess we'll see.

	return w.msg.WriteMsg(data)
}

func (w *etmWriter) Close() error {
	return w.msg.Close()
}

type etmReader struct {
	msgio.Reader
	io.Closer

	// buffer
	buf []byte

	// params
	msg msgio.ReadCloser // msgio for knowing where boundaries lie
	str cipher.Stream    // the stream cipher to encrypt with
	mac HMAC             // the mac to authenticate data with

	sync.Mutex
}

// NewETMReader Encrypt-Then-MAC
func NewETMReader(r io.Reader, s cipher.Stream, mac HMAC) msgio.ReadCloser {
	return &etmReader{msg: msgio.NewReader(r), str: s, mac: mac}
}

func (r *etmReader) NextMsgLen() (int, error) {
	return r.msg.NextMsgLen()
}

func (r *etmReader) drainBuf(buf []byte) int {
	if r.buf == nil {
		return 0
	}

	n := copy(buf, r.buf)
	r.buf = r.buf[n:]
	return n
}

func (r *etmReader) Read(buf []byte) (int, error) {
	r.Lock()
	defer r.Unlock()

	// first, check if we have anything in the buffer
	copied := r.drainBuf(buf)
	buf = buf[copied:]
	if copied > 0 {
		return copied, nil
		// return here to avoid complicating the rest...
		// user can call io.ReadFull.
	}

	// check the buffer has enough space for the next msg
	fullLen, err := r.msg.NextMsgLen()
	if err != nil {
		return 0, err
	}

	buf2 := buf
	changed := false
	// if not enough space, allocate a new buffer.
	if cap(buf) < fullLen {
		buf2 = make([]byte, fullLen)
		changed = true
	}
	buf2 = buf2[:fullLen]

	n, err := io.ReadFull(r.msg, buf2)
	if err != nil {
		return n, err
	}

	m, err := r.macCheckThenDecrypt(buf2)
	if err != nil {
		return 0, err
	}
	buf2 = buf2[:m]
	if !changed {
		return m, nil
	}

	n = copy(buf, buf2)
	if len(buf2) > len(buf) {
		r.buf = buf2[len(buf):] // had some left over? save it.
	}
	return n, nil
}

func (r *etmReader) ReadMsg() ([]byte, error) {
	r.Lock()
	defer r.Unlock()

	msg, err := r.msg.ReadMsg()
	if err != nil {
		return nil, err
	}

	n, err := r.macCheckThenDecrypt(msg)
	if err != nil {
		return nil, err
	}
	return msg[:n], nil
}

func (r *etmReader) macCheckThenDecrypt(m []byte) (int, error) {
	l := len(m)
	if l < r.mac.size {
		return 0, fmt.Errorf("buffer (%d) shorter than MAC size (%d)", l, r.mac.size)
	}

	mark := l - r.mac.size
	data := m[:mark]
	macd := m[mark:]

	r.mac.Write(data)
	expected := r.mac.Sum(nil)
	r.mac.Reset()

	// check mac. if failed, return error.
	if !hmac.Equal(macd, expected) {
		log.Error("MAC Invalid:", expected, "!=", macd)
		return 0, ErrMACInvalid
	}

	// ok seems good. decrypt. (can decrypt in place, yay!)
	// log.Debugf("DEC ciphertext (%d): %s %v", len(data), data, data)
	r.str.XORKeyStream(data, data)
	// log.Debugf("DEC plaintext (%d): %s %v", len(data), data, data)

	return mark, nil
}

func (w *etmReader) Close() error {
	return w.msg.Close()
}

// ReleaseMsg signals a buffer can be reused.
func (r *etmReader) ReleaseMsg(b []byte) {
	r.msg.ReleaseMsg(b)
}

// writeMsgCtx is used by the
func writeMsgCtx(ctx context.Context, w msgio.Writer, msg proto.Message) ([]byte, error) {
	enc, err := proto.Marshal(msg)
	if err != nil {
		return nil, err
	}

	// write in a goroutine so we can exit when our context is cancelled.
	done := make(chan error)
	go func(m []byte) {
		err := w.WriteMsg(m)
		done <- err
	}(enc)

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case e := <-done:
		return enc, e
	}
}

func readMsgCtx(ctx context.Context, r msgio.Reader, p proto.Message) ([]byte, error) {
	var msg []byte

	// read in a goroutine so we can exit when our context is cancelled.
	done := make(chan error)
	go func() {
		var err error
		msg, err = r.ReadMsg()
		done <- err
	}()

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case e := <-done:
		if e != nil {
			return nil, e
		}
	}

	return msg, proto.Unmarshal(msg, p)
}
<|endoftext|>