<commit_before>/*
 * Copyright (c) Clinton Freeman 2016
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
 * associated documentation files (the "Software"), to deal in the Software without restriction,
 * including without limitation the rights to use, copy, modify, merge, publish, distribute,
 * sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all copies or
 * substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
 * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

package main

import (
	"github.com/akualab/dmx"
	"github.com/kidoman/embd"
	_ "github.com/kidoman/embd/host/all"
	"log"
	"time"
)

// WeatherMachine holds connections to everything we need to manipulate the installation.
type WeatherMachine struct {
	stop    chan bool     // Channel for stopping the control elements of the installation.
	dmx     *dmx.DMX      // The DMX connection for writting messages to the Smoke machine and lights.
	config  Configuration // The configuration element for the installation.
	lastRun time.Time     // The last time the installation was run.
}

// ****************************************************************************
// ****************************************************************************
// Functions for manipulating the installation state; idle, warmup and running.
// ****************************************************************************
// ****************************************************************************

// stateFunctions are used to manipulate the WeatherMachine through the various states.
type stateFn func(state *WeatherMachine, msg HRMsg) stateFn

// idle is the state the weathermachine enters when sitting alone, with no one interacting with it.
func idle(state *WeatherMachine, msg HRMsg) (sF stateFn) {
	if msg.Contact {
		log.Printf("INFO: entering warmup")
		enableLight(state.config.S1Beat, state.config, state.dmx)
		go enablePump(state.config, state.stop)

		return warmup // skin contact has been made, enable light and enter warmup.
	}

	return idle // remain idle.
}

// warmup is the state the weathermachine enters when someone first touches it.
func warmup(state *WeatherMachine, msg HRMsg) stateFn {
	if msg.Contact && msg.HeartRate > 0 {
		// Wait for the fog to clear from the last run before running again.
		log.Printf("INFO: Waiting for fog to clear")
		d := (int64(state.config.FanDuration)*1000000) - time.Since(state.lastRun).Nanoseconds()
		time.Sleep(time.Nanosecond * time.Duration(d))

		log.Printf("INFO: entering running")
		go enableLightPulse(state.config, msg.HeartRate, state.stop, state.dmx)
		go enableSmoke(state.config, state.stop, state.dmx)
		go enableFan(state.config, state.stop)

		return running // skin contact and heart rate recieved, start the installation.
	} else if !msg.Contact {
		state.stop <- true // Pump starts at initial contact. If we lost contact between
		// then and now we need to shut it down.
		state.lastRun = time.Now()

		disableLight(state.config, state.dmx)

		log.Printf("INFO: entering idle")
		return idle // skin contact lost. Return to idle.
	}

	return warmup
}

// running is the state the weathermachine enters when someone is engaging with it.
func running(state *WeatherMachine, msg HRMsg) stateFn {
	if !msg.Contact {
		state.stop <- true
		state.stop <- true
		state.stop <- true
		state.stop <- true

		log.Printf("INFO: entering idle")
		return idle // skin contact lost. Return to idle.
	}

	return running // Keep the installation running.
}

// ****************************************************************************
// ****************************************************************************
// Functions for manipulating the physical installation; lights, smoke and fan.
// ****************************************************************************
// ****************************************************************************

// enableLight turns on the light via the supplied DMX connection 'dmx' with the supplied colour 'l'.
func enableLight(l LightColour, c Configuration, dmx *dmx.DMX) {
	log.Printf("INFO: Light on")
	embd.DigitalWrite(c.GPIOPinLight, embd.High)
	dmx.SetChannel(4, byte(l.Red))
	dmx.SetChannel(5, byte(l.Green))
	dmx.SetChannel(6, byte(l.Blue))
	dmx.SetChannel(7, byte(l.Amber))
	dmx.SetChannel(8, byte(l.Dimmer))
	dmx.Render()
}

// disableLight turns off the light via the supplied DMX connection 'dmx'.
func disableLight(c Configuration, dmx *dmx.DMX) {
	log.Printf("INFO: Light off")
	embd.DigitalWrite(c.GPIOPinLight, embd.Low)
	dmx.SetChannel(4, 0)
	dmx.SetChannel(5, 0)
	dmx.SetChannel(6, 0)
	dmx.SetChannel(7, 0)
	dmx.SetChannel(8, 0)
	dmx.Render()
}

// pulseLight pulses the light for a fixed duration.
func pulseLight(c Configuration, dmx *dmx.DMX) {
	enableLight(c.S1Beat, c, dmx)
	time.Sleep(time.Millisecond * time.Duration(c.S1Duration))
	disableLight(c, dmx)

	time.Sleep(time.Millisecond * time.Duration(c.S1Pause))

	enableLight(c.S2Beat, c, dmx)
	time.Sleep(time.Millisecond * time.Duration(c.S2Duration))
	disableLight(c, dmx)
}

// enableLightPulse starts the light pulsing by the frequency defined by hr. The light remains
// pulsing till being notified to stop on d.
func enableLightPulse(c Configuration, hr int, d chan bool, dmx *dmx.DMX) {
	// Perform the first heart beat straight away.
	pulseLight(c, dmx)

	dt := int((60000.0 / float32(hr)) * c.BeatRate)
	ticker := time.NewTicker(time.Millisecond * time.Duration(dt)).C

	// Sharp fixed length, pulse of light with variable off gap depending on HR.
	for {
		select {
		case <-ticker:
			pulseLight(c, dmx)

		case <-d:
			return
		}
	}
}

// pulsePump runs the pump for the duration specified in the configuration.
func pulsePump(c Configuration) {
	log.Printf("INFO: Pump on")
	embd.DigitalWrite(c.GPIOPinPump, embd.High)

	time.Sleep(time.Millisecond * time.Duration(c.PumpDuration))

	log.Printf("INFO: Pump Off")
	embd.DigitalWrite(c.GPIOPinPump, embd.Low)
}

// enablePump switches the relay on for the water pump after DeltaTPump milliseconds have expired
// in the configuration. Pump remains on till being notified to stop on d.
func enablePump(c Configuration, d chan bool) {
	dt := time.NewTimer(time.Millisecond * time.Duration(c.DeltaTPump)).C
	var ticker <-chan time.Time

	for {
		select {
		case <-dt:
			pulsePump(c)
			ticker = time.NewTicker(time.Millisecond * time.Duration(c.PumpInterval)).C

		case <-ticker:
			pulsePump(c)

		case <-d:
			return
		}
	}
}

// enableFan switches the relay on for the fan after DeltaTFan milliseconds have expired
// in the configuration. Fan remains on till being notified to stop on d.
func enableFan(c Configuration, d chan bool) {
	dt := time.NewTimer(time.Millisecond * time.Duration(c.DeltaTFan)).C

	for {
		select {
		case <-dt:
			log.Printf("INFO: Fan On")
			embd.DigitalWrite(c.GPIOPinFan, embd.High)

		case <-d:
			// Wait for the fan duration to clear the smoke chamber.
			ft := time.NewTimer(time.Millisecond * time.Duration(c.FanDuration)).C
			<-ft
			log.Printf("INFO: Fan Off")
			embd.DigitalWrite(c.GPIOPinFan, embd.Low)
			return
		}
	}
}

// puffSmoke enables the smoke machine via the supplied DMX connection 'dmx' for a period of
// time and intentsity supplied in configuration.
func puffSmoke(c Configuration, dmx *dmx.DMX) {
	log.Printf("INFO: Smoke on")
	dmx.SetChannel(1, byte(c.SmokeVolume))
	dmx.Render()

	time.Sleep(time.Millisecond * time.Duration(c.SmokeDuration))

	log.Printf("INFO: Smoke off")
	dmx.SetChannel(1, 0)
	dmx.Render()
}

// enableSmoke enages the DMX smoke machine by the SmokeVolume amount in the configuration.
// Smoke Machine remains on till being notified to stop on d.
func enableSmoke(c Configuration, d chan bool, dmx *dmx.DMX) {
	dt := time.NewTimer(time.Millisecond * time.Duration(c.DeltaTSmoke)).C
	var ticker <-chan time.Time

	for {
		select {
		case <-dt:
			puffSmoke(c, dmx)
			ticker = time.NewTicker(time.Millisecond * time.Duration(c.SmokeInterval)).C

		case <-ticker:
			puffSmoke(c, dmx)

		case <-d:
			return
		}
	}
}
<commit_msg>Update the last time the installation ran.<commit_after>/*
 * Copyright (c) Clinton Freeman 2016
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
 * associated documentation files (the "Software"), to deal in the Software without restriction,
 * including without limitation the rights to use, copy, modify, merge, publish, distribute,
 * sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all copies or
 * substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
 * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

package main

import (
	"github.com/akualab/dmx"
	"github.com/kidoman/embd"
	_ "github.com/kidoman/embd/host/all"
	"log"
	"time"
)

// WeatherMachine holds connections to everything we need to manipulate the installation.
type WeatherMachine struct {
	stop    chan bool     // Channel for stopping the control elements of the installation.
	dmx     *dmx.DMX      // The DMX connection for writting messages to the Smoke machine and lights.
	config  Configuration // The configuration element for the installation.
	lastRun time.Time     // The last time the installation was run.
}

// ****************************************************************************
// ****************************************************************************
// Functions for manipulating the installation state; idle, warmup and running.
// ****************************************************************************
// ****************************************************************************

// stateFunctions are used to manipulate the WeatherMachine through the various states.
type stateFn func(state *WeatherMachine, msg HRMsg) stateFn

// idle is the state the weathermachine enters when sitting alone, with no one interacting with it.
func idle(state *WeatherMachine, msg HRMsg) (sF stateFn) {
	if msg.Contact {
		log.Printf("INFO: entering warmup")
		enableLight(state.config.S1Beat, state.config, state.dmx)
		go enablePump(state.config, state.stop)

		return warmup // skin contact has been made, enable light and enter warmup.
	}

	return idle // remain idle.
}

// warmup is the state the weathermachine enters when someone first touches it.
func warmup(state *WeatherMachine, msg HRMsg) stateFn {
	if msg.Contact && msg.HeartRate > 0 {
		// Wait for the fog to clear from the last run before running again.
		log.Printf("INFO: Waiting for fog to clear")
		d := (int64(state.config.FanDuration)*1000000) - time.Since(state.lastRun).Nanoseconds()
		time.Sleep(time.Nanosecond * time.Duration(d))

		log.Printf("INFO: entering running")
		go enableLightPulse(state.config, msg.HeartRate, state.stop, state.dmx)
		go enableSmoke(state.config, state.stop, state.dmx)
		go enableFan(state.config, state.stop)

		return running // skin contact and heart rate recieved, start the installation.
	} else if !msg.Contact {
		state.stop <- true // Pump starts at initial contact. If we lost contact between
		// then and now we need to shut it down.
		state.lastRun = time.Now()

		disableLight(state.config, state.dmx)

		log.Printf("INFO: entering idle")
		return idle // skin contact lost. Return to idle.
	}

	return warmup
}

// running is the state the weathermachine enters when someone is engaging with it.
func running(state *WeatherMachine, msg HRMsg) stateFn {
	if !msg.Contact {
		state.stop <- true
		state.stop <- true
		state.stop <- true
		state.stop <- true
		state.lastRun = time.Now()

		log.Printf("INFO: entering idle")
		return idle // skin contact lost. Return to idle.
	}

	return running // Keep the installation running.
}

// ****************************************************************************
// ****************************************************************************
// Functions for manipulating the physical installation; lights, smoke and fan.
// ****************************************************************************
// ****************************************************************************

// enableLight turns on the light via the supplied DMX connection 'dmx' with the supplied colour 'l'.
func enableLight(l LightColour, c Configuration, dmx *dmx.DMX) {
	log.Printf("INFO: Light on")
	embd.DigitalWrite(c.GPIOPinLight, embd.High)
	dmx.SetChannel(4, byte(l.Red))
	dmx.SetChannel(5, byte(l.Green))
	dmx.SetChannel(6, byte(l.Blue))
	dmx.SetChannel(7, byte(l.Amber))
	dmx.SetChannel(8, byte(l.Dimmer))
	dmx.Render()
}

// disableLight turns off the light via the supplied DMX connection 'dmx'.
func disableLight(c Configuration, dmx *dmx.DMX) {
	log.Printf("INFO: Light off")
	embd.DigitalWrite(c.GPIOPinLight, embd.Low)
	dmx.SetChannel(4, 0)
	dmx.SetChannel(5, 0)
	dmx.SetChannel(6, 0)
	dmx.SetChannel(7, 0)
	dmx.SetChannel(8, 0)
	dmx.Render()
}

// pulseLight pulses the light for a fixed duration.
func pulseLight(c Configuration, dmx *dmx.DMX) {
	enableLight(c.S1Beat, c, dmx)
	time.Sleep(time.Millisecond * time.Duration(c.S1Duration))
	disableLight(c, dmx)

	time.Sleep(time.Millisecond * time.Duration(c.S1Pause))

	enableLight(c.S2Beat, c, dmx)
	time.Sleep(time.Millisecond * time.Duration(c.S2Duration))
	disableLight(c, dmx)
}

// enableLightPulse starts the light pulsing by the frequency defined by hr. The light remains
// pulsing till being notified to stop on d.
func enableLightPulse(c Configuration, hr int, d chan bool, dmx *dmx.DMX) {
	// Perform the first heart beat straight away.
	pulseLight(c, dmx)

	dt := int((60000.0 / float32(hr)) * c.BeatRate)
	ticker := time.NewTicker(time.Millisecond * time.Duration(dt)).C

	// Sharp fixed length, pulse of light with variable off gap depending on HR.
	for {
		select {
		case <-ticker:
			pulseLight(c, dmx)

		case <-d:
			return
		}
	}
}

// pulsePump runs the pump for the duration specified in the configuration.
func pulsePump(c Configuration) {
	log.Printf("INFO: Pump on")
	embd.DigitalWrite(c.GPIOPinPump, embd.High)

	time.Sleep(time.Millisecond * time.Duration(c.PumpDuration))

	log.Printf("INFO: Pump Off")
	embd.DigitalWrite(c.GPIOPinPump, embd.Low)
}

// enablePump switches the relay on for the water pump after DeltaTPump milliseconds have expired
// in the configuration. Pump remains on till being notified to stop on d.
func enablePump(c Configuration, d chan bool) {
	dt := time.NewTimer(time.Millisecond * time.Duration(c.DeltaTPump)).C
	var ticker <-chan time.Time

	for {
		select {
		case <-dt:
			pulsePump(c)
			ticker = time.NewTicker(time.Millisecond * time.Duration(c.PumpInterval)).C

		case <-ticker:
			pulsePump(c)

		case <-d:
			return
		}
	}
}

// enableFan switches the relay on for the fan after DeltaTFan milliseconds have expired
// in the configuration. Fan remains on till being notified to stop on d.
func enableFan(c Configuration, d chan bool) {
	dt := time.NewTimer(time.Millisecond * time.Duration(c.DeltaTFan)).C

	for {
		select {
		case <-dt:
			log.Printf("INFO: Fan On")
			embd.DigitalWrite(c.GPIOPinFan, embd.High)

		case <-d:
			// Wait for the fan duration to clear the smoke chamber.
			ft := time.NewTimer(time.Millisecond * time.Duration(c.FanDuration)).C
			<-ft
			log.Printf("INFO: Fan Off")
			embd.DigitalWrite(c.GPIOPinFan, embd.Low)
			return
		}
	}
}

// puffSmoke enables the smoke machine via the supplied DMX connection 'dmx' for a period of
// time and intentsity supplied in configuration.
func puffSmoke(c Configuration, dmx *dmx.DMX) {
	log.Printf("INFO: Smoke on")
	dmx.SetChannel(1, byte(c.SmokeVolume))
	dmx.Render()

	time.Sleep(time.Millisecond * time.Duration(c.SmokeDuration))

	log.Printf("INFO: Smoke off")
	dmx.SetChannel(1, 0)
	dmx.Render()
}

// enableSmoke enages the DMX smoke machine by the SmokeVolume amount in the configuration.
// Smoke Machine remains on till being notified to stop on d.
func enableSmoke(c Configuration, d chan bool, dmx *dmx.DMX) {
	dt := time.NewTimer(time.Millisecond * time.Duration(c.DeltaTSmoke)).C
	var ticker <-chan time.Time

	for {
		select {
		case <-dt:
			puffSmoke(c, dmx)
			ticker = time.NewTicker(time.Millisecond * time.Duration(c.SmokeInterval)).C

		case <-ticker:
			puffSmoke(c, dmx)

		case <-d:
			return
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 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 prometheus

import (
	"testing"
)

func BenchmarkCounterWithLabelValues(b *testing.B) {
	m := NewCounterVec(
		CounterOpts{
			Name: "benchmark_counter",
			Help: "A counter to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.WithLabelValues("eins", "zwei", "drei").Inc()
	}
}

func BenchmarkCounterWithMappedLabels(b *testing.B) {
	m := NewCounterVec(
		CounterOpts{
			Name: "benchmark_counter",
			Help: "A counter to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.With(Labels{"two": "zwei", "one": "eins", "three": "drei"}).Inc()
	}
}

func BenchmarkCounterWithPreparedMappedLabels(b *testing.B) {
	m := NewCounterVec(
		CounterOpts{
			Name: "benchmark_counter",
			Help: "A counter to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	labels := Labels{"two": "zwei", "one": "eins", "three": "drei"}
	for i := 0; i < b.N; i++ {
		m.With(labels).Inc()
	}
}

func BenchmarkCounterNoLabels(b *testing.B) {
	m := NewCounter(CounterOpts{
		Name: "benchmark_counter",
		Help: "A counter to benchmark it.",
	})
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.Inc()
	}
}

func BenchmarkGaugeWithLabelValues(b *testing.B) {
	m := NewGaugeVec(
		GaugeOpts{
			Name: "benchmark_gauge",
			Help: "A gauge to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.WithLabelValues("eins", "zwei", "drei").Set(3.1415)
	}
}

func BenchmarkGaugeNoLabels(b *testing.B) {
	m := NewGauge(GaugeOpts{
		Name: "benchmark_gauge",
		Help: "A gauge to benchmark it.",
	})
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.Set(3.1415)
	}
}

func BenchmarkSummaryWithLabelValues(b *testing.B) {
	m := NewSummaryVec(
		SummaryOpts{
			Name: "benchmark_summary",
			Help: "A summary to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.WithLabelValues("eins", "zwei", "drei").Observe(3.1415)
	}
}

func BenchmarkSummaryNoLabels(b *testing.B) {
	m := NewSummary(SummaryOpts{
		Name: "benchmark_summary",
		Help: "A summary to benchmark it.",
	},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.Observe(3.1415)
	}
}

func BenchmarkHistogramWithLabelValues(b *testing.B) {
	m := NewHistogramVec(
		HistogramOpts{
			Name: "benchmark_histogram",
			Help: "A histogram to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.WithLabelValues("eins", "zwei", "drei").Observe(3.1415)
	}
}

func BenchmarkHistogramNoLabels(b *testing.B) {
	m := NewHistogram(HistogramOpts{
		Name: "benchmark_histogram",
		Help: "A histogram to benchmark it.",
	},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.Observe(3.1415)
	}
}
<commit_msg>concurrency benchmark<commit_after>// Copyright 2014 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 prometheus

import (
	"sync"
	"testing"
)

func BenchmarkCounterWithLabelValues(b *testing.B) {
	m := NewCounterVec(
		CounterOpts{
			Name: "benchmark_counter",
			Help: "A counter to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.WithLabelValues("eins", "zwei", "drei").Inc()
	}
}

func BenchmarkCounterWithLabelValuesConcurrent(b *testing.B) {
	m := NewCounterVec(
		CounterOpts{
			Name: "benchmark_counter",
			Help: "A counter to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	wg := sync.WaitGroup{}
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			for j := 0; j < b.N/10; j++ {
				m.WithLabelValues("eins", "zwei", "drei").Inc()
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

func BenchmarkCounterWithMappedLabels(b *testing.B) {
	m := NewCounterVec(
		CounterOpts{
			Name: "benchmark_counter",
			Help: "A counter to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.With(Labels{"two": "zwei", "one": "eins", "three": "drei"}).Inc()
	}
}

func BenchmarkCounterWithPreparedMappedLabels(b *testing.B) {
	m := NewCounterVec(
		CounterOpts{
			Name: "benchmark_counter",
			Help: "A counter to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	labels := Labels{"two": "zwei", "one": "eins", "three": "drei"}
	for i := 0; i < b.N; i++ {
		m.With(labels).Inc()
	}
}

func BenchmarkCounterNoLabels(b *testing.B) {
	m := NewCounter(CounterOpts{
		Name: "benchmark_counter",
		Help: "A counter to benchmark it.",
	})
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.Inc()
	}
}

func BenchmarkGaugeWithLabelValues(b *testing.B) {
	m := NewGaugeVec(
		GaugeOpts{
			Name: "benchmark_gauge",
			Help: "A gauge to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.WithLabelValues("eins", "zwei", "drei").Set(3.1415)
	}
}

func BenchmarkGaugeNoLabels(b *testing.B) {
	m := NewGauge(GaugeOpts{
		Name: "benchmark_gauge",
		Help: "A gauge to benchmark it.",
	})
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.Set(3.1415)
	}
}

func BenchmarkSummaryWithLabelValues(b *testing.B) {
	m := NewSummaryVec(
		SummaryOpts{
			Name: "benchmark_summary",
			Help: "A summary to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.WithLabelValues("eins", "zwei", "drei").Observe(3.1415)
	}
}

func BenchmarkSummaryNoLabels(b *testing.B) {
	m := NewSummary(SummaryOpts{
		Name: "benchmark_summary",
		Help: "A summary to benchmark it.",
	},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.Observe(3.1415)
	}
}

func BenchmarkHistogramWithLabelValues(b *testing.B) {
	m := NewHistogramVec(
		HistogramOpts{
			Name: "benchmark_histogram",
			Help: "A histogram to benchmark it.",
		},
		[]string{"one", "two", "three"},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.WithLabelValues("eins", "zwei", "drei").Observe(3.1415)
	}
}

func BenchmarkHistogramNoLabels(b *testing.B) {
	m := NewHistogram(HistogramOpts{
		Name: "benchmark_histogram",
		Help: "A histogram to benchmark it.",
	},
	)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.Observe(3.1415)
	}
}
<|endoftext|>


<commit_before>package service

import (
	"errors"
	"fmt"
	"sync"

	msg "github.com/jbenet/go-ipfs/net/message"
	u "github.com/jbenet/go-ipfs/util"
	ctxc "github.com/jbenet/go-ipfs/util/ctxcloser"

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

var log = u.Logger("service")

// ErrNoResponse is returned by Service when a Request did not get a response,
// and no other error happened
var ErrNoResponse = errors.New("no response to request")

// Handler is an interface that objects must implement in order to handle
// a service's requests.
type Handler interface {

	// HandleMessage receives an incoming message, and potentially returns
	// a response message to send back.
	HandleMessage(context.Context, msg.NetMessage) msg.NetMessage
}

// Sender interface for network services.
type Sender interface {
	// SendMessage sends out a given message, without expecting a response.
	SendMessage(ctx context.Context, m msg.NetMessage) error

	// SendRequest sends out a given message, and awaits a response.
	// Set Deadlines or cancellations in the context.Context you pass in.
	SendRequest(ctx context.Context, m msg.NetMessage) (msg.NetMessage, error)
}

// Service is an interface for a net resource with both outgoing (sender) and
// incomig (SetHandler) requests.
type Service interface {
	Sender
	ctxc.ContextCloser

	// GetPipe
	GetPipe() *msg.Pipe

	// SetHandler assigns the request Handler for this service.
	SetHandler(Handler)
	GetHandler() Handler
}

// Service is a networking component that protocols can use to multiplex
// messages over the same channel, and to issue + handle requests.
type service struct {
	// Handler is the object registered to handle incoming requests.
	Handler     Handler
	HandlerLock sync.RWMutex

	// Requests are all the pending requests on this service.
	Requests     RequestMap
	RequestsLock sync.RWMutex

	// Message Pipe (connected to the outside world)
	*msg.Pipe
	ctxc.ContextCloser
}

// NewService creates a service object with given type ID and Handler
func NewService(ctx context.Context, h Handler) Service {
	s := &service{
		Handler:       h,
		Requests:      RequestMap{},
		Pipe:          msg.NewPipe(10),
		ContextCloser: ctxc.NewContextCloser(ctx, nil),
	}

	s.Children().Add(1)
	go s.handleIncomingMessages()
	return s
}

// GetPipe implements the mux.Protocol interface
func (s *service) GetPipe() *msg.Pipe {
	return s.Pipe
}

// sendMessage sends a message out (actual leg work. SendMessage is to export w/o rid)
func (s *service) sendMessage(ctx context.Context, m msg.NetMessage, rid RequestID) error {

	// serialize ServiceMessage wrapper
	data, err := wrapData(m.Data(), rid)
	if err != nil {
		return err
	}

	// log.Debug("Service send message [to = %s]", m.Peer())

	// send message
	m2 := msg.New(m.Peer(), data)
	select {
	case s.Outgoing <- m2:
	case <-ctx.Done():
		return ctx.Err()
	}

	return nil
}

// SendMessage sends a message out
func (s *service) SendMessage(ctx context.Context, m msg.NetMessage) error {
	return s.sendMessage(ctx, m, nil)
}

// SendRequest sends a request message out and awaits a response.
func (s *service) SendRequest(ctx context.Context, m msg.NetMessage) (msg.NetMessage, error) {

	// check if we should bail given our contexts
	select {
	default:
	case <-s.Closing():
		return nil, fmt.Errorf("service closed: %s", s.Context().Err())
	case <-ctx.Done():
		return nil, ctx.Err()
	}

	// create a request
	r, err := NewRequest(m.Peer().ID())
	if err != nil {
		return nil, err
	}

	// register Request
	s.RequestsLock.Lock()
	s.Requests[r.Key()] = r
	s.RequestsLock.Unlock()

	// defer deleting this request
	defer func() {
		s.RequestsLock.Lock()
		delete(s.Requests, r.Key())
		s.RequestsLock.Unlock()
	}()

	// check if we should bail after waiting for mutex
	select {
	default:
	case <-s.Closing():
		return nil, fmt.Errorf("service closed: %s", s.Context().Err())
	case <-ctx.Done():
		return nil, ctx.Err()
	}

	// Send message
	s.sendMessage(ctx, m, r.ID)

	// wait for response
	m = nil
	err = nil
	select {
	case m = <-r.Response:
	case <-s.Closed():
		err = fmt.Errorf("service closed: %s", s.Context().Err())
	case <-ctx.Done():
		err = ctx.Err()
	}

	if m == nil {
		return nil, ErrNoResponse
	}

	return m, err
}

// handleIncoming consumes the messages on the s.Incoming channel and
// routes them appropriately (to requests, or handler).
func (s *service) handleIncomingMessages() {
	defer s.Children().Done()

	for {
		select {
		case m, more := <-s.Incoming:
			if !more {
				return
			}
			s.Children().Add(1)
			go s.handleIncomingMessage(m)

		case <-s.Closing():
			return
		}
	}
}

func (s *service) handleIncomingMessage(m msg.NetMessage) {
	defer s.Children().Done()

	// unwrap the incoming message
	data, rid, err := unwrapData(m.Data())
	if err != nil {
		log.Errorf("service de-serializing error: %v", err)
		return
	}

	m2 := msg.New(m.Peer(), data)

	// if it's a request (or has no RequestID), handle it
	if rid == nil || rid.IsRequest() {
		handler := s.GetHandler()
		if handler == nil {
			log.Errorf("service dropped msg: %v", m)
			return // no handler, drop it.
		}

		// should this be "go HandleMessage ... ?"
		r1 := handler.HandleMessage(s.Context(), m2)

		// if handler gave us a response, send it back out!
		if r1 != nil {
			err := s.sendMessage(s.Context(), r1, rid.Response())
			if err != nil {
				log.Errorf("error sending response message: %v", err)
			}
		}
		return
	}

	// Otherwise, it is a response. handle it.
	if !rid.IsResponse() {
		log.Errorf("RequestID should identify a response here.")
	}

	key := RequestKey(m.Peer().ID(), RequestID(rid))
	s.RequestsLock.RLock()
	r, found := s.Requests[key]
	s.RequestsLock.RUnlock()

	if !found {
		log.Errorf("no request key %v (timeout?)", []byte(key))
		return
	}

	select {
	case r.Response <- m2:
	case <-s.Closing():
	}
}

// SetHandler assigns the request Handler for this service.
func (s *service) SetHandler(h Handler) {
	s.HandlerLock.Lock()
	defer s.HandlerLock.Unlock()
	s.Handler = h
}

// GetHandler returns the request Handler for this service.
func (s *service) GetHandler() Handler {
	s.HandlerLock.RLock()
	defer s.HandlerLock.RUnlock()
	return s.Handler
}
<commit_msg>stash<commit_after>package service

import (
	"errors"
	"fmt"
	"sync"

	msg "github.com/jbenet/go-ipfs/net/message"
	u "github.com/jbenet/go-ipfs/util"

	ctxgroup "github.com/jbenet/go-ctxgroup"
	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	router "github.com/jbenet/go-router"
)

var log = u.Logger("service")

// ErrNoResponse is returned by Service when a Request did not get a response,
// and no other error happened
var ErrNoResponse = errors.New("no response to request")

// Handler is an interface that objects must implement in order to handle
// a service's requests.
type Handler interface {

	// HandleMessage receives an incoming message, and potentially returns
	// a response message to send back.
	HandleMessage(context.Context, msg.NetMessage) msg.NetMessage
}

// Sender interface for network services.
type Sender interface {
	// SendMessage sends out a given message, without expecting a response.
	SendMessage(ctx context.Context, m msg.NetMessage) error

	// SendRequest sends out a given message, and awaits a response.
	// Set Deadlines or cancellations in the context.Context you pass in.
	SendRequest(ctx context.Context, m msg.NetMessage) (msg.NetMessage, error)
}

// Service is an interface for a net resource with both outgoing (sender) and
// incomig (SetHandler) requests.
type Service interface {
	Sender // can use it to send out msgs
	router.Node // it is a Node in the net topology.

	// SetUplink assigns the Node to send packets out
	SetUplink(router.Node)
	Uplink() router.Node

	// SetHandler assigns the request Handler for this service.
	SetHandler(Handler)
	GetHandler() Handler
}

// Service is a networking component that protocols can use to multiplex
// messages over the same channel, and to issue + handle requests.
type service struct {
	// Handler is the object registered to handle incoming requests.
	Handler     Handler
	HandlerLock sync.RWMutex

	// Requests are all the pending requests on this service.
	Requests     RequestMap
	RequestsLock sync.RWMutex

	// the connection to the outside world
	uplink     router.Node
	uplinkLock sync.RWMutex
	addr router.Address
}

// NewService creates a service object with given type ID and Handler
func NewService(addr router.Address, uplink router.Node, h Handler) Service {
	s := &service{
		Handler:  h,
		Requests: RequestMap{},
		uplink:   uplink,
		addr:     addr,
	}
	return s
}

// sendMessage sends a message out (actual leg work. SendMessage is to export w/o rid)
func (s *service) sendMessage(ctx context.Context, m msg.NetMessage, rid RequestID) error {

	// serialize ServiceMessage wrapper
	data, err := wrapData(m.Data(), rid)
	if err != nil {
		return err
	}

	// log.Debug("Service send message [to = %s]", m.Peer())

	// send message
	m2 := msg.New(m.Peer(), data)

	pkt := msg.Packet{
		Src:
	}

	select {
	case s.Outgoing <- m2:
	case <-ctx.Done():
		return ctx.Err()
	}

	pkt := msg.Packet{
		Src: m.
	}

	return nil
}

// SendMessage sends a message out
func (s *service) SendMessage(ctx context.Context, m msg.NetMessage) error {
	return s.sendMessage(ctx, m, nil)
}

// SendRequest sends a request message out and awaits a response.
func (s *service) SendRequest(ctx context.Context, m msg.NetMessage) (msg.NetMessage, error) {

	// check if we should bail given our contexts
	select {
	default:
	case <-s.Closing():
		return nil, fmt.Errorf("service closed: %s", s.Context().Err())
	case <-ctx.Done():
		return nil, ctx.Err()
	}

	// create a request
	r, err := NewRequest(m.Peer().ID())
	if err != nil {
		return nil, err
	}

	// register Request
	s.RequestsLock.Lock()
	s.Requests[r.Key()] = r
	s.RequestsLock.Unlock()

	// defer deleting this request
	defer func() {
		s.RequestsLock.Lock()
		delete(s.Requests, r.Key())
		s.RequestsLock.Unlock()
	}()

	// check if we should bail after waiting for mutex
	select {
	default:
	case <-s.Closing():
		return nil, fmt.Errorf("service closed: %s", s.Context().Err())
	case <-ctx.Done():
		return nil, ctx.Err()
	}

	// Send message
	s.sendMessage(ctx, m, r.ID)

	// wait for response
	m = nil
	err = nil
	select {
	case m = <-r.Response:
	case <-s.Closed():
		err = fmt.Errorf("service closed: %s", s.Context().Err())
	case <-ctx.Done():
		err = ctx.Err()
	}

	if m == nil {
		return nil, ErrNoResponse
	}

	return m, err
}

// handleIncoming consumes the messages on the s.Incoming channel and
// routes them appropriately (to requests, or handler).
func (s *service) handleIncomingMessages() {
	defer s.Children().Done()

	for {
		select {
		case m, more := <-s.Incoming:
			if !more {
				return
			}
			s.Children().Add(1)
			go s.handleIncomingMessage(m)

		case <-s.Closing():
			return
		}
	}
}

func (s *service) handleIncomingMessage(pkt *msg.Packet) error {

	// check the packet has a valid Context
	ctx := pkt.Context
	if ctx == nil {
		return fmt.Errorf("service got pkt without valid Context")
	}

	// check the source is a peer
	srcPeer, ok := pkt.Src.(peer.Peer)
	if !ok {
		return fmt.Errorf("service got pkt from non-Peer src: %v", pkt.Src)
	}

	// unwrap the incoming message
	data, rid, err := unwrapData(pkt.Data)
	if err != nil {
		return fmt.Errorf("service de-serializing error: %v", err)
	}

	// convert to msg.NetMessage, which the rest of the system expects.
	m2 := msg.New(srcPeer, data)

	// if it's a request (or has no RequestID), handle it
	if rid == nil || rid.IsRequest() {
		handler := s.GetHandler()
		if handler == nil {
			log.Errorf("service dropped msg: %v", m)
			log.Event()
			return nil
			// no handler, drop it.
		}

		// this go routine is developer friendliness to keep their stacks
		// separate (and more readable) from the network goroutine. If
		// problems arise and you'd like to see _the full_ stack of where
		// this message is coming from, just remove the goroutine part.
		response := make(chan msg.NetMessage)
		go func() msg.NetMessage {
			return handler.HandleMessage(ctx, m2)
		}()
		r1 := <-response
		// Note: HandleMessage *must* respect context. We could co-opt it
		// and do a select {} here on the context, BUT that would just drop
		// a packet and free up the goroutine to return to the network. the
		// problem is still there: the Service handler hasn't returned yet.

		// if handler gave us a response, send it out!
		if r1 != nil {
			if err := s.sendMessage(ctx, r1, rid.Response()); err != nil {
				return fmt.Errorf("error sending response message: %v", err)
			}
		}
		return
	}

	// Otherwise, it is a response. handle it.
	if !rid.IsResponse() {
		log.Errorf("RequestID should identify a response here.")
	}

	key := RequestKey(m.Peer().ID(), RequestID(rid))
	s.RequestsLock.RLock()
	r, found := s.Requests[key]
	s.RequestsLock.RUnlock()

	if !found {
		log.Errorf("no request key %v (timeout?)", []byte(key))
		return
	}

	select {
	case r.Response <- m2:
	case <-s.Closing():
	}
}

// Address is the router.Node address
func (s *service) Address() router.Address {
	return s.addr
}

// HandlePacket implements router.Node
// service only receives packets in HandlePacket
func (s *service) HandlePacket(p router.Packet, from router.Node) error {
	pkt, ok := p.(*msg.Packet)
	if !ok {
		return msg.ErrInvalidPayload
	}
}

// SetHandler assigns the request Handler for this service.
func (s *service) SetHandler(h Handler) {
	s.HandlerLock.Lock()
	defer s.HandlerLock.Unlock()
	s.Handler = h
}

// GetHandler returns the request Handler for this service.
func (s *service) GetHandler() Handler {
	s.HandlerLock.RLock()
	defer s.HandlerLock.RUnlock()
	return s.Handler
}
<|endoftext|>


<commit_before>// Copyright 2018 The Terraformer 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 aws

import (
	"context"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"os"
	"regexp"

	"github.com/aws/aws-sdk-go-v2/service/sts"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"

	"github.com/GoogleCloudPlatform/terraformer/terraformutils"
)

type AWSService struct { //nolint
	terraformutils.Service
}

var awsVariable = regexp.MustCompile(`(\${[0-9A-Za-z:]+})`)

var configCache *aws.Config

func (s *AWSService) generateConfig() (aws.Config, error) {
	if configCache != nil {
		return *configCache, nil
	}

	config, e := s.buildBaseConfig()

	if e != nil {
		return config, e
	}
	if s.Verbose {
		config.ClientLogMode = aws.LogRequestWithBody & aws.LogResponseWithBody
	}

	creds, e := config.Credentials.Retrieve(context.TODO())

	if e != nil {
		return config, e
	}

	// terraform cannot ask for MFA token, so we need to pass STS session token, which might contain credentials with MFA requirement
	accessKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
	if accessKey == "" {
		os.Setenv("AWS_ACCESS_KEY_ID", creds.AccessKeyID)
		os.Setenv("AWS_SECRET_ACCESS_KEY", creds.SecretAccessKey)

		if creds.SessionToken != "" {
			os.Setenv("AWS_SESSION_TOKEN", creds.SessionToken)
		}
	}
	configCache = &config
	return config, nil
}

func (s *AWSService) buildBaseConfig() (aws.Config, error) {
	var loadOptions []func(*config.LoadOptions) error
	if s.GetArgs()["profile"].(string) != "" {
		loadOptions = append(loadOptions, config.WithSharedConfigProfile(s.GetArgs()["profile"].(string)))
	}
	if s.GetArgs()["region"].(string) != "" {
		os.Setenv("AWS_REGION", s.GetArgs()["region"].(string))
	}
	loadOptions = append(loadOptions, config.WithAssumeRoleCredentialOptions(func(options *stscreds.AssumeRoleOptions) {
		options.TokenProvider = stscreds.StdinTokenProvider
	}))
	return config.LoadDefaultConfig(context.TODO(), loadOptions...)
}

// for CF interpolation and IAM Policy variables
func (*AWSService) escapeAwsInterpolation(str string) string {
	return awsVariable.ReplaceAllString(str, "$$$1")
}

func (s *AWSService) getAccountNumber(config aws.Config) (*string, error) {
	stsSvc := sts.NewFromConfig(config)
	identity, err := stsSvc.GetCallerIdentity(context.TODO(), &sts.GetCallerIdentityInput{})
	if err != nil {
		return nil, err
	}
	return identity.Account, nil
}
<commit_msg>#549 linting<commit_after>// Copyright 2018 The Terraformer 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 aws

import (
	"context"
	"os"
	"regexp"

	"github.com/aws/aws-sdk-go-v2/service/sts"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"

	"github.com/GoogleCloudPlatform/terraformer/terraformutils"
)

type AWSService struct { //nolint
	terraformutils.Service
}

var awsVariable = regexp.MustCompile(`(\${[0-9A-Za-z:]+})`)

var configCache *aws.Config

func (s *AWSService) generateConfig() (aws.Config, error) {
	if configCache != nil {
		return *configCache, nil
	}

	baseConfig, e := s.buildBaseConfig()

	if e != nil {
		return baseConfig, e
	}
	if s.Verbose {
		baseConfig.ClientLogMode = aws.LogRequestWithBody & aws.LogResponseWithBody
	}

	creds, e := baseConfig.Credentials.Retrieve(context.TODO())

	if e != nil {
		return baseConfig, e
	}

	// terraform cannot ask for MFA token, so we need to pass STS session token, which might contain credentials with MFA requirement
	accessKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
	if accessKey == "" {
		os.Setenv("AWS_ACCESS_KEY_ID", creds.AccessKeyID)
		os.Setenv("AWS_SECRET_ACCESS_KEY", creds.SecretAccessKey)

		if creds.SessionToken != "" {
			os.Setenv("AWS_SESSION_TOKEN", creds.SessionToken)
		}
	}
	configCache = &baseConfig
	return baseConfig, nil
}

func (s *AWSService) buildBaseConfig() (aws.Config, error) {
	var loadOptions []func(*config.LoadOptions) error
	if s.GetArgs()["profile"].(string) != "" {
		loadOptions = append(loadOptions, config.WithSharedConfigProfile(s.GetArgs()["profile"].(string)))
	}
	if s.GetArgs()["region"].(string) != "" {
		os.Setenv("AWS_REGION", s.GetArgs()["region"].(string))
	}
	loadOptions = append(loadOptions, config.WithAssumeRoleCredentialOptions(func(options *stscreds.AssumeRoleOptions) {
		options.TokenProvider = stscreds.StdinTokenProvider
	}))
	return config.LoadDefaultConfig(context.TODO(), loadOptions...)
}

// for CF interpolation and IAM Policy variables
func (*AWSService) escapeAwsInterpolation(str string) string {
	return awsVariable.ReplaceAllString(str, "$$$1")
}

func (s *AWSService) getAccountNumber(config aws.Config) (*string, error) {
	stsSvc := sts.NewFromConfig(config)
	identity, err := stsSvc.GetCallerIdentity(context.TODO(), &sts.GetCallerIdentityInput{})
	if err != nil {
		return nil, err
	}
	return identity.Account, nil
}
<|endoftext|>


<commit_before>// Package patreon implements the OAuth protocol for authenticating users through Patreon.
package patreon

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

	"github.com/markbates/goth"
	"golang.org/x/oauth2"
)

const (
	// AuthorizationURL specifies Patreon's OAuth2 authorization endpoint (see https://tools.ietf.org/html/rfc6749#section-3.1).
	// See Example_refreshToken for examples.
	authorizationURL = "https://www.patreon.com/oauth2/authorize"

	// AccessTokenURL specifies Patreon's OAuth2 token endpoint (see https://tools.ietf.org/html/rfc6749#section-3.2).
	// See Example_refreshToken for examples.
	accessTokenURL = "https://www.patreon.com/api/oauth2/token"

	profileURL = "https://www.patreon.com/api/oauth2/v2/identity"
)

//goland:noinspection GoUnusedConst
const (
	// ScopeIdentity provides read access to data about the user. See the /identity endpoint documentation for details about what data is available.
	ScopeIdentity = "identity"

	// ScopeIdentityEmail provides read access to the user’s email.
	ScopeIdentityEmail = "identity[email]"

	// ScopeIdentityMemberships provides read access to the user’s memberships.
	ScopeIdentityMemberships = "identity.memberships"

	// ScopeCampaigns provides read access to basic campaign data. See the /campaign endpoint documentation for details about what data is available.
	ScopeCampaigns = "campaigns"

	// ScopeCampaignsWebhook provides read, write, update, and delete access to the campaign’s webhooks created by the client.
	ScopeCampaignsWebhook = "w:campaigns.webhook"

	// ScopeCampaignsMembers provides read access to data about a campaign’s members. See the /members endpoint documentation for details about what data is available. Also allows the same information to be sent via webhooks created by your client.
	ScopeCampaignsMembers = "campaigns.members"

	// ScopeCampaignsMembersEmail provides read access to the member’s email. Also allows the same information to be sent via webhooks created by your client.
	ScopeCampaignsMembersEmail = "campaigns.members[email]"

	// ScopeCampaignsMembersAddress provides read access to the member’s address, if an address was collected in the pledge flow. Also allows the same information to be sent via webhooks created by your client.
	ScopeCampaignsMembersAddress = "campaigns.members.address"

	// ScopeCampaignsPosts provides read access to the posts on a campaign.
	ScopeCampaignsPosts = "campaigns.posts"
)

// New creates a new Patreon provider and sets up important connection details.
// You should always call `patreon.New` to get a new provider.  Never try to
// create one manually.
func New(clientKey, secret, callbackURL string, scopes ...string) *Provider {
	p := &Provider{
		ClientKey:    clientKey,
		Secret:       secret,
		CallbackURL:  callbackURL,
		providerName: "patreon",
	}
	p.config = newConfig(p, scopes)
	return p
}

// Provider is the implementation of `goth.Provider` for accessing Patreon.
type Provider struct {
	ClientKey    string
	Secret       string
	CallbackURL  string
	HTTPClient   *http.Client
	config       *oauth2.Config
	providerName string
}

// Name gets the name used to retrieve this provider.
func (p *Provider) Name() string {
	return p.providerName
}

// SetName is to update the name of the provider (needed in case of multiple providers of 1 type)
func (p *Provider) SetName(name string) {
	p.providerName = name
}

func (p *Provider) Client() *http.Client {
	return goth.HTTPClientWithFallBack(p.HTTPClient)
}

// Debug is a no-op for the Patreon package.
func (p *Provider) Debug(debug bool) {}

// BeginAuth asks Patreon for an authentication end-point.
func (p *Provider) BeginAuth(state string) (goth.Session, error) {
	url := p.config.AuthCodeURL(state)
	session := &Session{
		AuthURL: url,
	}
	return session, nil
}

// FetchUser will go to Patreon and access basic information about the user.
func (p *Provider) FetchUser(session goth.Session) (goth.User, error) {
	s := session.(*Session)
	user := goth.User{
		AccessToken:  s.AccessToken,
		Provider:     p.Name(),
		RefreshToken: s.RefreshToken,
		ExpiresAt:    s.ExpiresAt,
	}

	if user.AccessToken == "" {
		// data is not yet retrieved since accessToken is still empty
		return user, fmt.Errorf("%s cannot get user information without accessToken", p.providerName)
	}

	req, err := http.NewRequest("GET", profileURL, nil)
	if err != nil {
		return user, err
	}
	req.Header.Set("Authorization", "Bearer "+s.AccessToken)
	resp, err := p.Client().Do(req)
	if err != nil {
		return user, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return user, fmt.Errorf("%s responded with a %d trying to fetch user information", p.providerName, resp.StatusCode)
	}

	u := User{}
	err = json.NewDecoder(resp.Body).Decode(&u)
	if err != nil {
		return user, err
	}

	user.Name = u.Data.Attributes.FullName
	user.Email = u.Data.Attributes.Email
	user.UserID = u.Data.ID

	return user, err
}

func newConfig(provider *Provider, authURL, tokenURL string, scopes []string) *oauth2.Config {
	c := &oauth2.Config{
		ClientID:     provider.ClientKey,
		ClientSecret: provider.Secret,
		RedirectURL:  provider.CallbackURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  authURL,
			TokenURL: tokenURL,
		},
		Scopes: []string{},
	}

	if len(scopes) > 0 {
		for _, scope := range scopes {
			c.Scopes = append(c.Scopes, scope)
		}
	}
	return c
}

// RefreshTokenAvailable refresh token is provided by auth provider or not
func (p *Provider) RefreshTokenAvailable() bool {
	return true
}

// RefreshToken get new access token based on the refresh token
func (p *Provider) RefreshToken(refreshToken string) (*oauth2.Token, error) {
	token := &oauth2.Token{RefreshToken: refreshToken}
	ts := p.config.TokenSource(goth.ContextForClient(p.Client()), token)
	newToken, err := ts.Token()
	if err != nil {
		return nil, err
	}
	return newToken, err
}

type User struct {
	Data struct {
		Attributes struct {
			Email    string `json:"email"`
			FullName string `json:"full_name"`
		} `json:"attributes"`
		ID string `json:"id"`
	} `json:"data"`
}
<commit_msg>NewCustomisedURL<commit_after>// Package patreon implements the OAuth protocol for authenticating users through Patreon.
package patreon

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

	"github.com/markbates/goth"
	"golang.org/x/oauth2"
)

const (
	// AuthorizationURL specifies Patreon's OAuth2 authorization endpoint (see https://tools.ietf.org/html/rfc6749#section-3.1).
	// See Example_refreshToken for examples.
	authorizationURL = "https://www.patreon.com/oauth2/authorize"

	// AccessTokenURL specifies Patreon's OAuth2 token endpoint (see https://tools.ietf.org/html/rfc6749#section-3.2).
	// See Example_refreshToken for examples.
	accessTokenURL = "https://www.patreon.com/api/oauth2/token"

	profileURL = "https://www.patreon.com/api/oauth2/v2/identity"
)

//goland:noinspection GoUnusedConst
const (
	// ScopeIdentity provides read access to data about the user. See the /identity endpoint documentation for details about what data is available.
	ScopeIdentity = "identity"

	// ScopeIdentityEmail provides read access to the user’s email.
	ScopeIdentityEmail = "identity[email]"

	// ScopeIdentityMemberships provides read access to the user’s memberships.
	ScopeIdentityMemberships = "identity.memberships"

	// ScopeCampaigns provides read access to basic campaign data. See the /campaign endpoint documentation for details about what data is available.
	ScopeCampaigns = "campaigns"

	// ScopeCampaignsWebhook provides read, write, update, and delete access to the campaign’s webhooks created by the client.
	ScopeCampaignsWebhook = "w:campaigns.webhook"

	// ScopeCampaignsMembers provides read access to data about a campaign’s members. See the /members endpoint documentation for details about what data is available. Also allows the same information to be sent via webhooks created by your client.
	ScopeCampaignsMembers = "campaigns.members"

	// ScopeCampaignsMembersEmail provides read access to the member’s email. Also allows the same information to be sent via webhooks created by your client.
	ScopeCampaignsMembersEmail = "campaigns.members[email]"

	// ScopeCampaignsMembersAddress provides read access to the member’s address, if an address was collected in the pledge flow. Also allows the same information to be sent via webhooks created by your client.
	ScopeCampaignsMembersAddress = "campaigns.members.address"

	// ScopeCampaignsPosts provides read access to the posts on a campaign.
	ScopeCampaignsPosts = "campaigns.posts"
)

// New creates a new Patreon provider and sets up important connection details.
// You should always call `patreon.New` to get a new provider.  Never try to
// create one manually.
func New(clientKey, secret, callbackURL string, scopes ...string) *Provider {
	return NewCustomisedURL(clientKey, secret, callbackURL, authorizationURL, accessTokenURL, profileURL, scopes...)
}

// NewCustomisedURL is similar to New(...) but can be used to set custom URLs to connect to
func NewCustomisedURL(clientKey, secret, callbackURL, authURL, tokenURL, profileURL string, scopes ...string) *Provider {
	p := &Provider{
		ClientKey:    clientKey,
		Secret:       secret,
		CallbackURL:  callbackURL,
		providerName: "patreon",
		profileURL:   profileURL,
	}
	p.config = newConfig(p, authURL, tokenURL, scopes)
	return p
}

// Provider is the implementation of `goth.Provider` for accessing Patreon.
type Provider struct {
	ClientKey    string
	Secret       string
	CallbackURL  string
	HTTPClient   *http.Client
	config       *oauth2.Config
	providerName string
}

// Name gets the name used to retrieve this provider.
func (p *Provider) Name() string {
	return p.providerName
}

// SetName is to update the name of the provider (needed in case of multiple providers of 1 type)
func (p *Provider) SetName(name string) {
	p.providerName = name
}

func (p *Provider) Client() *http.Client {
	return goth.HTTPClientWithFallBack(p.HTTPClient)
}

// Debug is a no-op for the Patreon package.
func (p *Provider) Debug(debug bool) {}

// BeginAuth asks Patreon for an authentication end-point.
func (p *Provider) BeginAuth(state string) (goth.Session, error) {
	url := p.config.AuthCodeURL(state)
	session := &Session{
		AuthURL: url,
	}
	return session, nil
}

// FetchUser will go to Patreon and access basic information about the user.
func (p *Provider) FetchUser(session goth.Session) (goth.User, error) {
	s := session.(*Session)
	user := goth.User{
		AccessToken:  s.AccessToken,
		Provider:     p.Name(),
		RefreshToken: s.RefreshToken,
		ExpiresAt:    s.ExpiresAt,
	}

	if user.AccessToken == "" {
		// data is not yet retrieved since accessToken is still empty
		return user, fmt.Errorf("%s cannot get user information without accessToken", p.providerName)
	}

	req, err := http.NewRequest("GET", profileURL, nil)
	if err != nil {
		return user, err
	}
	req.Header.Set("Authorization", "Bearer "+s.AccessToken)
	resp, err := p.Client().Do(req)
	if err != nil {
		return user, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return user, fmt.Errorf("%s responded with a %d trying to fetch user information", p.providerName, resp.StatusCode)
	}

	u := User{}
	err = json.NewDecoder(resp.Body).Decode(&u)
	if err != nil {
		return user, err
	}

	user.Name = u.Data.Attributes.FullName
	user.Email = u.Data.Attributes.Email
	user.UserID = u.Data.ID

	return user, err
}

func newConfig(provider *Provider, authURL, tokenURL string, scopes []string) *oauth2.Config {
	c := &oauth2.Config{
		ClientID:     provider.ClientKey,
		ClientSecret: provider.Secret,
		RedirectURL:  provider.CallbackURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  authURL,
			TokenURL: tokenURL,
		},
		Scopes: []string{},
	}

	if len(scopes) > 0 {
		for _, scope := range scopes {
			c.Scopes = append(c.Scopes, scope)
		}
	}
	return c
}

// RefreshTokenAvailable refresh token is provided by auth provider or not
func (p *Provider) RefreshTokenAvailable() bool {
	return true
}

// RefreshToken get new access token based on the refresh token
func (p *Provider) RefreshToken(refreshToken string) (*oauth2.Token, error) {
	token := &oauth2.Token{RefreshToken: refreshToken}
	ts := p.config.TokenSource(goth.ContextForClient(p.Client()), token)
	newToken, err := ts.Token()
	if err != nil {
		return nil, err
	}
	return newToken, err
}

type User struct {
	Data struct {
		Attributes struct {
			Email    string `json:"email"`
			FullName string `json:"full_name"`
		} `json:"attributes"`
		ID string `json:"id"`
	} `json:"data"`
}
<|endoftext|>


<commit_before>package chef

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"reflect"
	"strings"
	"testing"
)

var testRequiredHeaders []string

var ConfigFilePath = "./TEST_CONFIG.json"

func init() {
	testRequiredHeaders = []string{
		"Accept",
		"X-Ops-Timestamp",
		"X-Ops-Userid",
		"X-Ops-Sign",
		"X-Ops-Content-Hash",
		"X-Ops-Authorization-1",
	}
}

func testConnectionWrapper() *Chef {
	t := new(testing.T)
	chef, err := Connect()
	if err != nil {
		t.Error(err)
	}
	chef.SSLNoVerify = true

	return chef
}

type testConfigFile struct {
	RequiredCookbook struct {
		Name    string `json:"name"`
		Version string `json:"version"`
	} `json:"required_cookbook"`
	RequiredNode struct {
		Name string `json:"name"`
	} `json:"required_node"`
	RequiredRecipe struct {
		Name string `json:"name"`
	} `json:"required_recipe"`
	RequiredRole struct {
		Name string `json:"name"`
	} `json:"required_role"`
	RequiredClient struct {
		Name string `json:"name"`
	} `json:"required_client"`
	RequiredEnvironment struct {
		Name string `json:"name"`
	} `json:"required_environment"`
	RequiredUser struct {
		Name string `json:"name"`
	} `json:"required_user"`
	RequiredData struct {
		Name string `json:"name"`
	} `json:"required_data"`
	RequiredPrincipal struct {
		Name string `json:"name"`
	} `json:"required_principal"`
	SearchData struct {
		Index string `json:"index"`
		Query string `json:"query"`
	} `json:"search_data"`
	TestCredentials struct {
		Host    string `json:"host"`
		Port    string `json:"port"`
		Version string `json:"version"`
		UserId  string `json:"user_name"`
		Key     string `json:"key"`
	} `json:"test_credentials"`
	KeyPath   string `json:"key_path"`
	KeyString string `json:"key_string"`
}

func testConfig() *testConfigFile {
	t := new(testing.T)
	file, err := ioutil.ReadFile(ConfigFilePath)
	if err != nil {
		t.Error(err)
	}

	var config *testConfigFile
	json.Unmarshal(file, &config)
	if config == nil {
		t.Error("Config is nil")
	}
	return config
}

func TestReadConfig(t *testing.T) {
	_ = testConfig()
}

func TestApiRequest(t *testing.T) {
	chef := testConnectionWrapper()
	method := "GET"
	endpoint := "cookbooks"
	requestURL := fmt.Sprintf("%s/%s", chef.Url, endpoint)
	req, err := http.NewRequest(method, requestURL, nil)
	if err != nil {
		t.Error(err)
	}
	chef.apiRequest(req, method, fmt.Sprintf("/%s", endpoint), "")

	for _, value := range testRequiredHeaders {
		if req.Header.Get(value) == "" {
			t.Error("Couldn't find header:", value)
		}
	}
}

func TestResponseBody(t *testing.T) {
	etsy, err := http.Get("https://www.etsy.com/")
	if err != nil {
		t.Error(err)
	}

	bytes, err := responseBody(etsy)
	if err != nil {
		t.Error(err)
	}

	etsyString := "Is code your craft? http://www.etsy.com/careers"
	if !strings.Contains(string(bytes), etsyString) {
		t.Error("Response body didn't return valid string")
	}
}

func TestGet(t *testing.T) {
	c := testConnectionWrapper()
	resp, err := c.Get("cookbooks")
	if err != nil {
		t.Error(err)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Error(err)
	}

	cookbooks := map[string]interface{}{}
	json.Unmarshal(body, &cookbooks)
	found := false
	config := testConfig()
	cookbook := config.RequiredCookbook.Name
	for name := range cookbooks {
		if name == cookbook {
			found = true
			break
		}
	}
	if !found {
		t.Error("Required cookbook not found")
	}
}

func TestPost(t *testing.T) {
  c := testConnectionWrapper()
  config := testConfig()
  cookbook := config.RequiredCookbook.Name
  run_list := strings.NewReader(fmt.Sprintf(`{ "run_list": [ "%s" ] }`, cookbook))
  resp, err := c.Post("/environments/_default/cookbook_versions", "application/json", run_list)
  if err != nil {
    t.Error(err)
  }

  body, err := ioutil.ReadAll(resp.Body)
  if err != nil {
    t.Error(err)
  }

  // This could or should be better. Be good to have another
  // test for unsolvable run_list
  cookbooks := map[string]interface{}{}
  json.Unmarshal(body, &cookbooks)
  found := false
  for name := range cookbooks {
    if name == cookbook {
      found = true
      break
    }
  }
  if !found {
    t.Error("Cookbook not solved")
  }
}

func TestConnect(t *testing.T) {
	_, err := Connect()
	if err != nil {
		t.Error(err)
	}
}

func TestApiRequestHeaders(t *testing.T) {
	chef := testConnectionWrapper()
	headers := chef.apiRequestHeaders("GET", "/cookbooks", "")
	count := 0
	for _, requiredHeader := range testRequiredHeaders {
		for header := range headers {
			if strings.ToLower(requiredHeader) == strings.ToLower(header) {
				count += 1
				break
			}
		}
	}
	if count != len(testRequiredHeaders) {
		t.Error("apiRequestHeaders didn't return all of testRequiredHeaders")
	}
}

func TestGenerateRequestAuthorization(t *testing.T) {
	chef := testConnectionWrapper()
	auth := chef.generateRequestAuthorization("GET", "/cookbooks", "", "2013-10-27T20:45:25Z")
	if len(auth[0]) != 60 {
		t.Error("Incorrect request authorization string")
	}
}

func TestPrivateEncrypt(t *testing.T) {
	chef := testConnectionWrapper()
	enc, err := chef.privateEncrypt([]byte("encrypt_this"))
	if err != nil {
		t.Error(err)
	}
	if len(enc) != 256 {
		t.Error("Wrong size of encrypted data")
	}
}

func TestBase64BlockEncode(t *testing.T) {
	toEncode := []byte("abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz")
	results := base64BlockEncode(toEncode)
	expected := []string{"YWJjZGVmZ2hpamtsbW5vcHFyc3R1dnd4eXphYmNkZWZnaGlqa2xtbm9wcXJz", "dHV2d3h5emFiY2RlZmdoaWprbG1ub3BxcnN0dXZ3eHl6"}
	if !reflect.DeepEqual(results, expected) {
		t.Error("Results not matching")
	}
}

func TestHashAndBase64(t *testing.T) {
	if len(hashAndBase64("hash_this")) != 28 {
		t.Error("Wrong length for hashAndBase64")
	}
}

func TestDo(t *testing.T) {
	chef := testConnectionWrapper()
	req, err := http.NewRequest("GET", "https://www.etsy.com/", nil)
	if err != nil {
		t.Error(err)
	}
	_, err = chef.Do(req)
	if err != nil {
		t.Error(err)
	}

}

func TestGenerateRequest(t *testing.T) {
	chef := testConnectionWrapper()
	_, err := chef.generateRequest("GET", "cookbooks", nil)
	if err != nil {
		t.Error(err)
	}
}

func TestKeyFromString(t *testing.T) {
	config := testConfig()
	_, err := keyFromString([]byte(config.KeyString))
	if err != nil {
		t.Error(err)
	}
}

func TestKeyFromFile(t *testing.T) {
	config := testConfig()
	_, err := keyFromFile(config.KeyPath)
	if err != nil {
		t.Error(err)
	}
}

func TestConnectCredentials(t *testing.T) {
	config := testConfig()
	host := config.TestCredentials.Host
	port := config.TestCredentials.Port
	version := config.TestCredentials.Version
	userid := config.TestCredentials.UserId
	key := config.TestCredentials.Key
	_, err := ConnectCredentials(host, port, version, userid, key)
	if err != nil {
		t.Error(err)
	}
}

func TestSplitWhitespace(t *testing.T) {
	str := "c   h   e   f"
	if !reflect.DeepEqual(splitWhitespace(str), []string{"c", "h", "e", "f"}) {
		t.Error("splitWhitespace slices not equal")
	}
}

func TestFilterQuotes(t *testing.T) {
	known := map[string]string{
		`'this`: "this",
		`this'`: "this",
		`"this`: "this",
		`this"`: "this",
	}

	for bad, good := range known {
		if filterQuotes(bad) != good {
			t.Error("filterQuotes didn't produce an expected string")
		}
	}
}
<commit_msg>Fail more gracefully when you can't connect<commit_after>package chef

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"reflect"
	"strings"
	"testing"
)

var testRequiredHeaders []string

var ConfigFilePath = "./TEST_CONFIG.json"

func init() {
	testRequiredHeaders = []string{
		"Accept",
		"X-Ops-Timestamp",
		"X-Ops-Userid",
		"X-Ops-Sign",
		"X-Ops-Content-Hash",
		"X-Ops-Authorization-1",
	}
}

func testConnectionWrapper() *Chef {
	t := new(testing.T)
	chef, err := Connect()
	if err != nil {
		t.Fatal(err)
	}
	chef.SSLNoVerify = true

	return chef
}

type testConfigFile struct {
	RequiredCookbook struct {
		Name    string `json:"name"`
		Version string `json:"version"`
	} `json:"required_cookbook"`
	RequiredNode struct {
		Name string `json:"name"`
	} `json:"required_node"`
	RequiredRecipe struct {
		Name string `json:"name"`
	} `json:"required_recipe"`
	RequiredRole struct {
		Name string `json:"name"`
	} `json:"required_role"`
	RequiredClient struct {
		Name string `json:"name"`
	} `json:"required_client"`
	RequiredEnvironment struct {
		Name string `json:"name"`
	} `json:"required_environment"`
	RequiredUser struct {
		Name string `json:"name"`
	} `json:"required_user"`
	RequiredData struct {
		Name string `json:"name"`
	} `json:"required_data"`
	RequiredPrincipal struct {
		Name string `json:"name"`
	} `json:"required_principal"`
	SearchData struct {
		Index string `json:"index"`
		Query string `json:"query"`
	} `json:"search_data"`
	TestCredentials struct {
		Host    string `json:"host"`
		Port    string `json:"port"`
		Version string `json:"version"`
		UserId  string `json:"user_name"`
		Key     string `json:"key"`
	} `json:"test_credentials"`
	KeyPath   string `json:"key_path"`
	KeyString string `json:"key_string"`
}

func testConfig() *testConfigFile {
	t := new(testing.T)
	file, err := ioutil.ReadFile(ConfigFilePath)
	if err != nil {
		t.Error(err)
	}

	var config *testConfigFile
	json.Unmarshal(file, &config)
	if config == nil {
		t.Error("Config is nil")
	}
	return config
}

func TestReadConfig(t *testing.T) {
	_ = testConfig()
}

func TestApiRequest(t *testing.T) {
	chef := testConnectionWrapper()
	method := "GET"
	endpoint := "cookbooks"
	requestURL := fmt.Sprintf("%s/%s", chef.Url, endpoint)
	req, err := http.NewRequest(method, requestURL, nil)
	if err != nil {
		t.Error(err)
	}
	chef.apiRequest(req, method, fmt.Sprintf("/%s", endpoint), "")

	for _, value := range testRequiredHeaders {
		if req.Header.Get(value) == "" {
			t.Error("Couldn't find header:", value)
		}
	}
}

func TestResponseBody(t *testing.T) {
	etsy, err := http.Get("https://www.etsy.com/")
	if err != nil {
		t.Error(err)
	}

	bytes, err := responseBody(etsy)
	if err != nil {
		t.Error(err)
	}

	etsyString := "Is code your craft? http://www.etsy.com/careers"
	if !strings.Contains(string(bytes), etsyString) {
		t.Error("Response body didn't return valid string")
	}
}

func TestGet(t *testing.T) {
	c := testConnectionWrapper()
	resp, err := c.Get("cookbooks")
	if err != nil {
		t.Error(err)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Error(err)
	}

	cookbooks := map[string]interface{}{}
	json.Unmarshal(body, &cookbooks)
	found := false
	config := testConfig()
	cookbook := config.RequiredCookbook.Name
	for name := range cookbooks {
		if name == cookbook {
			found = true
			break
		}
	}
	if !found {
		t.Error("Required cookbook not found")
	}
}

func TestPost(t *testing.T) {
  c := testConnectionWrapper()
  config := testConfig()
  cookbook := config.RequiredCookbook.Name
  run_list := strings.NewReader(fmt.Sprintf(`{ "run_list": [ "%s" ] }`, cookbook))
  resp, err := c.Post("/environments/_default/cookbook_versions", "application/json", run_list)
  if err != nil {
    t.Error(err)
  }

  body, err := ioutil.ReadAll(resp.Body)
  if err != nil {
    t.Error(err)
  }

  // This could or should be better. Be good to have another
  // test for unsolvable run_list
  cookbooks := map[string]interface{}{}
  json.Unmarshal(body, &cookbooks)
  found := false
  for name := range cookbooks {
    if name == cookbook {
      found = true
      break
    }
  }
  if !found {
    t.Error("Cookbook not solved")
  }
}

func TestConnect(t *testing.T) {
	_, err := Connect()
	if err != nil {
		t.Error(err)
	}
}

func TestApiRequestHeaders(t *testing.T) {
	chef := testConnectionWrapper()
	headers := chef.apiRequestHeaders("GET", "/cookbooks", "")
	count := 0
	for _, requiredHeader := range testRequiredHeaders {
		for header := range headers {
			if strings.ToLower(requiredHeader) == strings.ToLower(header) {
				count += 1
				break
			}
		}
	}
	if count != len(testRequiredHeaders) {
		t.Error("apiRequestHeaders didn't return all of testRequiredHeaders")
	}
}

func TestGenerateRequestAuthorization(t *testing.T) {
	chef := testConnectionWrapper()
	auth := chef.generateRequestAuthorization("GET", "/cookbooks", "", "2013-10-27T20:45:25Z")
	if len(auth[0]) != 60 {
		t.Error("Incorrect request authorization string")
	}
}

func TestPrivateEncrypt(t *testing.T) {
	chef := testConnectionWrapper()
	enc, err := chef.privateEncrypt([]byte("encrypt_this"))
	if err != nil {
		t.Error(err)
	}
	if len(enc) != 256 {
		t.Error("Wrong size of encrypted data")
	}
}

func TestBase64BlockEncode(t *testing.T) {
	toEncode := []byte("abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz")
	results := base64BlockEncode(toEncode)
	expected := []string{"YWJjZGVmZ2hpamtsbW5vcHFyc3R1dnd4eXphYmNkZWZnaGlqa2xtbm9wcXJz", "dHV2d3h5emFiY2RlZmdoaWprbG1ub3BxcnN0dXZ3eHl6"}
	if !reflect.DeepEqual(results, expected) {
		t.Error("Results not matching")
	}
}

func TestHashAndBase64(t *testing.T) {
	if len(hashAndBase64("hash_this")) != 28 {
		t.Error("Wrong length for hashAndBase64")
	}
}

func TestDo(t *testing.T) {
	chef := testConnectionWrapper()
	req, err := http.NewRequest("GET", "https://www.etsy.com/", nil)
	if err != nil {
		t.Error(err)
	}
	_, err = chef.Do(req)
	if err != nil {
		t.Error(err)
	}

}

func TestGenerateRequest(t *testing.T) {
	chef := testConnectionWrapper()
	_, err := chef.generateRequest("GET", "cookbooks", nil)
	if err != nil {
		t.Error(err)
	}
}

func TestKeyFromString(t *testing.T) {
	config := testConfig()
	_, err := keyFromString([]byte(config.KeyString))
	if err != nil {
		t.Error(err)
	}
}

func TestKeyFromFile(t *testing.T) {
	config := testConfig()
	_, err := keyFromFile(config.KeyPath)
	if err != nil {
		t.Error(err)
	}
}

func TestConnectCredentials(t *testing.T) {
	config := testConfig()
	host := config.TestCredentials.Host
	port := config.TestCredentials.Port
	version := config.TestCredentials.Version
	userid := config.TestCredentials.UserId
	key := config.TestCredentials.Key
	_, err := ConnectCredentials(host, port, version, userid, key)
	if err != nil {
		t.Error(err)
	}
}

func TestSplitWhitespace(t *testing.T) {
	str := "c   h   e   f"
	if !reflect.DeepEqual(splitWhitespace(str), []string{"c", "h", "e", "f"}) {
		t.Error("splitWhitespace slices not equal")
	}
}

func TestFilterQuotes(t *testing.T) {
	known := map[string]string{
		`'this`: "this",
		`this'`: "this",
		`"this`: "this",
		`this"`: "this",
	}

	for bad, good := range known {
		if filterQuotes(bad) != good {
			t.Error("filterQuotes didn't produce an expected string")
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 James Blizzard. All rights reserved.
//
// API helper functions

package apitools

import (
    "net/http"
    "encoding/json"
    "github.com/gorilla/mux"
    "github.com/jimblizz/logger"
    "strconv"
)

type ApiResponse struct {
    ResponseCode int    `json:"code"`
    Data interface{}    `json:"data"`
}

type Api struct {
    Routes map[string]string `json:"routes"`
    Log *logger.Logger
}

type File struct {
    Name		string
    ContentType         string
    BodyLength 	        float64
    Body       	        []byte
}

func New(log *logger.Logger) *Api {
    var a Api
    a.Log = log
    a.Routes = make(map[string]string)
    return &a
}

func (a Api) RegisterRoute (r *mux.Router, f func(http.ResponseWriter, *http.Request), method string, path string, description string) {
    r.HandleFunc(path, f).Methods(method)
    a.Routes[path] = description
}

func (a Api) SendOk(w http.ResponseWriter, results interface{}) () {
    w.Header().Set("Content-Type", "application/json; charset=UTF-8")
    w.WriteHeader(http.StatusOK)

    err := json.NewEncoder(w).Encode(ApiResponse{200, results})
    if err != nil {
        a.Log.Error("Error sending response", err)
    }
}

func (a Api) SendBad(w http.ResponseWriter, message string) {
    w.Header().Set("Content-Type", "application/json; charset=UTF-8")
    w.WriteHeader(http.StatusBadRequest)

    err := json.NewEncoder(w).Encode(ApiResponse{400, message})
    if err != nil {
        a.Log.Error("Error sending response", err)
    }
}

func (a Api) SendAccessDenied(w http.ResponseWriter, message string) {
    w.Header().Set("Content-Type", "application/json; charset=UTF-8")
    w.WriteHeader(http.StatusForbidden)

    err := json.NewEncoder(w).Encode(ApiResponse{403, message})
    if err != nil {
        a.Log.Error("Error sending response", err)
    }
}

func (a Api) SendNotFound(w http.ResponseWriter, message string) {
    w.Header().Set("Content-Type", "application/json; charset=UTF-8")
    w.WriteHeader(http.StatusNotFound)

    err := json.NewEncoder(w).Encode(ApiResponse{404, message})
    if err != nil {
        a.Log.Error("Error sending response", err)
    }
}

func (a Api) SendError(w http.ResponseWriter, message string) {
    w.Header().Set("Content-Type", "application/json; charset=UTF-8")
    w.WriteHeader(http.StatusInternalServerError)

    err := json.NewEncoder(w).Encode(ApiResponse{500, message})
    if err != nil {
        a.Log.Error("Error sending response", err)
    }
}

func (a Api) SendDownload(w http.ResponseWriter, file File, disposition string) () {
    if disposition == "" {
        disposition = "attachment"
    }

    w.Header().Set("Content-Disposition", disposition + "; filename=\"" + file.Name + "\"")
    w.Header().Set("Content-Type", file.ContentType)
    w.Header().Set("Content-Length", strconv.Itoa(len(file.Body)))
    w.Header().Set("Cache-Control", "no-cache")

    w.Write(file.Body)
}
<commit_msg>Improve SendDownload support<commit_after>// Copyright 2016 James Blizzard. All rights reserved.
//
// API helper functions

package apitools

import (
    "net/http"
    "encoding/json"
    "github.com/gorilla/mux"
    "github.com/jimblizz/logger"
    "strconv"
)

type ApiResponse struct {
    ResponseCode int    `json:"code"`
    Data interface{}    `json:"data"`
}

type Api struct {
    Routes map[string]string `json:"routes"`
    Log *logger.Logger
}

type File struct {
    Name		string
    ContentType         string
    BodyLength 	        float64
    Body       	        []byte
}

func New(log *logger.Logger) *Api {
    var a Api
    a.Log = log
    a.Routes = make(map[string]string)
    return &a
}

func (a Api) RegisterRoute (r *mux.Router, f func(http.ResponseWriter, *http.Request), method string, path string, description string) {
    r.HandleFunc(path, f).Methods(method)
    a.Routes[path] = description
}

func (a Api) SendOk(w http.ResponseWriter, results interface{}) () {
    w.Header().Set("Content-Type", "application/json; charset=UTF-8")
    w.WriteHeader(http.StatusOK)

    err := json.NewEncoder(w).Encode(ApiResponse{200, results})
    if err != nil {
        a.Log.Error("Error sending response", err)
    }
}

func (a Api) SendBad(w http.ResponseWriter, message string) {
    w.Header().Set("Content-Type", "application/json; charset=UTF-8")
    w.WriteHeader(http.StatusBadRequest)

    err := json.NewEncoder(w).Encode(ApiResponse{400, message})
    if err != nil {
        a.Log.Error("Error sending response", err)
    }
}

func (a Api) SendAccessDenied(w http.ResponseWriter, message string) {
    w.Header().Set("Content-Type", "application/json; charset=UTF-8")
    w.WriteHeader(http.StatusForbidden)

    err := json.NewEncoder(w).Encode(ApiResponse{403, message})
    if err != nil {
        a.Log.Error("Error sending response", err)
    }
}

func (a Api) SendNotFound(w http.ResponseWriter, message string) {
    w.Header().Set("Content-Type", "application/json; charset=UTF-8")
    w.WriteHeader(http.StatusNotFound)

    err := json.NewEncoder(w).Encode(ApiResponse{404, message})
    if err != nil {
        a.Log.Error("Error sending response", err)
    }
}

func (a Api) SendError(w http.ResponseWriter, message string) {
    w.Header().Set("Content-Type", "application/json; charset=UTF-8")
    w.WriteHeader(http.StatusInternalServerError)

    err := json.NewEncoder(w).Encode(ApiResponse{500, message})
    if err != nil {
        a.Log.Error("Error sending response", err)
    }
}

func (a Api) SendDownload(w http.ResponseWriter, file File, disposition string) () {
    if disposition == "" {
        disposition = "attachment"
    }

    w.Header().Set("Content-Disposition", disposition + "; filename=\"" + file.Name + "\"")
    w.Header().Set("Cache-Control", "no-cache")

    reader := bytes.NewReader(file.Body)
    http.ServeContent(w, r, file.Name, time.Now(), reader)
}
<|endoftext|>


<commit_before>package python

import (
	"crypto/sha1"
	"encoding/hex"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
)

var dockerEnv = os.Getenv("IN_DOCKER_CONTAINER")

func runCmdLogError(cmd *exec.Cmd) {
	log.Printf("Running %v", cmd.Args)
	err := runCmdStderr(cmd)
	if err != nil {
		log.Printf("Error running `%s`: %s", strings.Join(cmd.Args, " "), err)
	}
}

func runCmdStderr(cmd *exec.Cmd) error {
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stderr
	return cmd.Run()
}

// getProgramPath returns path to toolchain (assuming that exe file is located in .bin, path is <path-to-.bin>/..
// In `docker` mode, it will return empty string because there is no virtualenv.
func getProgramPath() (string, error) {
	if dockerEnv == "" {
		path := filepath.Join(filepath.Dir(os.Args[0]), "..")
		return filepath.Abs(path)
	}
	return "", nil
}

// getVENVBinPath returns toolchains Python virtualenv path. If toolchain is ran in
// `docker` mode, it will return empty string because there is no virtualenv.
func getVENVBinPath() (string, error) {
	if dockerEnv == "" {
		path, err := getProgramPath()
		if (err != nil) {
			return ``, err
		}
		return filepath.Abs(filepath.Join(path, ".env", getEnvBinDir()))
	}
	return "", nil
}

// Returns binaries directory of virtualenv which may be different on Windows and Unix
func getEnvBinDir() string {
	if runtime.GOOS == "windows" {
		return "Scripts"
	} else {
		return "bin"
	}
}

func getHash(text string) string {
	hasher := sha1.New()
	hasher.Write([]byte(text))
	return hex.EncodeToString(hasher.Sum(nil))[:8]
}
<commit_msg>fix docker mode<commit_after>package python

import (
	"crypto/sha1"
	"encoding/hex"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
)

var dockerEnv = os.Getenv("IN_DOCKER_CONTAINER")

func runCmdLogError(cmd *exec.Cmd) {
	log.Printf("Running %v", cmd.Args)
	err := runCmdStderr(cmd)
	if err != nil {
		log.Printf("Error running `%s`: %s", strings.Join(cmd.Args, " "), err)
	}
}

func runCmdStderr(cmd *exec.Cmd) error {
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stderr
	return cmd.Run()
}

// getProgramPath returns path to toolchain (assuming that exe file is located in .bin, path is <path-to-.bin>/..
// In `docker` mode, it will return empty string because there is no virtualenv.
func getProgramPath() (string, error) {
	if dockerEnv == "" {
		path := filepath.Join(filepath.Dir(os.Args[0]), "..")
		return filepath.Abs(path)
	}
	return "", nil
}

// getVENVBinPath returns toolchains Python virtualenv path. If toolchain is ran in
// `docker` mode, it will return empty string because there is no virtualenv.
func getVENVBinPath() (string, error) {
	if dockerEnv == "" {
		path, err := getProgramPath()
		if err != nil {
			return "", err
		}
		return filepath.Abs(filepath.Join(path, ".env", getEnvBinDir()))
	}
	return "/venv/bin", nil
}

// Returns binaries directory of virtualenv which may be different on Windows and Unix
func getEnvBinDir() string {
	if runtime.GOOS == "windows" {
		return "Scripts"
	} else {
		return "bin"
	}
}

func getHash(text string) string {
	hasher := sha1.New()
	hasher.Write([]byte(text))
	return hex.EncodeToString(hasher.Sum(nil))[:8]
}
<|endoftext|>


<commit_before>package search

import (
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"

	"time"

	"github.com/NyaaPantsu/nyaa/config"
	"github.com/NyaaPantsu/nyaa/models"
	"github.com/NyaaPantsu/nyaa/models/torrents"
	"github.com/NyaaPantsu/nyaa/utils/cache"
	"github.com/NyaaPantsu/nyaa/utils/log"
	"github.com/NyaaPantsu/nyaa/utils/search/structs"
	"github.com/gin-gonic/gin"
)

var searchOperator string
var useTSQuery bool

// Configure : initialize search
func Configure(conf *config.SearchConfig) (err error) {
	useTSQuery = false
	// Postgres needs ILIKE for case-insensitivity
	if models.ORM.Dialect().GetName() == "postgres" {
		searchOperator = "ILIKE ?"
		//useTSQuery = true
		// !!DISABLED!! because this makes search a lot stricter
		// (only matches at word borders)
	} else {
		searchOperator = "LIKE ?"
	}
	return
}

func stringIsASCII(input string) bool {
	for _, char := range input {
		if char > 127 {
			return false
		}
	}
	return true
}

// ByQueryNoUser : search torrents according to request without user
func ByQueryNoUser(c *gin.Context, pagenum int) (search structs.TorrentParam, tor []models.Torrent, count int, err error) {
	search, tor, count, err = ByQuery(c, pagenum, true, false, false, false)
	return
}

// ByQueryWithUser : search torrents according to request with user
func ByQueryWithUser(c *gin.Context, pagenum int) (search structs.TorrentParam, tor []models.Torrent, count int, err error) {
	search, tor, count, err = ByQuery(c, pagenum, true, true, false, false)
	return
}

// ByQueryNoCount : search torrents according to request without user and count
func ByQueryNoCount(c *gin.Context, pagenum int) (search structs.TorrentParam, tor []models.Torrent, err error) {
	search, tor, _, err = ByQuery(c, pagenum, false, false, false, false)
	return
}

// ByQueryDeleted : search deleted torrents according to request with user and count
func ByQueryDeleted(c *gin.Context, pagenum int) (search structs.TorrentParam, tor []models.Torrent, count int, err error) {
	search, tor, count, err = ByQuery(c, pagenum, true, true, true, false)
	return
}

// ByQueryNoHidden : search torrents and filter those hidden
func ByQueryNoHidden(c *gin.Context, pagenum int) (search structs.TorrentParam, tor []models.Torrent, count int, err error) {
	search, tor, count, err = ByQuery(c, pagenum, true, false, false, true)
	return
}

// TODO Clean this up
// Some fields are postgres specific (countAll, withUser)
// elasticsearch always provide a count to how many hits
// ES doesn't store users
// deleted is unused because es doesn't index deleted torrents
func ByQuery(c *gin.Context, pagenum int, countAll bool, withUser bool, deleted bool, hidden bool) (structs.TorrentParam, []models.Torrent, int, error) {
	var err error
	if config.Get().Search.EnableElasticSearch && models.ElasticSearchClient != nil && !deleted {
		var torrentParam structs.TorrentParam
		torrentParam.FromRequest(c)
		torrentParam.Offset = uint32(pagenum)
		torrentParam.Hidden = hidden
		torrentParam.Full = withUser
		if found, ok := cache.C.Get(torrentParam.Identifier()); ok {
			torrentCache := found.(*structs.TorrentCache)
			return torrentParam, torrentCache.Torrents, torrentCache.Count, nil
		}
		totalHits, tor, err := torrentParam.Find(models.ElasticSearchClient)
		cache.C.Set(torrentParam.Identifier(), &structs.TorrentCache{tor, int(totalHits)}, 5*time.Minute)
		// Convert back to non-json torrents
		return torrentParam, tor, int(totalHits), err
	}
	log.Errorf("Unable to create elasticsearch client: %s", err)
	log.Errorf("Falling back to postgresql query")
	return byQueryPostgres(c, pagenum, countAll, withUser, deleted, hidden)
}

func byQueryPostgres(c *gin.Context, pagenum int, countAll bool, withUser bool, deleted bool, hidden bool) (
	search structs.TorrentParam, tor []models.Torrent, count int, err error,
) {
	search.FromRequest(c)
	search.Offset = uint32(pagenum)
	search.Hidden = hidden
	search.Deleted = deleted
	search.Full = withUser

	orderBy := search.Sort.ToDBField()
	if search.Sort == structs.Date {
		search.NotNull = search.Sort.ToDBField() + " IS NOT NULL"
	}
	if found, ok := cache.C.Get(search.Identifier()); ok {
		torrentCache := found.(*structs.TorrentCache)
		tor = torrentCache.Torrents
		count = torrentCache.Count
		return
	}

	orderBy += " "

	switch search.Order {
	case true:
		orderBy += "asc"
		if models.ORM.Dialect().GetName() == "postgres" {
			orderBy += " NULLS FIRST"
		}
	case false:
		orderBy += "desc"
		if models.ORM.Dialect().GetName() == "postgres" {
			orderBy += " NULLS LAST"
		}
	}

	parameters := structs.WhereParams{
		Params: make([]interface{}, 0, 64),
	}
	conditions := make([]string, 0, 64)
	if len(search.Category) > 0 {
		conditionsOr := make([]string, len(search.Category))
		for key, val := range search.Category {
			if val.Main > 0 {
				conditionsOr[key] = "(category = ?"
				parameters.Params = append(parameters.Params, val.Main)
				if val.Sub > 0 {
					conditionsOr[key] += " AND sub_category = ?"
					parameters.Params = append(parameters.Params, val.Sub)
				}
				conditionsOr[key] += ")"
			}
		}
		conditions = append(conditions, strings.Join(conditionsOr, " OR "))
	}
	if len(search.Languages) > 0 {
		langs := ""
		for key, val := range search.Languages {
			langs += val.Code
			if key+1 < len(search.Languages) {
				langs += ","
			}
		}
		conditions = append(conditions, "language "+searchOperator)
		parameters.Params = append(parameters.Params, "%"+langs+"%")
	}

	if c.Query("userID") != "" {
		if search.UserID > 0 {
			conditions = append(conditions, "uploader = ?")
			parameters.Params = append(parameters.Params, search.UserID)
			if search.Hidden {
				conditions = append(conditions, "hidden = ?")
				parameters.Params = append(parameters.Params, false)
			}
		} else if search.UserID == 0 {
			conditions = append(conditions, "(uploader = ? OR hidden = ?)")
			parameters.Params = append(parameters.Params, search.UserID)
			parameters.Params = append(parameters.Params, true)
		}
	}
	if search.FromID != 0 {
		conditions = append(conditions, "torrent_id > ?")
		parameters.Params = append(parameters.Params, search.FromID)
	}
	if search.FromDate != "" {
		conditions = append(conditions, "date >= ?")
		parameters.Params = append(parameters.Params, search.FromDate)
	}
	if search.ToDate != "" {
		conditions = append(conditions, "date <= ?")
		parameters.Params = append(parameters.Params, search.ToDate)
	}
	if search.Status != 0 {
		if search.Status == structs.FilterRemakes {
			conditions = append(conditions, "status <> ?")
		} else {
			conditions = append(conditions, "status >= ?")
		}
		parameters.Params = append(parameters.Params, strconv.Itoa(int(search.Status)+1))
	}
	if len(search.NotNull) > 0 {
		conditions = append(conditions, search.NotNull)
	}
	if search.MinSize > 0 {
		conditions = append(conditions, "filesize >= ?")
		parameters.Params = append(parameters.Params, uint64(search.MinSize))
	}
	if search.MaxSize > 0 {
		conditions = append(conditions, "filesize <= ?")
		parameters.Params = append(parameters.Params, uint64(search.MaxSize))
	}

	querySplit := strings.Fields(search.NameLike)
	for _, word := range querySplit {
		firstRune, _ := utf8.DecodeRuneInString(word)
		if len(word) == 1 && unicode.IsPunct(firstRune) {
			// some queries have a single punctuation character
			// which causes a full scan instead of using the index
			// and yields no meaningful results.
			// due to len() == 1 we're just looking at 1-byte/ascii
			// punctuation characters.
			continue
		}

		if useTSQuery && stringIsASCII(word) {
			conditions = append(conditions, "torrent_name @@ plainto_tsquery(?)")
			parameters.Params = append(parameters.Params, word)
		} else {
			// TODO: possible to make this faster?
			conditions = append(conditions, "torrent_name "+searchOperator)
			parameters.Params = append(parameters.Params, "%"+word+"%")
		}
	}

	parameters.Conditions = strings.Join(conditions[:], " AND ")

	log.Infof("SQL query is :: %s\n", parameters.Conditions)

	if deleted {
		tor, count, err = torrents.FindDeleted(&parameters, orderBy, int(search.Max), int(search.Max*(search.Offset-1)))
	} else if countAll && !withUser {
		tor, count, err = torrents.FindOrderBy(&parameters, orderBy, int(search.Max), int(search.Max*(search.Offset-1)))
	} else if countAll && withUser {
		tor, count, err = torrents.FindWithUserOrderBy(&parameters, orderBy, int(search.Max), int(search.Max*(search.Offset-1)))
	} else {
		tor, err = torrents.FindOrderByNoCount(&parameters, orderBy, int(search.Max), int(search.Max*(search.Offset-1)))
	}
	cache.C.Set(search.Identifier(), &structs.TorrentCache{tor, count}, 5*time.Minute)
	return
}

// AuthorizedQuery return a seach byquery according to the bool. If false, it doesn't look for hidden torrents, else it looks for every torrents
func AuthorizedQuery(c *gin.Context, pagenum int, authorized bool) (structs.TorrentParam, []models.Torrent, int, error) {
	if !authorized {
		return ByQuery(c, pagenum, true, true, false, true)
	}
	return ByQuery(c, pagenum, true, true, false, false)
}
<commit_msg>Remove cache on empty results query<commit_after>package search

import (
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"

	"time"

	"github.com/NyaaPantsu/nyaa/config"
	"github.com/NyaaPantsu/nyaa/models"
	"github.com/NyaaPantsu/nyaa/models/torrents"
	"github.com/NyaaPantsu/nyaa/utils/cache"
	"github.com/NyaaPantsu/nyaa/utils/log"
	"github.com/NyaaPantsu/nyaa/utils/search/structs"
	"github.com/gin-gonic/gin"
)

var searchOperator string
var useTSQuery bool

// Configure : initialize search
func Configure(conf *config.SearchConfig) (err error) {
	useTSQuery = false
	// Postgres needs ILIKE for case-insensitivity
	if models.ORM.Dialect().GetName() == "postgres" {
		searchOperator = "ILIKE ?"
		//useTSQuery = true
		// !!DISABLED!! because this makes search a lot stricter
		// (only matches at word borders)
	} else {
		searchOperator = "LIKE ?"
	}
	return
}

func stringIsASCII(input string) bool {
	for _, char := range input {
		if char > 127 {
			return false
		}
	}
	return true
}

// ByQueryNoUser : search torrents according to request without user
func ByQueryNoUser(c *gin.Context, pagenum int) (search structs.TorrentParam, tor []models.Torrent, count int, err error) {
	search, tor, count, err = ByQuery(c, pagenum, true, false, false, false)
	return
}

// ByQueryWithUser : search torrents according to request with user
func ByQueryWithUser(c *gin.Context, pagenum int) (search structs.TorrentParam, tor []models.Torrent, count int, err error) {
	search, tor, count, err = ByQuery(c, pagenum, true, true, false, false)
	return
}

// ByQueryNoCount : search torrents according to request without user and count
func ByQueryNoCount(c *gin.Context, pagenum int) (search structs.TorrentParam, tor []models.Torrent, err error) {
	search, tor, _, err = ByQuery(c, pagenum, false, false, false, false)
	return
}

// ByQueryDeleted : search deleted torrents according to request with user and count
func ByQueryDeleted(c *gin.Context, pagenum int) (search structs.TorrentParam, tor []models.Torrent, count int, err error) {
	search, tor, count, err = ByQuery(c, pagenum, true, true, true, false)
	return
}

// ByQueryNoHidden : search torrents and filter those hidden
func ByQueryNoHidden(c *gin.Context, pagenum int) (search structs.TorrentParam, tor []models.Torrent, count int, err error) {
	search, tor, count, err = ByQuery(c, pagenum, true, false, false, true)
	return
}

// TODO Clean this up
// Some fields are postgres specific (countAll, withUser)
// elasticsearch always provide a count to how many hits
// ES doesn't store users
// deleted is unused because es doesn't index deleted torrents
func ByQuery(c *gin.Context, pagenum int, countAll bool, withUser bool, deleted bool, hidden bool) (structs.TorrentParam, []models.Torrent, int, error) {
	var err error
	if config.Get().Search.EnableElasticSearch && models.ElasticSearchClient != nil && !deleted {
		var torrentParam structs.TorrentParam
		torrentParam.FromRequest(c)
		torrentParam.Offset = uint32(pagenum)
		torrentParam.Hidden = hidden
		torrentParam.Full = withUser
		if found, ok := cache.C.Get(torrentParam.Identifier()); ok {
			torrentCache := found.(*structs.TorrentCache)
			return torrentParam, torrentCache.Torrents, torrentCache.Count, nil
		}
		totalHits, tor, err := torrentParam.Find(models.ElasticSearchClient)
		if totalHits > 0 {
			cache.C.Set(torrentParam.Identifier(), &structs.TorrentCache{tor, int(totalHits)}, 5*time.Minute)
		}
		// Convert back to non-json torrents
		return torrentParam, tor, int(totalHits), err
	}
	log.Errorf("Unable to create elasticsearch client: %s", err)
	log.Errorf("Falling back to postgresql query")
	return byQueryPostgres(c, pagenum, countAll, withUser, deleted, hidden)
}

func byQueryPostgres(c *gin.Context, pagenum int, countAll bool, withUser bool, deleted bool, hidden bool) (
	search structs.TorrentParam, tor []models.Torrent, count int, err error,
) {
	search.FromRequest(c)
	search.Offset = uint32(pagenum)
	search.Hidden = hidden
	search.Deleted = deleted
	search.Full = withUser

	orderBy := search.Sort.ToDBField()
	if search.Sort == structs.Date {
		search.NotNull = search.Sort.ToDBField() + " IS NOT NULL"
	}
	if found, ok := cache.C.Get(search.Identifier()); ok {
		torrentCache := found.(*structs.TorrentCache)
		tor = torrentCache.Torrents
		count = torrentCache.Count
		return
	}

	orderBy += " "

	switch search.Order {
	case true:
		orderBy += "asc"
		if models.ORM.Dialect().GetName() == "postgres" {
			orderBy += " NULLS FIRST"
		}
	case false:
		orderBy += "desc"
		if models.ORM.Dialect().GetName() == "postgres" {
			orderBy += " NULLS LAST"
		}
	}

	parameters := structs.WhereParams{
		Params: make([]interface{}, 0, 64),
	}
	conditions := make([]string, 0, 64)
	if len(search.Category) > 0 {
		conditionsOr := make([]string, len(search.Category))
		for key, val := range search.Category {
			if val.Main > 0 {
				conditionsOr[key] = "(category = ?"
				parameters.Params = append(parameters.Params, val.Main)
				if val.Sub > 0 {
					conditionsOr[key] += " AND sub_category = ?"
					parameters.Params = append(parameters.Params, val.Sub)
				}
				conditionsOr[key] += ")"
			}
		}
		conditions = append(conditions, strings.Join(conditionsOr, " OR "))
	}
	if len(search.Languages) > 0 {
		langs := ""
		for key, val := range search.Languages {
			langs += val.Code
			if key+1 < len(search.Languages) {
				langs += ","
			}
		}
		conditions = append(conditions, "language "+searchOperator)
		parameters.Params = append(parameters.Params, "%"+langs+"%")
	}

	if c.Query("userID") != "" {
		if search.UserID > 0 {
			conditions = append(conditions, "uploader = ?")
			parameters.Params = append(parameters.Params, search.UserID)
			if search.Hidden {
				conditions = append(conditions, "hidden = ?")
				parameters.Params = append(parameters.Params, false)
			}
		} else if search.UserID == 0 {
			conditions = append(conditions, "(uploader = ? OR hidden = ?)")
			parameters.Params = append(parameters.Params, search.UserID)
			parameters.Params = append(parameters.Params, true)
		}
	}
	if search.FromID != 0 {
		conditions = append(conditions, "torrent_id > ?")
		parameters.Params = append(parameters.Params, search.FromID)
	}
	if search.FromDate != "" {
		conditions = append(conditions, "date >= ?")
		parameters.Params = append(parameters.Params, search.FromDate)
	}
	if search.ToDate != "" {
		conditions = append(conditions, "date <= ?")
		parameters.Params = append(parameters.Params, search.ToDate)
	}
	if search.Status != 0 {
		if search.Status == structs.FilterRemakes {
			conditions = append(conditions, "status <> ?")
		} else {
			conditions = append(conditions, "status >= ?")
		}
		parameters.Params = append(parameters.Params, strconv.Itoa(int(search.Status)+1))
	}
	if len(search.NotNull) > 0 {
		conditions = append(conditions, search.NotNull)
	}
	if search.MinSize > 0 {
		conditions = append(conditions, "filesize >= ?")
		parameters.Params = append(parameters.Params, uint64(search.MinSize))
	}
	if search.MaxSize > 0 {
		conditions = append(conditions, "filesize <= ?")
		parameters.Params = append(parameters.Params, uint64(search.MaxSize))
	}

	querySplit := strings.Fields(search.NameLike)
	for _, word := range querySplit {
		firstRune, _ := utf8.DecodeRuneInString(word)
		if len(word) == 1 && unicode.IsPunct(firstRune) {
			// some queries have a single punctuation character
			// which causes a full scan instead of using the index
			// and yields no meaningful results.
			// due to len() == 1 we're just looking at 1-byte/ascii
			// punctuation characters.
			continue
		}

		if useTSQuery && stringIsASCII(word) {
			conditions = append(conditions, "torrent_name @@ plainto_tsquery(?)")
			parameters.Params = append(parameters.Params, word)
		} else {
			// TODO: possible to make this faster?
			conditions = append(conditions, "torrent_name "+searchOperator)
			parameters.Params = append(parameters.Params, "%"+word+"%")
		}
	}

	parameters.Conditions = strings.Join(conditions[:], " AND ")

	log.Infof("SQL query is :: %s\n", parameters.Conditions)

	if deleted {
		tor, count, err = torrents.FindDeleted(&parameters, orderBy, int(search.Max), int(search.Max*(search.Offset-1)))
	} else if countAll && !withUser {
		tor, count, err = torrents.FindOrderBy(&parameters, orderBy, int(search.Max), int(search.Max*(search.Offset-1)))
	} else if countAll && withUser {
		tor, count, err = torrents.FindWithUserOrderBy(&parameters, orderBy, int(search.Max), int(search.Max*(search.Offset-1)))
	} else {
		tor, err = torrents.FindOrderByNoCount(&parameters, orderBy, int(search.Max), int(search.Max*(search.Offset-1)))
	}
	if len(tor) > 0 {
		cache.C.Set(search.Identifier(), &structs.TorrentCache{tor, count}, 5*time.Minute)
	}
	return
}

// AuthorizedQuery return a seach byquery according to the bool. If false, it doesn't look for hidden torrents, else it looks for every torrents
func AuthorizedQuery(c *gin.Context, pagenum int, authorized bool) (structs.TorrentParam, []models.Torrent, int, error) {
	if !authorized {
		return ByQuery(c, pagenum, true, true, false, true)
	}
	return ByQuery(c, pagenum, true, true, false, false)
}
<|endoftext|>