<commit_before>package main

import (
	"log"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"../godump978"
	rtl "github.com/jpoirier/gortlsdr"
)

type UAT struct {
	dev     *rtl.Context
	indexID int
}

type ES struct {
	dev     *rtl.Context
	indexID int
}

var UATDev *UAT
var ESDev *ES

var uat_shutdown chan int
var uat_wg *sync.WaitGroup = &sync.WaitGroup{}

var es_shutdown chan int
var es_wg *sync.WaitGroup = &sync.WaitGroup{}

var maxSignalStrength int

func (e *ES) read() {
	defer es_wg.Done()
	log.Println("Entered ES read() ...")
	cmd := exec.Command("/usr/bin/dump1090", "--net", "--device-index", strconv.Itoa(e.indexID))
	err := cmd.Start()
	if err != nil {
		log.Printf("Error executing /usr/bin/dump1090: %s\n", err.Error())
		return
	}
	log.Println("Executed /usr/bin/dump1090 successfully...")
	for {
		select {
		default:
			time.Sleep(1 * time.Second)
		case <-es_shutdown:
			log.Println("ES read(): shutdown msg received, calling cmd.Process.Kill() ...")
			err := cmd.Process.Kill()
			if err != nil {
				log.Println("\t couldn't kill dump1090: %s", err.Error)
			} else {
				cmd.Wait()
				log.Println("\t kill successful...")
			}
			return
		}
	}
}

func (u *UAT) read() {
	defer uat_wg.Done()
	log.Println("Entered UAT read() ...")
	var buffer = make([]uint8, rtl.DefaultBufLength)
	for {
		select {
		default:
			nRead, err := u.dev.ReadSync(buffer, rtl.DefaultBufLength)
			if err != nil {
				//log.Printf("\tReadSync Failed - error: %s\n", err)
				break
			}
			// log.Printf("\tReadSync %d\n", nRead)
			if nRead > 0 {
				buf := buffer[:nRead]
				godump978.InChan <- buf
			}
		case <-uat_shutdown:
			log.Println("UAT read(): shutdown msg received...")
			return
		}
	}
}

func (e *ES) sdrConfig() (err error) {
	return
}

// Read 978MHz from SDR.
func (u *UAT) sdrConfig() (err error) {
	log.Printf("===== UAT Device name: %s =====\n", rtl.GetDeviceName(u.indexID))
	if u.dev, err = rtl.Open(u.indexID); err != nil {
		log.Printf("\tUAT Open Failed...\n")
		return
	}
	log.Printf("\tGetTunerType: %s\n", u.dev.GetTunerType())

	//---------- Set Tuner Gain ----------
	err = u.dev.SetTunerGainMode(true)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetTunerGainMode Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetTunerGainMode Successful\n")
	}

	tgain := 480
	err = u.dev.SetTunerGain(tgain)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetTunerGain Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetTunerGain Successful\n")
	}

	//---------- Get/Set Sample Rate ----------
	samplerate := 2083334
	err = u.dev.SetSampleRate(samplerate)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetSampleRate Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetSampleRate - rate: %d\n", samplerate)
	}
	log.Printf("\tGetSampleRate: %d\n", u.dev.GetSampleRate())

	//---------- Get/Set Xtal Freq ----------
	rtlFreq, tunerFreq, err := u.dev.GetXtalFreq()
	if err != nil {
		u.dev.Close()
		log.Printf("\tGetXtalFreq Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tGetXtalFreq - Rtl: %d, Tuner: %d\n", rtlFreq, tunerFreq)
	}

	newRTLFreq := 28800000
	newTunerFreq := 28800000
	err = u.dev.SetXtalFreq(newRTLFreq, newTunerFreq)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetXtalFreq Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetXtalFreq - Center freq: %d, Tuner freq: %d\n",
			newRTLFreq, newTunerFreq)
	}

	//---------- Get/Set Center Freq ----------
	err = u.dev.SetCenterFreq(978000000)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetCenterFreq 978MHz Failed, error: %s\n", err)
		return
	} else {
		log.Printf("\tSetCenterFreq 978MHz Successful\n")
	}

	log.Printf("\tGetCenterFreq: %d\n", u.dev.GetCenterFreq())

	//---------- Set Bandwidth ----------
	bw := 1000000
	log.Printf("\tSetting Bandwidth: %d\n", bw)
	if err = u.dev.SetTunerBw(bw); err != nil {
		u.dev.Close()
		log.Printf("\tSetTunerBw %d Failed, error: %s\n", bw, err)
		return
	} else {
		log.Printf("\tSetTunerBw %d Successful\n", bw)
	}

	if err = u.dev.ResetBuffer(); err != nil {
		u.dev.Close()
		log.Printf("\tResetBuffer Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tResetBuffer Successful\n")
	}
	//---------- Get/Set Freq Correction ----------
	freqCorr := u.dev.GetFreqCorrection()
	log.Printf("\tGetFreqCorrection: %d\n", freqCorr)
	err = u.dev.SetFreqCorrection(globalSettings.PPM)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetFreqCorrection %d Failed, error: %s\n", globalSettings.PPM, err)
		return
	} else {
		log.Printf("\tSetFreqCorrection %d Successful\n", globalSettings.PPM)
	}
	return
}

// Read from the godump978 channel - on or off.
func uatReader() {
	log.Println("Entered uatReader() ...")
	for {
		uat := <-godump978.OutChan
		o, msgtype := parseInput(uat)
		if o != nil && msgtype != 0 {
			relayMessage(msgtype, o)
		}
	}
}

func (u *UAT) shutdown() {
	log.Println("Entered UAT shutdown() ...")
	close(uat_shutdown) // signal to shutdown
	log.Println("UAT shutdown(): closing device ...")
	u.dev.Close() // preempt the blocking ReadSync call
	log.Println("UAT shutdown(): calling uat_wg.Wait() ...")
	uat_wg.Wait() // Wait for the goroutine to shutdown
	log.Println("UAT shutdown(): uat_wg.Wait() returned...")
}

func (e *ES) shutdown() {
	log.Println("Entered ES shutdown() ...")
	close(es_shutdown) // signal to shutdown
	log.Println("ES shutdown(): calling es_wg.Wait() ...")
	es_wg.Wait() // Wait for the goroutine to shutdown
	log.Println("ES shutdown(): es_wg.Wait() returned...")
}

var devMap = map[int]string{0: "", 1: ""}

// Watch for config/device changes.
func sdrWatcher() {
	for {
		time.Sleep(1 * time.Second)
		count := rtl.GetDeviceCount()
		atomic.StoreUint32(&globalStatus.Devices, uint32(count))
		log.Println("DeviceCount...", count)

		// support two and only two dongles
		if count > 2 {
			count = 2
		}

		// cleanup if necessary
		if count < 1 || (!globalSettings.UAT_Enabled && !globalSettings.ES_Enabled) {
			log.Println("count == 0, doing cleanup if necessary...")
			if UATDev != nil {
				UATDev.shutdown()
				UATDev = nil
			}
			if ESDev != nil {
				ESDev.shutdown()
				ESDev = nil
			}
			continue
		}

		if count == 1 {
			if UATDev != nil && ESDev == nil {
				UATDev.indexID = 0
			} else if UATDev == nil && ESDev != nil {
				ESDev.indexID = 0
			}
		}

		ids := []string{"", ""}
		for i := 0; i < count; i++ {
			// manufact, product, serial, err
			_, _, s, _ := rtl.GetDeviceUsbStrings(i)
			ids[i] = s
		}

		// UAT specific handling
		// When count is one, favor UAT in the case where the user
		// has enabled both UAT and ES via the web interface.
		id := 0
		if globalSettings.UAT_Enabled {
			// log.Println("globalSettings.UAT_Enabled == true")
			if count == 1 {
				if ESDev != nil {
					ESDev.shutdown()
					ESDev = nil
				}
			} else { // count == 2
				if UATDev == nil && ESDev != nil {
					if ESDev.indexID == 0 {
						id = 1
					}
				}
			}

			if UATDev == nil {
				// log.Println("\tUATDev == nil")
				// preference check based on stratux
				// hardware serial when it exists
				serial := ids[id]
				if strings.HasPrefix(serial, "stratux:1090") {
					log.Println("Settings conflict: 978UAT set via WebUI but hardware serial says stratux:1090")
				} else {
					UATDev = &UAT{indexID: id}
					if err := UATDev.sdrConfig(); err != nil {
						log.Printf("UATDev = &UAT{indexID: id} failed: %s\n", err)
						UATDev = nil
					} else {
						uat_shutdown = make(chan int)
						uat_wg.Add(1)
						go UATDev.read()
					}
				}
			}
		} else if UATDev != nil {
			UATDev.shutdown()
			UATDev = nil
			if count == 1 && ESDev != nil {
				ESDev.indexID = 0
			}
		}

		// ES specific handling
		id = 0
		if globalSettings.ES_Enabled {
			// log.Println("globalSettings.ES_Enabled == true")
			if count == 1 {
				if globalSettings.UAT_Enabled {
					// defer to the UAT handler
					goto End
				}
			} else { // count == 2
				if ESDev == nil && UATDev != nil {
					if UATDev.indexID == 0 {
						id = 1
					}
				}
			}

			if ESDev == nil {
				// log.Println("\tESDev == nil")
				// preference check based on stratux
				// hardware serial when it exists
				serial := ids[id]
				if strings.HasPrefix(serial, "stratux:978") {
					log.Println("Settings conflict: 1090ES set via WebUI but hardware serial says stratux:978")
				} else {
					ESDev = &ES{indexID: id}
					if err := ESDev.sdrConfig(); err != nil {
						log.Printf("ESDev = &ES{indexID: id} failed: %s\n", err)
						ESDev = nil
					} else {
						es_shutdown = make(chan int)
						es_wg.Add(1)
						go ESDev.read()
					}
				}
			}
		} else if ESDev != nil {
			ESDev.shutdown()
			ESDev = nil
			if count == 1 && UATDev != nil {
				UATDev.indexID = 0
			}
		}
	End:
	}
}

func sdrInit() {
	go sdrWatcher()
	go uatReader()
	godump978.Dump978Init()
	go godump978.ProcessDataFromChannel()
}
<commit_msg>commnet out device count logging<commit_after>package main

import (
	"log"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"../godump978"
	rtl "github.com/jpoirier/gortlsdr"
)

type UAT struct {
	dev     *rtl.Context
	indexID int
}

type ES struct {
	dev     *rtl.Context
	indexID int
}

var UATDev *UAT
var ESDev *ES

var uat_shutdown chan int
var uat_wg *sync.WaitGroup = &sync.WaitGroup{}

var es_shutdown chan int
var es_wg *sync.WaitGroup = &sync.WaitGroup{}

var maxSignalStrength int

func (e *ES) read() {
	defer es_wg.Done()
	log.Println("Entered ES read() ...")
	cmd := exec.Command("/usr/bin/dump1090", "--net", "--device-index", strconv.Itoa(e.indexID))
	err := cmd.Start()
	if err != nil {
		log.Printf("Error executing /usr/bin/dump1090: %s\n", err.Error())
		return
	}
	log.Println("Executed /usr/bin/dump1090 successfully...")
	for {
		select {
		default:
			time.Sleep(1 * time.Second)
		case <-es_shutdown:
			log.Println("ES read(): shutdown msg received, calling cmd.Process.Kill() ...")
			err := cmd.Process.Kill()
			if err != nil {
				log.Println("\t couldn't kill dump1090: %s", err.Error)
			} else {
				cmd.Wait()
				log.Println("\t kill successful...")
			}
			return
		}
	}
}

func (u *UAT) read() {
	defer uat_wg.Done()
	log.Println("Entered UAT read() ...")
	var buffer = make([]uint8, rtl.DefaultBufLength)
	for {
		select {
		default:
			nRead, err := u.dev.ReadSync(buffer, rtl.DefaultBufLength)
			if err != nil {
				//log.Printf("\tReadSync Failed - error: %s\n", err)
				break
			}
			// log.Printf("\tReadSync %d\n", nRead)
			if nRead > 0 {
				buf := buffer[:nRead]
				godump978.InChan <- buf
			}
		case <-uat_shutdown:
			log.Println("UAT read(): shutdown msg received...")
			return
		}
	}
}

func (e *ES) sdrConfig() (err error) {
	return
}

// Read 978MHz from SDR.
func (u *UAT) sdrConfig() (err error) {
	log.Printf("===== UAT Device name: %s =====\n", rtl.GetDeviceName(u.indexID))
	if u.dev, err = rtl.Open(u.indexID); err != nil {
		log.Printf("\tUAT Open Failed...\n")
		return
	}
	log.Printf("\tGetTunerType: %s\n", u.dev.GetTunerType())

	//---------- Set Tuner Gain ----------
	err = u.dev.SetTunerGainMode(true)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetTunerGainMode Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetTunerGainMode Successful\n")
	}

	tgain := 480
	err = u.dev.SetTunerGain(tgain)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetTunerGain Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetTunerGain Successful\n")
	}

	//---------- Get/Set Sample Rate ----------
	samplerate := 2083334
	err = u.dev.SetSampleRate(samplerate)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetSampleRate Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetSampleRate - rate: %d\n", samplerate)
	}
	log.Printf("\tGetSampleRate: %d\n", u.dev.GetSampleRate())

	//---------- Get/Set Xtal Freq ----------
	rtlFreq, tunerFreq, err := u.dev.GetXtalFreq()
	if err != nil {
		u.dev.Close()
		log.Printf("\tGetXtalFreq Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tGetXtalFreq - Rtl: %d, Tuner: %d\n", rtlFreq, tunerFreq)
	}

	newRTLFreq := 28800000
	newTunerFreq := 28800000
	err = u.dev.SetXtalFreq(newRTLFreq, newTunerFreq)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetXtalFreq Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetXtalFreq - Center freq: %d, Tuner freq: %d\n",
			newRTLFreq, newTunerFreq)
	}

	//---------- Get/Set Center Freq ----------
	err = u.dev.SetCenterFreq(978000000)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetCenterFreq 978MHz Failed, error: %s\n", err)
		return
	} else {
		log.Printf("\tSetCenterFreq 978MHz Successful\n")
	}

	log.Printf("\tGetCenterFreq: %d\n", u.dev.GetCenterFreq())

	//---------- Set Bandwidth ----------
	bw := 1000000
	log.Printf("\tSetting Bandwidth: %d\n", bw)
	if err = u.dev.SetTunerBw(bw); err != nil {
		u.dev.Close()
		log.Printf("\tSetTunerBw %d Failed, error: %s\n", bw, err)
		return
	} else {
		log.Printf("\tSetTunerBw %d Successful\n", bw)
	}

	if err = u.dev.ResetBuffer(); err != nil {
		u.dev.Close()
		log.Printf("\tResetBuffer Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tResetBuffer Successful\n")
	}
	//---------- Get/Set Freq Correction ----------
	freqCorr := u.dev.GetFreqCorrection()
	log.Printf("\tGetFreqCorrection: %d\n", freqCorr)
	err = u.dev.SetFreqCorrection(globalSettings.PPM)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetFreqCorrection %d Failed, error: %s\n", globalSettings.PPM, err)
		return
	} else {
		log.Printf("\tSetFreqCorrection %d Successful\n", globalSettings.PPM)
	}
	return
}

// Read from the godump978 channel - on or off.
func uatReader() {
	log.Println("Entered uatReader() ...")
	for {
		uat := <-godump978.OutChan
		o, msgtype := parseInput(uat)
		if o != nil && msgtype != 0 {
			relayMessage(msgtype, o)
		}
	}
}

func (u *UAT) shutdown() {
	log.Println("Entered UAT shutdown() ...")
	close(uat_shutdown) // signal to shutdown
	log.Println("UAT shutdown(): closing device ...")
	u.dev.Close() // preempt the blocking ReadSync call
	log.Println("UAT shutdown(): calling uat_wg.Wait() ...")
	uat_wg.Wait() // Wait for the goroutine to shutdown
	log.Println("UAT shutdown(): uat_wg.Wait() returned...")
}

func (e *ES) shutdown() {
	log.Println("Entered ES shutdown() ...")
	close(es_shutdown) // signal to shutdown
	log.Println("ES shutdown(): calling es_wg.Wait() ...")
	es_wg.Wait() // Wait for the goroutine to shutdown
	log.Println("ES shutdown(): es_wg.Wait() returned...")
}

var devMap = map[int]string{0: "", 1: ""}

// Watch for config/device changes.
func sdrWatcher() {
	for {
		time.Sleep(1 * time.Second)
		count := rtl.GetDeviceCount()
		atomic.StoreUint32(&globalStatus.Devices, uint32(count))
		// log.Println("DeviceCount...", count)

		// support two and only two dongles
		if count > 2 {
			count = 2
		}

		// cleanup if necessary
		if count < 1 || (!globalSettings.UAT_Enabled && !globalSettings.ES_Enabled) {
			log.Println("count == 0, doing cleanup if necessary...")
			if UATDev != nil {
				UATDev.shutdown()
				UATDev = nil
			}
			if ESDev != nil {
				ESDev.shutdown()
				ESDev = nil
			}
			continue
		}

		if count == 1 {
			if UATDev != nil && ESDev == nil {
				UATDev.indexID = 0
			} else if UATDev == nil && ESDev != nil {
				ESDev.indexID = 0
			}
		}

		ids := []string{"", ""}
		for i := 0; i < count; i++ {
			// manufact, product, serial, err
			_, _, s, _ := rtl.GetDeviceUsbStrings(i)
			ids[i] = s
		}

		// UAT specific handling
		// When count is one, favor UAT in the case where the user
		// has enabled both UAT and ES via the web interface.
		id := 0
		if globalSettings.UAT_Enabled {
			// log.Println("globalSettings.UAT_Enabled == true")
			if count == 1 {
				if ESDev != nil {
					ESDev.shutdown()
					ESDev = nil
				}
			} else { // count == 2
				if UATDev == nil && ESDev != nil {
					if ESDev.indexID == 0 {
						id = 1
					}
				}
			}

			if UATDev == nil {
				// log.Println("\tUATDev == nil")
				// preference check based on stratux
				// hardware serial when it exists
				serial := ids[id]
				if strings.HasPrefix(serial, "stratux:1090") {
					log.Println("Settings conflict: 978UAT set via WebUI but hardware serial says stratux:1090")
				} else {
					UATDev = &UAT{indexID: id}
					if err := UATDev.sdrConfig(); err != nil {
						log.Printf("UATDev = &UAT{indexID: id} failed: %s\n", err)
						UATDev = nil
					} else {
						uat_shutdown = make(chan int)
						uat_wg.Add(1)
						go UATDev.read()
					}
				}
			}
		} else if UATDev != nil {
			UATDev.shutdown()
			UATDev = nil
			if count == 1 && ESDev != nil {
				ESDev.indexID = 0
			}
		}

		// ES specific handling
		id = 0
		if globalSettings.ES_Enabled {
			// log.Println("globalSettings.ES_Enabled == true")
			if count == 1 {
				if globalSettings.UAT_Enabled {
					// defer to the UAT handler
					goto End
				}
			} else { // count == 2
				if ESDev == nil && UATDev != nil {
					if UATDev.indexID == 0 {
						id = 1
					}
				}
			}

			if ESDev == nil {
				// log.Println("\tESDev == nil")
				// preference check based on stratux
				// hardware serial when it exists
				serial := ids[id]
				if strings.HasPrefix(serial, "stratux:978") {
					log.Println("Settings conflict: 1090ES set via WebUI but hardware serial says stratux:978")
				} else {
					ESDev = &ES{indexID: id}
					if err := ESDev.sdrConfig(); err != nil {
						log.Printf("ESDev = &ES{indexID: id} failed: %s\n", err)
						ESDev = nil
					} else {
						es_shutdown = make(chan int)
						es_wg.Add(1)
						go ESDev.read()
					}
				}
			}
		} else if ESDev != nil {
			ESDev.shutdown()
			ESDev = nil
			if count == 1 && UATDev != nil {
				UATDev.indexID = 0
			}
		}
	End:
	}
}

func sdrInit() {
	go sdrWatcher()
	go uatReader()
	godump978.Dump978Init()
	go godump978.ProcessDataFromChannel()
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"log"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"../godump978"
	rtl "github.com/jpoirier/gortlsdr"
)

type UAT struct {
	dev     *rtl.Context
	indexID int
}

type ES struct {
	dev     *rtl.Context
	indexID int
}

var UATDev *UAT
var ESDev *ES

var uat_shutdown chan int
var uat_wg *sync.WaitGroup = &sync.WaitGroup{}

var es_shutdown chan int
var es_wg *sync.WaitGroup = &sync.WaitGroup{}

var maxSignalStrength int

func readToChan(fp io.ReadCloser, ch chan []byte) {
	for {
		buf := make([]byte, 1024)
		n, err := fp.Read(buf)
		if n > 0 {
			ch <- buf[:n]
		} else if err != nil {
			return
		}
	}
}

func (e *ES) read() {
	defer es_wg.Done()
	log.Println("Entered ES read() ...")
	cmd := exec.Command("/usr/bin/dump1090", "--net", "--device-index", strconv.Itoa(e.indexID))
	stdout, _ := cmd.StdoutPipe()
	stderr, _ := cmd.StderrPipe()

	outputChan := make(chan []byte, 1024)

	go readToChan(stdout, outputChan)
	go readToChan(stderr, outputChan)

	err := cmd.Start()
	if err != nil {
		log.Printf("Error executing /usr/bin/dump1090: %s\n", err.Error())
		return
	}
	log.Println("Executed /usr/bin/dump1090 successfully...")

	for {
		select {
		case buf := <-outputChan:
			replayLog(string(buf), MSGCLASS_DUMP1090)

		case <-es_shutdown:
			log.Println("ES read(): shutdown msg received, calling cmd.Process.Kill() ...")
			err := cmd.Process.Kill()
			if err != nil {
				log.Println("\t couldn't kill dump1090: %s", err.Error)
			} else {
				cmd.Wait()
				log.Println("\t kill successful...")
			}
			return
		default:
			time.Sleep(1 * time.Second)

		}
	}
}

func (u *UAT) read() {
	defer uat_wg.Done()
	log.Println("Entered UAT read() ...")
	var buffer = make([]uint8, rtl.DefaultBufLength)
	for {
		select {
		default:
			nRead, err := u.dev.ReadSync(buffer, rtl.DefaultBufLength)
			if err != nil {
				//log.Printf("\tReadSync Failed - error: %s\n", err)
				break
			}
			// log.Printf("\tReadSync %d\n", nRead)
			if nRead > 0 {
				buf := buffer[:nRead]
				godump978.InChan <- buf
			}
		case <-uat_shutdown:
			log.Println("UAT read(): shutdown msg received...")
			return
		}
	}
}

func (e *ES) sdrConfig() (err error) {
	return
}

// Read 978MHz from SDR.
func (u *UAT) sdrConfig() (err error) {
	log.Printf("===== UAT Device name: %s =====\n", rtl.GetDeviceName(u.indexID))
	if u.dev, err = rtl.Open(u.indexID); err != nil {
		log.Printf("\tUAT Open Failed...\n")
		return
	}
	log.Printf("\tGetTunerType: %s\n", u.dev.GetTunerType())

	//---------- Set Tuner Gain ----------
	err = u.dev.SetTunerGainMode(true)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetTunerGainMode Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetTunerGainMode Successful\n")
	}

	tgain := 480
	err = u.dev.SetTunerGain(tgain)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetTunerGain Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetTunerGain Successful\n")
	}

	//---------- Get/Set Sample Rate ----------
	samplerate := 2083334
	err = u.dev.SetSampleRate(samplerate)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetSampleRate Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetSampleRate - rate: %d\n", samplerate)
	}
	log.Printf("\tGetSampleRate: %d\n", u.dev.GetSampleRate())

	//---------- Get/Set Xtal Freq ----------
	rtlFreq, tunerFreq, err := u.dev.GetXtalFreq()
	if err != nil {
		u.dev.Close()
		log.Printf("\tGetXtalFreq Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tGetXtalFreq - Rtl: %d, Tuner: %d\n", rtlFreq, tunerFreq)
	}

	newRTLFreq := 28800000
	newTunerFreq := 28800000
	err = u.dev.SetXtalFreq(newRTLFreq, newTunerFreq)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetXtalFreq Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetXtalFreq - Center freq: %d, Tuner freq: %d\n",
			newRTLFreq, newTunerFreq)
	}

	//---------- Get/Set Center Freq ----------
	err = u.dev.SetCenterFreq(978000000)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetCenterFreq 978MHz Failed, error: %s\n", err)
		return
	} else {
		log.Printf("\tSetCenterFreq 978MHz Successful\n")
	}

	log.Printf("\tGetCenterFreq: %d\n", u.dev.GetCenterFreq())

	//---------- Set Bandwidth ----------
	bw := 1000000
	log.Printf("\tSetting Bandwidth: %d\n", bw)
	if err = u.dev.SetTunerBw(bw); err != nil {
		u.dev.Close()
		log.Printf("\tSetTunerBw %d Failed, error: %s\n", bw, err)
		return
	} else {
		log.Printf("\tSetTunerBw %d Successful\n", bw)
	}

	if err = u.dev.ResetBuffer(); err != nil {
		u.dev.Close()
		log.Printf("\tResetBuffer Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tResetBuffer Successful\n")
	}
	//---------- Get/Set Freq Correction ----------
	freqCorr := u.dev.GetFreqCorrection()
	log.Printf("\tGetFreqCorrection: %d\n", freqCorr)
	err = u.dev.SetFreqCorrection(globalSettings.PPM)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetFreqCorrection %d Failed, error: %s\n", globalSettings.PPM, err)
		return
	} else {
		log.Printf("\tSetFreqCorrection %d Successful\n", globalSettings.PPM)
	}
	return
}

// Read from the godump978 channel - on or off.
func uatReader() {
	log.Println("Entered uatReader() ...")
	for {
		uat := <-godump978.OutChan
		o, msgtype := parseInput(uat)
		if o != nil && msgtype != 0 {
			relayMessage(msgtype, o)
		}
	}
}

func (u *UAT) writeID() error {
	info, err := u.dev.GetHwInfo()
	if err != nil {
		return err
	}
	info.Serial = "stratux:978"
	return u.dev.SetHwInfo(info)
}

func (e *ES) writeID() error {
	info, err := e.dev.GetHwInfo()
	if err != nil {
		return err
	}
	info.Serial = "stratux:1090"
	return e.dev.SetHwInfo(info)
}

func (u *UAT) shutdown() {
	log.Println("Entered UAT shutdown() ...")
	close(uat_shutdown) // signal to shutdown
	log.Println("UAT shutdown(): closing device ...")
	u.dev.Close() // preempt the blocking ReadSync call
	log.Println("UAT shutdown(): calling uat_wg.Wait() ...")
	uat_wg.Wait() // Wait for the goroutine to shutdown
	log.Println("UAT shutdown(): uat_wg.Wait() returned...")
}

func (e *ES) shutdown() {
	log.Println("Entered ES shutdown() ...")
	close(es_shutdown) // signal to shutdown
	log.Println("ES shutdown(): calling es_wg.Wait() ...")
	es_wg.Wait() // Wait for the goroutine to shutdown
	log.Println("ES shutdown(): es_wg.Wait() returned...")
}

var devMap = map[int]string{0: "", 1: ""}

// Watch for config/device changes.
func sdrWatcher() {
	for {
		time.Sleep(1 * time.Second)
		count := rtl.GetDeviceCount()
		atomic.StoreUint32(&globalStatus.Devices, uint32(count))
		// log.Println("DeviceCount...", count)

		// support two and only two dongles
		if count > 2 {
			count = 2
		}

		// cleanup if necessary
		if count < 1 || (!globalSettings.UAT_Enabled && !globalSettings.ES_Enabled) {
			log.Println("count == 0, doing cleanup if necessary...")
			if UATDev != nil {
				UATDev.shutdown()
				UATDev = nil
			}
			if ESDev != nil {
				ESDev.shutdown()
				ESDev = nil
			}
			continue
		}

		if count == 1 {
			if UATDev != nil && ESDev == nil {
				UATDev.indexID = 0
			} else if UATDev == nil && ESDev != nil {
				ESDev.indexID = 0
			}
		}

		ids := []string{"", ""}
		for i := 0; i < count; i++ {
			// manufact, product, serial, err
			_, _, s, _ := rtl.GetDeviceUsbStrings(i)
			ids[i] = s
		}

		// UAT specific handling
		// When count is one, favor UAT in the case where the user
		// has enabled both UAT and ES via the web interface.
		id := 0
		if globalSettings.UAT_Enabled {
			// log.Println("globalSettings.UAT_Enabled == true")
			if count == 1 {
				if ESDev != nil {
					ESDev.shutdown()
					ESDev = nil
				}
			} else { // count == 2
				if UATDev == nil && ESDev != nil {
					if ESDev.indexID == 0 {
						id = 1
					}
				}
			}

			if UATDev == nil {
				// log.Println("\tUATDev == nil")
				// preference check based on stratux
				// hardware serial when it exists
				serial := ids[id]
				if strings.HasPrefix(serial, "stratux:1090") {
					log.Println("Settings conflict: 978UAT set via WebUI but hardware serial says stratux:1090")
				} else {
					UATDev = &UAT{indexID: id}
					if err := UATDev.sdrConfig(); err != nil {
						log.Printf("UATDev = &UAT{indexID: id} failed: %s\n", err)
						UATDev = nil
					} else {
						uat_shutdown = make(chan int)
						uat_wg.Add(1)
						go UATDev.read()
					}
				}
			}
		} else if UATDev != nil {
			UATDev.shutdown()
			UATDev = nil
			if count == 1 && ESDev != nil {
				ESDev.indexID = 0
			}
		}

		// ES specific handling
		id = 0
		if globalSettings.ES_Enabled {
			// log.Println("globalSettings.ES_Enabled == true")
			if count == 1 {
				if globalSettings.UAT_Enabled {
					// defer to the UAT handler
					goto End
				}
			} else { // count == 2
				if ESDev == nil && UATDev != nil {
					if UATDev.indexID == 0 {
						id = 1
					}
				}
			}

			if ESDev == nil {
				// log.Println("\tESDev == nil")
				// preference check based on stratux
				// hardware serial when it exists
				serial := ids[id]
				if strings.HasPrefix(serial, "stratux:978") {
					log.Println("Settings conflict: 1090ES set via WebUI but hardware serial says stratux:978")
				} else {
					ESDev = &ES{indexID: id}
					if err := ESDev.sdrConfig(); err != nil {
						log.Printf("ESDev = &ES{indexID: id} failed: %s\n", err)
						ESDev = nil
					} else {
						es_shutdown = make(chan int)
						es_wg.Add(1)
						go ESDev.read()
					}
				}
			}
		} else if ESDev != nil {
			ESDev.shutdown()
			ESDev = nil
			if count == 1 && UATDev != nil {
				UATDev.indexID = 0
			}
		}
	End:
	}
}

func sdrInit() {
	go sdrWatcher()
	go uatReader()
	godump978.Dump978Init()
	go godump978.ProcessDataFromChannel()
}
<commit_msg>Pause UAT reading 75% of the time when no UAT messages are being received.<commit_after>package main

import (
	"io"
	"log"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"../godump978"
	rtl "github.com/jpoirier/gortlsdr"
)

type UAT struct {
	dev     *rtl.Context
	indexID int
}

type ES struct {
	dev     *rtl.Context
	indexID int
}

var UATDev *UAT
var ESDev *ES

var uat_shutdown chan int
var uat_wg *sync.WaitGroup = &sync.WaitGroup{}

var es_shutdown chan int
var es_wg *sync.WaitGroup = &sync.WaitGroup{}

var maxSignalStrength int

func readToChan(fp io.ReadCloser, ch chan []byte) {
	for {
		buf := make([]byte, 1024)
		n, err := fp.Read(buf)
		if n > 0 {
			ch <- buf[:n]
		} else if err != nil {
			return
		}
	}
}

func (e *ES) read() {
	defer es_wg.Done()
	log.Println("Entered ES read() ...")
	cmd := exec.Command("/usr/bin/dump1090", "--net", "--device-index", strconv.Itoa(e.indexID))
	stdout, _ := cmd.StdoutPipe()
	stderr, _ := cmd.StderrPipe()

	outputChan := make(chan []byte, 1024)

	go readToChan(stdout, outputChan)
	go readToChan(stderr, outputChan)

	err := cmd.Start()
	if err != nil {
		log.Printf("Error executing /usr/bin/dump1090: %s\n", err.Error())
		return
	}
	log.Println("Executed /usr/bin/dump1090 successfully...")

	for {
		select {
		case buf := <-outputChan:
			replayLog(string(buf), MSGCLASS_DUMP1090)

		case <-es_shutdown:
			log.Println("ES read(): shutdown msg received, calling cmd.Process.Kill() ...")
			err := cmd.Process.Kill()
			if err != nil {
				log.Println("\t couldn't kill dump1090: %s", err.Error)
			} else {
				cmd.Wait()
				log.Println("\t kill successful...")
			}
			return
		default:
			time.Sleep(1 * time.Second)

		}
	}
}

func (u *UAT) read() {
	defer uat_wg.Done()
	log.Println("Entered UAT read() ...")
	var buffer = make([]uint8, rtl.DefaultBufLength)

	slowReadTimer := time.NewTicker(5 * time.Second)
	for {
		select {
		default:
			nRead, err := u.dev.ReadSync(buffer, rtl.DefaultBufLength)
			if err != nil {
				//log.Printf("\tReadSync Failed - error: %s\n", err)
				break
			}
			// log.Printf("\tReadSync %d\n", nRead)
			if nRead > 0 {
				buf := buffer[:nRead]
				godump978.InChan <- buf
			}
		case <-uat_shutdown:
			log.Println("UAT read(): shutdown msg received...")
			return
		case <-slowReadTimer.C:
			// Check (every 5 seconds) if we haven't received any UAT messages in the last minute. If that's the case, sleep for 15 seconds.
			// Puts the RTL-SDR reading and UAT message parsing on a 25% duty cycle unless we're receiving messages, then it goes back to 100%.
			if globalStatus.UAT_messages_last_minute == 0 {
				time.Sleep(15 * time.Second)
			}
		}
	}
}

func (e *ES) sdrConfig() (err error) {
	return
}

// Read 978MHz from SDR.
func (u *UAT) sdrConfig() (err error) {
	log.Printf("===== UAT Device name: %s =====\n", rtl.GetDeviceName(u.indexID))
	if u.dev, err = rtl.Open(u.indexID); err != nil {
		log.Printf("\tUAT Open Failed...\n")
		return
	}
	log.Printf("\tGetTunerType: %s\n", u.dev.GetTunerType())

	//---------- Set Tuner Gain ----------
	err = u.dev.SetTunerGainMode(true)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetTunerGainMode Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetTunerGainMode Successful\n")
	}

	tgain := 480
	err = u.dev.SetTunerGain(tgain)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetTunerGain Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetTunerGain Successful\n")
	}

	//---------- Get/Set Sample Rate ----------
	samplerate := 2083334
	err = u.dev.SetSampleRate(samplerate)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetSampleRate Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetSampleRate - rate: %d\n", samplerate)
	}
	log.Printf("\tGetSampleRate: %d\n", u.dev.GetSampleRate())

	//---------- Get/Set Xtal Freq ----------
	rtlFreq, tunerFreq, err := u.dev.GetXtalFreq()
	if err != nil {
		u.dev.Close()
		log.Printf("\tGetXtalFreq Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tGetXtalFreq - Rtl: %d, Tuner: %d\n", rtlFreq, tunerFreq)
	}

	newRTLFreq := 28800000
	newTunerFreq := 28800000
	err = u.dev.SetXtalFreq(newRTLFreq, newTunerFreq)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetXtalFreq Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tSetXtalFreq - Center freq: %d, Tuner freq: %d\n",
			newRTLFreq, newTunerFreq)
	}

	//---------- Get/Set Center Freq ----------
	err = u.dev.SetCenterFreq(978000000)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetCenterFreq 978MHz Failed, error: %s\n", err)
		return
	} else {
		log.Printf("\tSetCenterFreq 978MHz Successful\n")
	}

	log.Printf("\tGetCenterFreq: %d\n", u.dev.GetCenterFreq())

	//---------- Set Bandwidth ----------
	bw := 1000000
	log.Printf("\tSetting Bandwidth: %d\n", bw)
	if err = u.dev.SetTunerBw(bw); err != nil {
		u.dev.Close()
		log.Printf("\tSetTunerBw %d Failed, error: %s\n", bw, err)
		return
	} else {
		log.Printf("\tSetTunerBw %d Successful\n", bw)
	}

	if err = u.dev.ResetBuffer(); err != nil {
		u.dev.Close()
		log.Printf("\tResetBuffer Failed - error: %s\n", err)
		return
	} else {
		log.Printf("\tResetBuffer Successful\n")
	}
	//---------- Get/Set Freq Correction ----------
	freqCorr := u.dev.GetFreqCorrection()
	log.Printf("\tGetFreqCorrection: %d\n", freqCorr)
	err = u.dev.SetFreqCorrection(globalSettings.PPM)
	if err != nil {
		u.dev.Close()
		log.Printf("\tSetFreqCorrection %d Failed, error: %s\n", globalSettings.PPM, err)
		return
	} else {
		log.Printf("\tSetFreqCorrection %d Successful\n", globalSettings.PPM)
	}
	return
}

// Read from the godump978 channel - on or off.
func uatReader() {
	log.Println("Entered uatReader() ...")
	for {
		uat := <-godump978.OutChan
		o, msgtype := parseInput(uat)
		if o != nil && msgtype != 0 {
			relayMessage(msgtype, o)
		}
	}
}

func (u *UAT) writeID() error {
	info, err := u.dev.GetHwInfo()
	if err != nil {
		return err
	}
	info.Serial = "stratux:978"
	return u.dev.SetHwInfo(info)
}

func (e *ES) writeID() error {
	info, err := e.dev.GetHwInfo()
	if err != nil {
		return err
	}
	info.Serial = "stratux:1090"
	return e.dev.SetHwInfo(info)
}

func (u *UAT) shutdown() {
	log.Println("Entered UAT shutdown() ...")
	close(uat_shutdown) // signal to shutdown
	log.Println("UAT shutdown(): closing device ...")
	u.dev.Close() // preempt the blocking ReadSync call
	log.Println("UAT shutdown(): calling uat_wg.Wait() ...")
	uat_wg.Wait() // Wait for the goroutine to shutdown
	log.Println("UAT shutdown(): uat_wg.Wait() returned...")
}

func (e *ES) shutdown() {
	log.Println("Entered ES shutdown() ...")
	close(es_shutdown) // signal to shutdown
	log.Println("ES shutdown(): calling es_wg.Wait() ...")
	es_wg.Wait() // Wait for the goroutine to shutdown
	log.Println("ES shutdown(): es_wg.Wait() returned...")
}

var devMap = map[int]string{0: "", 1: ""}

// Watch for config/device changes.
func sdrWatcher() {
	for {
		time.Sleep(1 * time.Second)
		count := rtl.GetDeviceCount()
		atomic.StoreUint32(&globalStatus.Devices, uint32(count))
		// log.Println("DeviceCount...", count)

		// support two and only two dongles
		if count > 2 {
			count = 2
		}

		// cleanup if necessary
		if count < 1 || (!globalSettings.UAT_Enabled && !globalSettings.ES_Enabled) {
			log.Println("count == 0, doing cleanup if necessary...")
			if UATDev != nil {
				UATDev.shutdown()
				UATDev = nil
			}
			if ESDev != nil {
				ESDev.shutdown()
				ESDev = nil
			}
			continue
		}

		if count == 1 {
			if UATDev != nil && ESDev == nil {
				UATDev.indexID = 0
			} else if UATDev == nil && ESDev != nil {
				ESDev.indexID = 0
			}
		}

		ids := []string{"", ""}
		for i := 0; i < count; i++ {
			// manufact, product, serial, err
			_, _, s, _ := rtl.GetDeviceUsbStrings(i)
			ids[i] = s
		}

		// UAT specific handling
		// When count is one, favor UAT in the case where the user
		// has enabled both UAT and ES via the web interface.
		id := 0
		if globalSettings.UAT_Enabled {
			// log.Println("globalSettings.UAT_Enabled == true")
			if count == 1 {
				if ESDev != nil {
					ESDev.shutdown()
					ESDev = nil
				}
			} else { // count == 2
				if UATDev == nil && ESDev != nil {
					if ESDev.indexID == 0 {
						id = 1
					}
				}
			}

			if UATDev == nil {
				// log.Println("\tUATDev == nil")
				// preference check based on stratux
				// hardware serial when it exists
				serial := ids[id]
				if strings.HasPrefix(serial, "stratux:1090") {
					log.Println("Settings conflict: 978UAT set via WebUI but hardware serial says stratux:1090")
				} else {
					UATDev = &UAT{indexID: id}
					if err := UATDev.sdrConfig(); err != nil {
						log.Printf("UATDev = &UAT{indexID: id} failed: %s\n", err)
						UATDev = nil
					} else {
						uat_shutdown = make(chan int)
						uat_wg.Add(1)
						go UATDev.read()
					}
				}
			}
		} else if UATDev != nil {
			UATDev.shutdown()
			UATDev = nil
			if count == 1 && ESDev != nil {
				ESDev.indexID = 0
			}
		}

		// ES specific handling
		id = 0
		if globalSettings.ES_Enabled {
			// log.Println("globalSettings.ES_Enabled == true")
			if count == 1 {
				if globalSettings.UAT_Enabled {
					// defer to the UAT handler
					goto End
				}
			} else { // count == 2
				if ESDev == nil && UATDev != nil {
					if UATDev.indexID == 0 {
						id = 1
					}
				}
			}

			if ESDev == nil {
				// log.Println("\tESDev == nil")
				// preference check based on stratux
				// hardware serial when it exists
				serial := ids[id]
				if strings.HasPrefix(serial, "stratux:978") {
					log.Println("Settings conflict: 1090ES set via WebUI but hardware serial says stratux:978")
				} else {
					ESDev = &ES{indexID: id}
					if err := ESDev.sdrConfig(); err != nil {
						log.Printf("ESDev = &ES{indexID: id} failed: %s\n", err)
						ESDev = nil
					} else {
						es_shutdown = make(chan int)
						es_wg.Add(1)
						go ESDev.read()
					}
				}
			}
		} else if ESDev != nil {
			ESDev.shutdown()
			ESDev = nil
			if count == 1 && UATDev != nil {
				UATDev.indexID = 0
			}
		}
	End:
	}
}

func sdrInit() {
	go sdrWatcher()
	go uatReader()
	godump978.Dump978Init()
	go godump978.ProcessDataFromChannel()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os"

	"github.com/squaremo/ambergreen/common/backends"
	"github.com/squaremo/ambergreen/common/data"

	"github.com/gorilla/mux"
)

func main() {
	etcd := os.Getenv("ETCD_ADDRESS")
	if etcd == "" {
		etcd = "http://localhost:4001"
	}

	prom := os.Getenv("PROM_ADDRESS")
	if prom == "" {
		prom = "http://localhost:9090"
	}

	back := backends.NewBackend([]string{etcd})
	if err := back.Ping(); err != nil {
		log.Fatal(err)
	}
	log.Printf("Connected to backend\n")
	api := &api{back, prom}

	router := mux.NewRouter()

	router.HandleFunc("/", homePage)
	router.HandleFunc("/index.html", homePage)
	router.PathPrefix("/res/").HandlerFunc(handleResource)

	router.HandleFunc("/api/{service}/", api.listInstances)
	router.HandleFunc("/api/", api.listServices)

	router.PathPrefix("/stats/").HandlerFunc(api.proxyStats)

	http.ListenAndServe("0.0.0.0:7070", router)
}

func handleResource(w http.ResponseWriter, r *http.Request) {
	file := r.URL.Path[1:]
	http.ServeFile(w, r, file)
}

func homePage(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "index.html")
}

//=== API handlers

type api struct {
	backend *backends.Backend
	promURL string
}

type serviceDetails struct {
	Name    string       `json:"name"`
	Details data.Service `json:"details"`
}

type service struct {
	Name     string            `json:"name"`
	Children []instanceDetails `json:"children"`
	Details  data.Service      `json:"details"`
}

type instanceDetails struct {
	Name    string        `json:"name"`
	Details data.Instance `json:"details"`
}

func (api *api) listServices(w http.ResponseWriter, r *http.Request) {
	var currentService serviceDetails
	services := []serviceDetails{}

	api.backend.ForeachServiceInstance(func(name string, details data.Service) {
		currentService = serviceDetails{
			Name:    name,
			Details: details,
		}
		services = append(services, currentService)
	}, func(_ string, _ data.Instance) {
	})
	json.NewEncoder(w).Encode(services)
}

func (api *api) listInstances(w http.ResponseWriter, r *http.Request) {
	args := mux.Vars(r)
	serviceName := args["service"]
	details, err := api.backend.GetServiceDetails(serviceName)
	if err != nil {
		http.NotFound(w, r)
	}
	children := []instanceDetails{}
	api.backend.ForeachInstance(serviceName, func(name string, details data.Instance) {
		instance := instanceDetails{
			Name:    name,
			Details: details,
		}
		children = append(children, instance)
	})
	service := service{
		Name:     serviceName,
		Details:  details,
		Children: children,
	}
	json.NewEncoder(w).Encode(service)
}

/* Proxy for prometheus, as a stop-gap */

func (api *api) proxyStats(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path[len("/stats"):] + "?" + r.URL.RawQuery
	log.Println(path)
	resp, err := http.Get(api.promURL + path)
	if err != nil {
		http.Error(w, "Error contacting prometheus server: "+err.Error(), 500)
		return
	}
	defer resp.Body.Close()
	resp.Header.Write(w)
	io.Copy(w, resp.Body)
}
<commit_msg>Write proxied header with ResponseWriter method<commit_after>package main

import (
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os"

	"github.com/squaremo/ambergreen/common/backends"
	"github.com/squaremo/ambergreen/common/data"

	"github.com/gorilla/mux"
)

func main() {
	etcd := os.Getenv("ETCD_ADDRESS")
	if etcd == "" {
		etcd = "http://localhost:4001"
	}

	prom := os.Getenv("PROM_ADDRESS")
	if prom == "" {
		prom = "http://localhost:9090"
	}

	back := backends.NewBackend([]string{etcd})
	if err := back.Ping(); err != nil {
		log.Fatal(err)
	}
	log.Printf("Connected to backend\n")
	api := &api{back, prom}

	router := mux.NewRouter()

	router.HandleFunc("/", homePage)
	router.HandleFunc("/index.html", homePage)
	router.PathPrefix("/res/").HandlerFunc(handleResource)

	router.HandleFunc("/api/{service}/", api.listInstances)
	router.HandleFunc("/api/", api.listServices)

	router.PathPrefix("/stats/").HandlerFunc(api.proxyStats)

	http.ListenAndServe("0.0.0.0:7070", router)
}

func handleResource(w http.ResponseWriter, r *http.Request) {
	file := r.URL.Path[1:]
	http.ServeFile(w, r, file)
}

func homePage(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "index.html")
}

//=== API handlers

type api struct {
	backend *backends.Backend
	promURL string
}

type serviceDetails struct {
	Name    string       `json:"name"`
	Details data.Service `json:"details"`
}

type service struct {
	Name     string            `json:"name"`
	Children []instanceDetails `json:"children"`
	Details  data.Service      `json:"details"`
}

type instanceDetails struct {
	Name    string        `json:"name"`
	Details data.Instance `json:"details"`
}

func (api *api) listServices(w http.ResponseWriter, r *http.Request) {
	var currentService serviceDetails
	services := []serviceDetails{}

	api.backend.ForeachServiceInstance(func(name string, details data.Service) {
		currentService = serviceDetails{
			Name:    name,
			Details: details,
		}
		services = append(services, currentService)
	}, func(_ string, _ data.Instance) {
	})
	json.NewEncoder(w).Encode(services)
}

func (api *api) listInstances(w http.ResponseWriter, r *http.Request) {
	args := mux.Vars(r)
	serviceName := args["service"]
	details, err := api.backend.GetServiceDetails(serviceName)
	if err != nil {
		http.NotFound(w, r)
	}
	children := []instanceDetails{}
	api.backend.ForeachInstance(serviceName, func(name string, details data.Instance) {
		instance := instanceDetails{
			Name:    name,
			Details: details,
		}
		children = append(children, instance)
	})
	service := service{
		Name:     serviceName,
		Details:  details,
		Children: children,
	}
	json.NewEncoder(w).Encode(service)
}

/* Proxy for prometheus, as a stop-gap */

func (api *api) proxyStats(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path[len("/stats"):] + "?" + r.URL.RawQuery
	log.Println(path)
	resp, err := http.Get(api.promURL + path)
	if err != nil {
		http.Error(w, "Error contacting prometheus server: "+err.Error(), 500)
		return
	}
	defer resp.Body.Close()
	for k, vs := range resp.Header {
		w.Header()[k] = vs
	}
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}
<|endoftext|>


<commit_before>/*
*  Copyright (c) 2012 by Christoph Hack <christoph@tux21b.org>
*                        Stefan Lendl <st.lendl@gmail.com>
*  All rights reserved. Distributed under the Simplified BSD License.
 */

package main

import (
	"fmt"
	irc "github.com/fluffle/goirc/client"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"time"
)

var (
	timezone          = time.UTC
	wolframURL        = "http://www.wolframalpha.com/input/?i=%s"
	mathbinPreviewURL = "http://mathbin.net/preview.cgi?body=%s"
	mathbinURL        = "http://mathbin.net/%s"
	reMathbin         = regexp.MustCompile(`equation_previews/[\d_]+\.png`)
)

func init() {
	tz, err := time.LoadLocation("Europe/Vienna")
	if err != nil {
		log.Fatal(err)
	}
	timezone = tz
}

var tmpl = template.Must(template.New("weierbot").Parse(`<!doctype html>

<title>{{.Title}} - Weierbot</title>

{{if .Prev}}Previous{{end}}
{{if .Next}}Next{{end}}

Weierbot

{{.Title}}

{{.Chatlog}}

`))

var reURL = regexp.MustCompile(`https?://[-A-Za-z0-9+&@#/%?=~_|!:,.;]*[-A-Za-z0-9+&@#/%=~_|]`)
var reTime = regexp.MustCompile(`(?m)^\[\S+\]`)

type WeierBot struct {
	server     string
	client     *irc.Conn
	log        chan Message
	disconnect chan bool
}

func NewWeierBot(server, nick, password string, channels []string) *WeierBot {
	w := &WeierBot{
		server:     server,
		client:     irc.SimpleClient(nick),
		log:        make(chan Message, 512),
		disconnect: make(chan bool),
	}
	w.client.Me.Name = nick
	w.client.Me.Ident = nick

	w.client.AddHandler("connected", func(conn *irc.Conn, line *irc.Line) {
		conn.Privmsg("NickServ", "identify "+password)
		for _, ch := range channels {
			log.Printf("Joining %s\n", ch)
			conn.Join(ch)
		}
	})
	w.client.AddHandler("disconnected", func(conn *irc.Conn, line *irc.Line) {
		w.disconnect <- true
	})
	w.client.AddHandler("join", func(conn *irc.Conn, line *irc.Line) {
		if len(line.Args) >= 1 && strings.HasPrefix(line.Args[0], "#") {
			conn.Privmsg(line.Nick, fmt.Sprintf("Hallo %s, willkommen auf %s. "+
				"Dieser Channel wird unter http://weierbot.tux21b.org/ mitgespeichert.",
				line.Nick, line.Args[0]))
		}
	})
	w.client.AddHandler("privmsg", func(conn *irc.Conn, line *irc.Line) {
		if len(line.Args) < 2 {
			return
		}
		w.handleMessage(line.Args[0], Message{
			Nick:    line.Nick,
			Time:    line.Time,
			Message: line.Args[1],
		})
	})

	go w.writeLog()

	return w
}

func createMathbinImage(latexcode string) string {
	resp, err := http.Get(fmt.Sprintf(mathbinPreviewURL, url.QueryEscape("[EQ]"+latexcode+"[/EQ]")))
	if err != nil {
		return "error"
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)

	result := reMathbin.FindString(string(body))

	return fmt.Sprintf(mathbinURL, result)
}

func (bot *WeierBot) handleMessage(target string, msg Message) {
	if strings.HasPrefix(target, "#") {
		bot.log <- msg
	} else {
		target = msg.Nick
	}

	switch {
	case strings.HasPrefix(msg.Message, "!wolfram "):
		bot.send(target, fmt.Sprintf(wolframURL, url.QueryEscape(msg.Message[9:])))
	case msg.Message == "!log":
		bot.send(target, "http://weierbot.tux21b.org/")
	case strings.HasPrefix(msg.Message, "!mathbin "):
		imgurl := createMathbinImage(msg.Message[9:])
		bot.send(target, fmt.Sprintf(imgurl))
	}
}

func (bot *WeierBot) send(target, message string) {
	bot.client.Privmsg(target, message)
	if strings.HasPrefix(target, "#") {
		bot.log <- Message{Nick: bot.client.Me.Nick, Time: time.Now(),
			Message: message}
	}
}

func (bot *WeierBot) writeLog() {
	var (
		file     *os.File
		format   = "2006-01-02.log"
		filename string
	)
	defer func() {
		if file != nil {
			file.Close()
		}
	}()
	for msg := range bot.log {
		msg.Time = msg.Time.In(timezone)
		current := msg.Time.Format(format)
		if file == nil || current != filename {
			if file != nil {
				file.Close()
			}
			f, err := os.OpenFile(current,
				os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
			if err != nil {
				log.Println(err)
			} else {
				file, filename = f, current
			}
		}
		fmt.Fprintf(file, "%s\n", msg)
	}
}

func (bot *WeierBot) ServeIRC() {
	tries := 0
	for {
		tries++
		if err := bot.client.Connect(bot.server); err != nil {
			log.Println(err)
			if tries >= 5 {
				log.Fatal("maximum number of tries exceeded")
			}
			<-time.After(5 * time.Second)
			continue
		}
		tries = 0
		log.Printf("Connected to %s\n", bot.server)
		<-bot.disconnect
	}
}

func (bot *WeierBot) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	files, err := filepath.Glob("*.log")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	sort.Strings(files)

	if r.URL.Path == "/" && len(files) > 0 {
		http.Redirect(w, r, "/"+filepath.Base(files[len(files)-1]),
			http.StatusTemporaryRedirect)
	}
	index := sort.SearchStrings(files, r.URL.Path[1:])
	if index < 0 || index >= len(files) {
		http.NotFound(w, r)
		return
	}

	chatlog, err := ioutil.ReadFile(files[index])
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	chatlog = []byte(template.HTMLEscapeString(string(chatlog)))
	chatlog = reURL.ReplaceAllFunc(chatlog, func(m []byte) []byte {
		return []byte(fmt.Sprintf("%s", m, m))
	})
	chatlog = reTime.ReplaceAllFunc(chatlog, func(m []byte) []byte {
		return []byte(fmt.Sprintf("%s", m))
	})

	data := struct {
		Title      string
		Prev, Next string
		Chatlog    template.HTML
	}{
		Title:   filepath.Base(files[index]),
		Chatlog: template.HTML(chatlog),
	}
	if index > 0 {
		data.Prev = files[index-1]
	}
	if index < len(files)-1 {
		data.Next = files[index+1]
	}
	if err := tmpl.Execute(w, data); err != nil {
		log.Println(err.Error())
	}
}

type Message struct {
	Time    time.Time
	Nick    string
	Message string
}

func (m Message) String() string {
	return fmt.Sprintf("[%s] %s: %s", m.Time.Format(time.Kitchen), m.Nick, m.Message)
}

func main() {
	bot := NewWeierBot("irc.euirc.net", "weierbot2", "secret", []string{"#tm-test"})

	go func() {
		if err := http.ListenAndServe(":8005", bot); err != nil {
			log.Fatal(err)
		}
	}()

	bot.ServeIRC()
}
<commit_msg>output correct error<commit_after>/*
*  Copyright (c) 2012 by Christoph Hack <christoph@tux21b.org>
*                        Stefan Lendl <st.lendl@gmail.com>
*  All rights reserved. Distributed under the Simplified BSD License.
 */

package main

import (
	"fmt"
	irc "github.com/fluffle/goirc/client"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"time"
)

var (
	timezone          = time.UTC
	wolframURL        = "http://www.wolframalpha.com/input/?i=%s"
	mathbinPreviewURL = "http://mathbin.net/preview.cgi?body=%s"
	mathbinURL        = "http://mathbin.net/%s"
	reMathbin         = regexp.MustCompile(`equation_previews/[\d_]+\.png`)
)

func init() {
	tz, err := time.LoadLocation("Europe/Vienna")
	if err != nil {
		log.Fatal(err)
	}
	timezone = tz
}

var tmpl = template.Must(template.New("weierbot").Parse(`<!doctype html>

<title>{{.Title}} - Weierbot</title>

{{if .Prev}}Previous{{end}}
{{if .Next}}Next{{end}}

Weierbot

{{.Title}}

{{.Chatlog}}

`))

var reURL = regexp.MustCompile(`https?://[-A-Za-z0-9+&@#/%?=~_|!:,.;]*[-A-Za-z0-9+&@#/%=~_|]`)
var reTime = regexp.MustCompile(`(?m)^\[\S+\]`)

type WeierBot struct {
	server     string
	client     *irc.Conn
	log        chan Message
	disconnect chan bool
}

func NewWeierBot(server, nick, password string, channels []string) *WeierBot {
	w := &WeierBot{
		server:     server,
		client:     irc.SimpleClient(nick),
		log:        make(chan Message, 512),
		disconnect: make(chan bool),
	}
	w.client.Me.Name = nick
	w.client.Me.Ident = nick

	w.client.AddHandler("connected", func(conn *irc.Conn, line *irc.Line) {
		conn.Privmsg("NickServ", "identify "+password)
		for _, ch := range channels {
			log.Printf("Joining %s\n", ch)
			conn.Join(ch)
		}
	})
	w.client.AddHandler("disconnected", func(conn *irc.Conn, line *irc.Line) {
		w.disconnect <- true
	})
	w.client.AddHandler("join", func(conn *irc.Conn, line *irc.Line) {
		if len(line.Args) >= 1 && strings.HasPrefix(line.Args[0], "#") {
			conn.Privmsg(line.Nick, fmt.Sprintf("Hallo %s, willkommen auf %s. "+
				"Dieser Channel wird unter http://weierbot.tux21b.org/ mitgespeichert.",
				line.Nick, line.Args[0]))
		}
	})
	w.client.AddHandler("privmsg", func(conn *irc.Conn, line *irc.Line) {
		if len(line.Args) < 2 {
			return
		}
		w.handleMessage(line.Args[0], Message{
			Nick:    line.Nick,
			Time:    line.Time,
			Message: line.Args[1],
		})
	})

	go w.writeLog()

	return w
}

func createMathbinImage(latexcode string) string {
	resp, err := http.Get(fmt.Sprintf(mathbinPreviewURL, url.QueryEscape("[EQ]"+latexcode+"[/EQ]")))
	if err != nil {
		return err.Error()
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)

	result := reMathbin.FindString(string(body))

	return fmt.Sprintf(mathbinURL, result)
}

func (bot *WeierBot) handleMessage(target string, msg Message) {
	if strings.HasPrefix(target, "#") {
		bot.log <- msg
	} else {
		target = msg.Nick
	}

	switch {
	case strings.HasPrefix(msg.Message, "!wolfram "):
		bot.send(target, fmt.Sprintf(wolframURL, url.QueryEscape(msg.Message[9:])))
	case msg.Message == "!log":
		bot.send(target, "http://weierbot.tux21b.org/")
	case strings.HasPrefix(msg.Message, "!mathbin "):
		imgurl := createMathbinImage(msg.Message[9:])
		bot.send(target, fmt.Sprintf(imgurl))
	}
}

func (bot *WeierBot) send(target, message string) {
	bot.client.Privmsg(target, message)
	if strings.HasPrefix(target, "#") {
		bot.log <- Message{Nick: bot.client.Me.Nick, Time: time.Now(),
			Message: message}
	}
}

func (bot *WeierBot) writeLog() {
	var (
		file     *os.File
		format   = "2006-01-02.log"
		filename string
	)
	defer func() {
		if file != nil {
			file.Close()
		}
	}()
	for msg := range bot.log {
		msg.Time = msg.Time.In(timezone)
		current := msg.Time.Format(format)
		if file == nil || current != filename {
			if file != nil {
				file.Close()
			}
			f, err := os.OpenFile(current,
				os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
			if err != nil {
				log.Println(err)
			} else {
				file, filename = f, current
			}
		}
		fmt.Fprintf(file, "%s\n", msg)
	}
}

func (bot *WeierBot) ServeIRC() {
	tries := 0
	for {
		tries++
		if err := bot.client.Connect(bot.server); err != nil {
			log.Println(err)
			if tries >= 5 {
				log.Fatal("maximum number of tries exceeded")
			}
			<-time.After(5 * time.Second)
			continue
		}
		tries = 0
		log.Printf("Connected to %s\n", bot.server)
		<-bot.disconnect
	}
}

func (bot *WeierBot) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	files, err := filepath.Glob("*.log")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	sort.Strings(files)

	if r.URL.Path == "/" && len(files) > 0 {
		http.Redirect(w, r, "/"+filepath.Base(files[len(files)-1]),
			http.StatusTemporaryRedirect)
	}
	index := sort.SearchStrings(files, r.URL.Path[1:])
	if index < 0 || index >= len(files) {
		http.NotFound(w, r)
		return
	}

	chatlog, err := ioutil.ReadFile(files[index])
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	chatlog = []byte(template.HTMLEscapeString(string(chatlog)))
	chatlog = reURL.ReplaceAllFunc(chatlog, func(m []byte) []byte {
		return []byte(fmt.Sprintf("%s", m, m))
	})
	chatlog = reTime.ReplaceAllFunc(chatlog, func(m []byte) []byte {
		return []byte(fmt.Sprintf("%s", m))
	})

	data := struct {
		Title      string
		Prev, Next string
		Chatlog    template.HTML
	}{
		Title:   filepath.Base(files[index]),
		Chatlog: template.HTML(chatlog),
	}
	if index > 0 {
		data.Prev = files[index-1]
	}
	if index < len(files)-1 {
		data.Next = files[index+1]
	}
	if err := tmpl.Execute(w, data); err != nil {
		log.Println(err.Error())
	}
}

type Message struct {
	Time    time.Time
	Nick    string
	Message string
}

func (m Message) String() string {
	return fmt.Sprintf("[%s] %s: %s", m.Time.Format(time.Kitchen), m.Nick, m.Message)
}

func main() {
	bot := NewWeierBot("irc.euirc.net", "weierbot2", "secret", []string{"#tm-test"})

	go func() {
		if err := http.ListenAndServe(":8005", bot); err != nil {
			log.Fatal(err)
		}
	}()

	bot.ServeIRC()
}
<|endoftext|>


<commit_before>// Copyright 2014 <chaishushan{AT}gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package image

import (
	"image"
	"image/color"
	"image/draw"
	"reflect"
)

var (
	_ Image = (*Gray)(nil)
	_ Image = (*Gray16)(nil)
	_ Image = (*Gray32f)(nil)
	_ Image = (*GrayA)(nil)
	_ Image = (*GrayA32)(nil)
	_ Image = (*GrayA64f)(nil)
	_ Image = (*RGB)(nil)
	_ Image = (*RGB48)(nil)
	_ Image = (*RGB96f)(nil)
	_ Image = (*RGBA)(nil)
	_ Image = (*RGBA64)(nil)
	_ Image = (*RGBA128f)(nil)
	_ Image = (*Unknown)(nil)
)

type Image interface {
	// Get original type, such as *image.Gray, *image.RGBA, etc.
	BaseType() image.Image

	// Pix holds the image's pixels, as pixel values in big-endian format. The pixel at
	// (x, y) starts at Pix[(y-Rect.Min.Y)*Stride + (x-Rect.Min.X)*Channels*sizeof(DataType)].
	Pix() []byte
	// Stride is the Pix stride (in bytes) between vertically adjacent pixels.
	Stride() int
	// Rect is the image's bounds.
	Rect() image.Rectangle

	// 1:Gray, 2:GrayA, 3:RGB, 4:RGBA
	Channels() int
	// Uint8/Uint16/Int32/Int64/Float32/Float64
	Depth() reflect.Kind

	draw.Image
}

func newRGBA64FromImage(m image.Image) *RGBA64 {
	b := m.Bounds()
	rgba64 := NewRGBA64(b)
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			pr, pg, pb, pa := m.At(x, y).RGBA()
			rgba64.SetRGBA64(x, y, color.RGBA64{
				R: uint16(pr),
				G: uint16(pg),
				B: uint16(pb),
				A: uint16(pa),
			})
		}
	}
	return rgba64
}

func AsImage(m image.Image) Image {
	if p, ok := m.(Image); ok {
		return p
	}

	switch m := m.(type) {
	case *image.Gray:
		return &Gray{m}
	case *image.Gray16:
		return &Gray16{m}
	case *image.RGBA:
		return &RGBA{m}
	case *image.RGBA64:
		return &RGBA64{m}
	}

	return newRGBA64FromImage(m)
}

func CloneImage(m image.Image) Image {
	if m, ok := m.(Image); ok {
		switch channels, depth := m.Channels(), m.Depth(); {
		case channels == 1 && depth == reflect.Uint8:
			return new(Gray).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 1 && depth == reflect.Uint16:
			return new(Gray16).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 1 && depth == reflect.Float32:
			return new(Gray32f).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 2 && depth == reflect.Uint8:
			return new(GrayA).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 2 && depth == reflect.Uint16:
			return new(GrayA32).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 2 && depth == reflect.Float32:
			return new(GrayA64f).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 3 && depth == reflect.Uint8:
			return new(RGB).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 3 && depth == reflect.Uint16:
			return new(RGB48).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 3 && depth == reflect.Float32:
			return new(RGB96f).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 4 && depth == reflect.Uint8:
			return new(RGBA).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 4 && depth == reflect.Uint16:
			return new(RGBA64).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 4 && depth == reflect.Float32:
			return new(RGBA128f).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		}

		return new(Unknown).Init(
			append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect(),
			m.Channels(), m.Depth(),
		)
	}

	switch m := m.(type) {
	case *image.Gray:
		return new(Gray).Init(append([]uint8(nil), m.Pix...), m.Stride, m.Rect)
	case *image.Gray16:
		return new(Gray16).Init(append([]uint8(nil), m.Pix...), m.Stride, m.Rect)
	case *image.RGBA:
		return new(RGBA).Init(append([]uint8(nil), m.Pix...), m.Stride, m.Rect)
	case *image.RGBA64:
		return new(RGBA64).Init(append([]uint8(nil), m.Pix...), m.Stride, m.Rect)
	}

	return newRGBA64FromImage(m)
}

func CopyImage(buf, src image.Image) (dst Image) {
	if buf == nil {
		return CloneImage(src)
	}

	switch buf := buf.(type) {
	case *image.Gray:
		return (&Gray{buf}).CopyFrom(src)
	case *image.Gray16:
		return (&Gray16{buf}).CopyFrom(src)
	case *image.RGBA:
		return (&RGBA{buf}).CopyFrom(src)
	case *image.RGBA64:
		return (&RGBA64{buf}).CopyFrom(src)

	case *Gray:
		return buf.CopyFrom(src)
	case *Gray16:
		return buf.CopyFrom(src)
	case *Gray32f:
		return buf.CopyFrom(src)

	case *GrayA:
		return buf.CopyFrom(src)
	case *GrayA32:
		return buf.CopyFrom(src)
	case *GrayA64f:
		return buf.CopyFrom(src)

	case *RGB:
		return buf.CopyFrom(src)
	case *RGB48:
		return buf.CopyFrom(src)
	case *RGB96f:
		return buf.CopyFrom(src)

	case *RGBA:
		return buf.CopyFrom(src)
	case *RGBA64:
		return buf.CopyFrom(src)
	case *RGBA128f:
		return buf.CopyFrom(src)

	case *Unknown:
		return buf.CopyFrom(src)
	}

	return CloneImage(src)
}

func ConvertCopyImage(m image.Image, channels int, depth reflect.Kind) (Image, error) {
	buf, err := NewImage(m.Bounds(), channels, depth)
	if err != nil {
		return nil, err
	}
	p := CopyImage(buf, m)
	return p, nil
}

func ConvertImage(m image.Image, channels int, depth reflect.Kind) (Image, error) {
	if p, ok := m.(Image); ok {
		if channels == p.Channels() && depth == p.Depth() {
			return p, nil
		}
	}
	buf, err := NewImage(m.Bounds(), channels, depth)
	if err != nil {
		return nil, err
	}
	p := CopyImage(buf, m)
	return p, nil
}

func NewImage(r image.Rectangle, channels int, depth reflect.Kind) (m Image, err error) {
	switch {
	case channels == 1 && depth == reflect.Uint8:
		m = NewGray(r)
		return
	case channels == 1 && depth == reflect.Uint16:
		m = NewGray16(r)
		return
	case channels == 1 && depth == reflect.Float32:
		m = NewGray32f(r)
		return

	case channels == 2 && depth == reflect.Uint8:
		m = NewGrayA(r)
		return
	case channels == 2 && depth == reflect.Uint16:
		m = NewGrayA32(r)
		return
	case channels == 2 && depth == reflect.Float32:
		m = NewGrayA64f(r)
		return

	case channels == 3 && depth == reflect.Uint8:
		m = NewRGB(r)
		return
	case channels == 3 && depth == reflect.Uint16:
		m = NewRGB48(r)
		return
	case channels == 3 && depth == reflect.Float32:
		m = NewRGB96f(r)
		return

	case channels == 4 && depth == reflect.Uint8:
		m = NewRGBA(r)
		return
	case channels == 4 && depth == reflect.Uint16:
		m = NewRGBA64(r)
		return
	case channels == 4 && depth == reflect.Float32:
		m = NewRGBA128f(r)
		return

	default:
		m, err = NewUnknown(r, channels, depth)
		return
	}
}
<commit_msg>image: change api, use machine order instead of big-endian<commit_after>// Copyright 2014 <chaishushan{AT}gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package image

import (
	"image"
	"image/color"
	"image/draw"
	"reflect"
)

var (
	_ Image = (*Gray)(nil)
	_ Image = (*Gray16)(nil)
	_ Image = (*Gray32f)(nil)
	_ Image = (*GrayA)(nil)
	_ Image = (*GrayA32)(nil)
	_ Image = (*GrayA64f)(nil)
	_ Image = (*RGB)(nil)
	_ Image = (*RGB48)(nil)
	_ Image = (*RGB96f)(nil)
	_ Image = (*RGBA)(nil)
	_ Image = (*RGBA64)(nil)
	_ Image = (*RGBA128f)(nil)
	_ Image = (*Unknown)(nil)
)

type Image interface {
	// Get original type, such as *image.Gray, *image.RGBA, etc.
	BaseType() image.Image

	// Pix holds the image's pixels, as pixel values in machine order format. The pixel at
	// (x, y) starts at []Type(Pix[(y-Rect.Min.Y)*Stride:])[(x-Rect.Min.X)*Channels:].
	Pix() []byte
	// Stride is the Pix stride (in bytes) between vertically adjacent pixels.
	Stride() int
	// Rect is the image's bounds.
	Rect() image.Rectangle

	// 1:Gray, 2:GrayA, 3:RGB, 4:RGBA
	Channels() int
	// Uint8/Uint16/Int32/Int64/Float32/Float64
	Depth() reflect.Kind

	draw.Image
}

func newRGBA64FromImage(m image.Image) *RGBA64 {
	b := m.Bounds()
	rgba64 := NewRGBA64(b)
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			pr, pg, pb, pa := m.At(x, y).RGBA()
			rgba64.SetRGBA64(x, y, color.RGBA64{
				R: uint16(pr),
				G: uint16(pg),
				B: uint16(pb),
				A: uint16(pa),
			})
		}
	}
	return rgba64
}

func AsImage(m image.Image) Image {
	if p, ok := m.(Image); ok {
		return p
	}

	switch m := m.(type) {
	case *image.Gray:
		return &Gray{m}
	case *image.Gray16:
		return &Gray16{m}
	case *image.RGBA:
		return &RGBA{m}
	case *image.RGBA64:
		return &RGBA64{m}
	}

	return newRGBA64FromImage(m)
}

func CloneImage(m image.Image) Image {
	if m, ok := m.(Image); ok {
		switch channels, depth := m.Channels(), m.Depth(); {
		case channels == 1 && depth == reflect.Uint8:
			return new(Gray).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 1 && depth == reflect.Uint16:
			return new(Gray16).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 1 && depth == reflect.Float32:
			return new(Gray32f).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 2 && depth == reflect.Uint8:
			return new(GrayA).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 2 && depth == reflect.Uint16:
			return new(GrayA32).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 2 && depth == reflect.Float32:
			return new(GrayA64f).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 3 && depth == reflect.Uint8:
			return new(RGB).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 3 && depth == reflect.Uint16:
			return new(RGB48).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 3 && depth == reflect.Float32:
			return new(RGB96f).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 4 && depth == reflect.Uint8:
			return new(RGBA).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 4 && depth == reflect.Uint16:
			return new(RGBA64).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		case channels == 4 && depth == reflect.Float32:
			return new(RGBA128f).Init(append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect())
		}

		return new(Unknown).Init(
			append([]uint8(nil), m.Pix()...), m.Stride(), m.Rect(),
			m.Channels(), m.Depth(),
		)
	}

	switch m := m.(type) {
	case *image.Gray:
		return new(Gray).Init(append([]uint8(nil), m.Pix...), m.Stride, m.Rect)
	case *image.Gray16:
		return new(Gray16).Init(append([]uint8(nil), m.Pix...), m.Stride, m.Rect)
	case *image.RGBA:
		return new(RGBA).Init(append([]uint8(nil), m.Pix...), m.Stride, m.Rect)
	case *image.RGBA64:
		return new(RGBA64).Init(append([]uint8(nil), m.Pix...), m.Stride, m.Rect)
	}

	return newRGBA64FromImage(m)
}

func CopyImage(buf, src image.Image) (dst Image) {
	if buf == nil {
		return CloneImage(src)
	}

	switch buf := buf.(type) {
	case *image.Gray:
		return (&Gray{buf}).CopyFrom(src)
	case *image.Gray16:
		return (&Gray16{buf}).CopyFrom(src)
	case *image.RGBA:
		return (&RGBA{buf}).CopyFrom(src)
	case *image.RGBA64:
		return (&RGBA64{buf}).CopyFrom(src)

	case *Gray:
		return buf.CopyFrom(src)
	case *Gray16:
		return buf.CopyFrom(src)
	case *Gray32f:
		return buf.CopyFrom(src)

	case *GrayA:
		return buf.CopyFrom(src)
	case *GrayA32:
		return buf.CopyFrom(src)
	case *GrayA64f:
		return buf.CopyFrom(src)

	case *RGB:
		return buf.CopyFrom(src)
	case *RGB48:
		return buf.CopyFrom(src)
	case *RGB96f:
		return buf.CopyFrom(src)

	case *RGBA:
		return buf.CopyFrom(src)
	case *RGBA64:
		return buf.CopyFrom(src)
	case *RGBA128f:
		return buf.CopyFrom(src)

	case *Unknown:
		return buf.CopyFrom(src)
	}

	return CloneImage(src)
}

func ConvertCopyImage(m image.Image, channels int, depth reflect.Kind) (Image, error) {
	buf, err := NewImage(m.Bounds(), channels, depth)
	if err != nil {
		return nil, err
	}
	p := CopyImage(buf, m)
	return p, nil
}

func ConvertImage(m image.Image, channels int, depth reflect.Kind) (Image, error) {
	if p, ok := m.(Image); ok {
		if channels == p.Channels() && depth == p.Depth() {
			return p, nil
		}
	}
	buf, err := NewImage(m.Bounds(), channels, depth)
	if err != nil {
		return nil, err
	}
	p := CopyImage(buf, m)
	return p, nil
}

func NewImage(r image.Rectangle, channels int, depth reflect.Kind) (m Image, err error) {
	switch {
	case channels == 1 && depth == reflect.Uint8:
		m = NewGray(r)
		return
	case channels == 1 && depth == reflect.Uint16:
		m = NewGray16(r)
		return
	case channels == 1 && depth == reflect.Float32:
		m = NewGray32f(r)
		return

	case channels == 2 && depth == reflect.Uint8:
		m = NewGrayA(r)
		return
	case channels == 2 && depth == reflect.Uint16:
		m = NewGrayA32(r)
		return
	case channels == 2 && depth == reflect.Float32:
		m = NewGrayA64f(r)
		return

	case channels == 3 && depth == reflect.Uint8:
		m = NewRGB(r)
		return
	case channels == 3 && depth == reflect.Uint16:
		m = NewRGB48(r)
		return
	case channels == 3 && depth == reflect.Float32:
		m = NewRGB96f(r)
		return

	case channels == 4 && depth == reflect.Uint8:
		m = NewRGBA(r)
		return
	case channels == 4 && depth == reflect.Uint16:
		m = NewRGBA64(r)
		return
	case channels == 4 && depth == reflect.Float32:
		m = NewRGBA128f(r)
		return

	default:
		m, err = NewUnknown(r, channels, depth)
		return
	}
}
<|endoftext|>


<commit_before>// Copyright 2011-2014 Dmitry Chestnykh. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package captcha

import (
	"bytes"
	"image"
	"image/color"
	"image/png"
	"io"
	"math"
)

const (
	// Standard width and height of a captcha image.
	StdWidth  = 240
	StdHeight = 80
	// Maximum absolute skew factor of a single digit.
	maxSkew = 0.7
	// Number of background circles.
	circleCount = 20
	distortCircleCount = 5
)

type Image struct {
	*image.RGBA
	numWidth  int
	numHeight int
	dotSize   int
	rng       siprng
}

var background color.Color = color.RGBA{144, 144, 144, 0xFF}

// NewImage returns a new captcha image of the given width and height with the
// given digits, where each digit must be in range 0-9.
func NewImage(id string, digits []byte, width, height int) *Image {
	m := new(Image)

	// Initialize PRNG.
	m.rng.Seed(deriveSeed(imageSeedPurpose, id, digits))

	prim := color.RGBA{
		uint8(m.rng.Intn(129)),
		uint8(m.rng.Intn(129)),
		uint8(m.rng.Intn(129)),
		0xFF,
	}

	m.RGBA = image.NewRGBA(image.Rect(0, 0, width, height))
	m.calculateSizes(width, height, len(digits))
	// Randomly position captcha inside the image.
	maxx := width - (m.numWidth+m.dotSize)*len(digits) - m.dotSize
	maxy := height - m.numHeight - m.dotSize*2
	var border int
	if width > height {
		border = height / 10
	} else {
		border = width / 10
	}
	x := m.rng.Int(border, maxx-border)
	y := m.rng.Int(border, maxy-border)
	// Draw digits.
	for _, n := range digits {
		m.drawDigit(getChar(n), x, y, prim)
		x += m.numWidth + m.dotSize
	}
	// Draw strike-through line.
	m.strikeThrough(prim)

	m.fillWithDistortCircles(distortCircleCount, width/4)
	// Apply wave distortion.
	m.distort(m.rng.Float(5, 10), m.rng.Float(100, 200))
	// Fill image with random circles.
	m.fillWithCircles(circleCount, m.dotSize)
	return m
}

/*
func (m *Image) getRandomPalette() color.Palette {
	p := make([]color.Color, circleCount+1)
	// Transparent color.
	p[0] = color.RGBA{0xFF, 0xFF, 0xFF, 0x00}
	// Primary color.
	prim := color.RGBA{
		uint8(m.rng.Intn(129)),
		uint8(m.rng.Intn(129)),
		uint8(m.rng.Intn(129)),
		0xFF,
	}
	p[1] = prim
	// Circle colors.
	for i := 2; i <= circleCount; i++ {
		p[i] = m.randomBrightness(prim, 255)
	}
	return p
}*/

// encodedPNG encodes an image to PNG and returns
// the result as a byte slice.
func (m *Image) encodedPNG() []byte {
	var buf bytes.Buffer
	if err := png.Encode(&buf, m); err != nil {
		panic(err.Error())
	}
	return buf.Bytes()
}

// WriteTo writes captcha image in PNG format into the given writer.
func (m *Image) WriteTo(w io.Writer) (int64, error) {
	n, err := w.Write(m.encodedPNG())
	return int64(n), err
}

func (m *Image) calculateSizes(width, height, ncount int) {
	// Goal: fit all digits inside the image.
	var border int
	if width > height {
		border = height / 8
	} else {
		border = width / 8
	}
	// Convert everything to floats for calculations.
	w := float64(width - border*2)
	h := float64(height - border*2)
	// fw takes into account 1-dot spacing between digits.
	fw := float64(fontWidth + 1)
	fh := float64(fontHeight)
	nc := float64(ncount)
	// Calculate the width of a single digit taking into account only the
	// width of the image.
	nw := w / nc
	// Calculate the height of a digit from this width.
	nh := nw * fh / fw
	// Digit too high?
	if nh > h {
		// Fit digits based on height.
		nh = h
		nw = fw / fh * nh
	}
	// Calculate dot size.
	m.dotSize = int((nh / fh))
	if m.dotSize < 1 {
		m.dotSize = 1
	}
	// Save everything, making the actual width smaller by 1 dot to account
	// for spacing between digits.
	m.numWidth = int(nw) - m.dotSize
	m.numHeight = int(nh)
}

func (m *Image) drawHorizLine(fromX, toX, y int, color color.Color) {
	for x := fromX; x <= toX; x++ {
		m.Set(x, y, color)
	}
}

func (m *Image) drawCircle(x, y, radius int, color color.Color) {
	f := 1 - radius
	dfx := 1
	dfy := -2 * radius
	xo := 0
	yo := radius

	m.Set(x, y+radius, color)
	m.Set(x, y-radius, color)
	m.drawHorizLine(x-radius, x+radius, y, color)

	for xo < yo {
		if f >= 0 {
			yo--
			dfy += 2
			f += dfy
		}
		xo++
		dfx += 2
		f += dfx
		m.drawHorizLine(x-xo, x+xo, y+yo, color)
		m.drawHorizLine(x-xo, x+xo, y-yo, color)
		m.drawHorizLine(x-yo, x+yo, y+xo, color)
		m.drawHorizLine(x-yo, x+yo, y-xo, color)
	}
}

func (m *Image) colorDistortCircle(xc, yc, r int) {
	d := 2*r
	r2 := r*r

	dg := int(m.rng.Intn(240)) - 120
	if dg < 0 {
		dg -= 20
	} else {
		dg += 20
	}

	dr := int(m.rng.Intn(240)) - 120
	if dr < 0 {
		dr -= 20
	} else {
		dr += 20
	}

	db := int(m.rng.Intn(240)) - 120
	if db < 0 {
		db -= 20
	} else {
		db += 20
	}

	for y := 0; y <= d; y++ {
		for x := 0; x <= d; x++ {
			x2 := (x-r)*(x-r)
			y2 := (y-r)*(y-r)

			if x2+y2 <= r2 {
				xp := xc - r + x
				yp := yc - r + y
				r,g,b,a := m.At(xp, yp).RGBA()
				if a != 0x00 { //background
					m.Set(xp, yp, color.RGBA{
						uint8(int(r)+dr),
						uint8(int(g)+dg),
						uint8(int(b)+db),
						0xFF,
					})
				}
			}
		}
	}
}

func (m *Image) fillWithCircles(n, maxradius int) {
	maxx := m.Bounds().Max.X
	maxy := m.Bounds().Max.Y
	for i := 0; i < n; i++ {
		color := color.RGBA{
			uint8(m.rng.Intn(129)),
			uint8(m.rng.Intn(129)),
			uint8(m.rng.Intn(129)),
			0xFF,
		}
		r := m.rng.Int(1, maxradius)
		m.drawCircle(m.rng.Int(r, maxx-r), m.rng.Int(r, maxy-r), r, color)
	}
}

func (m *Image) fillWithDistortCircles(n, maxradius int) {
	maxx := m.Bounds().Max.X
	maxy := m.Bounds().Max.Y
	for i := 0; i < n; i++ {
		r := m.rng.Int(maxradius/10, maxradius-1)
		m.colorDistortCircle(m.rng.Int(0, maxx), m.rng.Int(0, maxy), r)
	}
}

func (m *Image) strikeThrough(color color.Color) {
	maxx := m.Bounds().Max.X
	maxy := m.Bounds().Max.Y
	y := m.rng.Int(maxy/3, maxy-maxy/3)
	amplitude := m.rng.Float(5, 20)
	period := m.rng.Float(80, 180)
	dx := 2.0 * math.Pi / period
	for x := 0; x < maxx; x++ {
		xo := amplitude * math.Cos(float64(y)*dx)
		yo := amplitude * math.Sin(float64(x)*dx)
		for yn := 0; yn < m.dotSize; yn++ {
			r := m.rng.Int(0, m.dotSize)
			m.drawCircle(x+int(xo), y+int(yo)+(yn*m.dotSize), r/2, color)
		}
	}
}

func (m *Image) drawDigit(digit []byte, x, y int, color color.Color) {
	skf := m.rng.Float(-maxSkew, maxSkew)
	xs := float64(x)
	r := m.dotSize / 2
	y += m.rng.Int(-r, r)
	for yo := 0; yo < fontHeight; yo++ {
		for xo := 0; xo < fontWidth; xo++ {
			if digit[yo*fontWidth+xo] != blackChar {
				continue
			}
			m.drawCircle(x+xo*m.dotSize, y+yo*m.dotSize, r, color)
		}
		xs += skf
		x = int(xs)
	}
}

func (m *Image) distort(amplude float64, period float64) {
	w := m.Bounds().Max.X
	h := m.Bounds().Max.Y
	r32,g32,b32,_ := background.RGBA()
	r:=uint8(r32)
	g:=uint8(g32)
	b:=uint8(b32)
	dbg := 32.0/float32(w);

	oldm := m
	newm := image.NewRGBA(image.Rect(0, 0, w, h))

	dx := 2.0 * math.Pi / period
	for x := 0; x < w; x++ {
		dbgx:=uint8(float32(x)*dbg)
		bg:=color.RGBA{r-dbgx, g-dbgx, b-dbgx, 0xFF}
		for y := 0; y < h; y++ {
			xo := amplude * math.Sin(float64(y)*dx)
			yo := amplude * math.Cos(float64(x)*dx)
			c := oldm.At(x+int(xo), y+int(yo))
			newm.Set(x, y, c)
			_,_,_, a := c.RGBA()
			if a == 0x00 {
				newm.Set(x, y, bg)
			}
		}
	}
	m.RGBA = newm
}

func (m *Image) randomBrightness(c color.RGBA, max uint8) color.RGBA {
	minc := min3(c.R, c.G, c.B)
	maxc := max3(c.R, c.G, c.B)
	if maxc > max {
		return c
	}
	n := m.rng.Intn(int(max-maxc)) - int(minc)
	return color.RGBA{
		uint8(int(c.R) + n),
		uint8(int(c.G) + n),
		uint8(int(c.B) + n),
		uint8(c.A),
	}
}

func min3(x, y, z uint8) (m uint8) {
	m = x
	if y < m {
		m = y
	}
	if z < m {
		m = z
	}
	return
}

func max3(x, y, z uint8) (m uint8) {
	m = x
	if y > m {
		m = y
	}
	if z > m {
		m = z
	}
	return
}
<commit_msg>reduce color variety<commit_after>// Copyright 2011-2014 Dmitry Chestnykh. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package captcha

import (
	"bytes"
	"image"
	"image/color"
	"image/png"
	"io"
	"math"
)

const (
	// Standard width and height of a captcha image.
	StdWidth  = 240
	StdHeight = 80
	// Maximum absolute skew factor of a single digit.
	maxSkew = 0.7
	// Number of background circles.
	circleCount = 20
	distortCircleCount = 5
)

type Image struct {
	*image.RGBA
	numWidth  int
	numHeight int
	dotSize   int
	rng       siprng
}

var background color.Color = color.RGBA{144, 144, 144, 0xFF}

// NewImage returns a new captcha image of the given width and height with the
// given digits, where each digit must be in range 0-9.
func NewImage(id string, digits []byte, width, height int) *Image {
	m := new(Image)

	// Initialize PRNG.
	m.rng.Seed(deriveSeed(imageSeedPurpose, id, digits))

	prim := color.RGBA{
		uint8(m.rng.Intn(129)),
		uint8(m.rng.Intn(129)),
		uint8(m.rng.Intn(129)),
		0xFF,
	}

	m.RGBA = image.NewRGBA(image.Rect(0, 0, width, height))
	m.calculateSizes(width, height, len(digits))
	// Randomly position captcha inside the image.
	maxx := width - (m.numWidth+m.dotSize)*len(digits) - m.dotSize
	maxy := height - m.numHeight - m.dotSize*2
	var border int
	if width > height {
		border = height / 10
	} else {
		border = width / 10
	}
	x := m.rng.Int(border, maxx-border)
	y := m.rng.Int(border, maxy-border)
	// Draw digits.
	for _, n := range digits {
		m.drawDigit(getChar(n), x, y, prim)
		x += m.numWidth + m.dotSize
	}
	// Draw strike-through line.
	m.strikeThrough(prim)

	m.fillWithDistortCircles(distortCircleCount, width/4)
	// Apply wave distortion.
	m.distort(m.rng.Float(5, 10), m.rng.Float(100, 200))
	// Fill image with random circles.
	m.fillWithCircles(circleCount, m.dotSize)
	return m
}

/*
func (m *Image) getRandomPalette() color.Palette {
	p := make([]color.Color, circleCount+1)
	// Transparent color.
	p[0] = color.RGBA{0xFF, 0xFF, 0xFF, 0x00}
	// Primary color.
	prim := color.RGBA{
		uint8(m.rng.Intn(129)),
		uint8(m.rng.Intn(129)),
		uint8(m.rng.Intn(129)),
		0xFF,
	}
	p[1] = prim
	// Circle colors.
	for i := 2; i <= circleCount; i++ {
		p[i] = m.randomBrightness(prim, 255)
	}
	return p
}*/

// encodedPNG encodes an image to PNG and returns
// the result as a byte slice.
func (m *Image) encodedPNG() []byte {
	var buf bytes.Buffer
	if err := png.Encode(&buf, m); err != nil {
		panic(err.Error())
	}
	return buf.Bytes()
}

// WriteTo writes captcha image in PNG format into the given writer.
func (m *Image) WriteTo(w io.Writer) (int64, error) {
	n, err := w.Write(m.encodedPNG())
	return int64(n), err
}

func (m *Image) calculateSizes(width, height, ncount int) {
	// Goal: fit all digits inside the image.
	var border int
	if width > height {
		border = height / 8
	} else {
		border = width / 8
	}
	// Convert everything to floats for calculations.
	w := float64(width - border*2)
	h := float64(height - border*2)
	// fw takes into account 1-dot spacing between digits.
	fw := float64(fontWidth + 1)
	fh := float64(fontHeight)
	nc := float64(ncount)
	// Calculate the width of a single digit taking into account only the
	// width of the image.
	nw := w / nc
	// Calculate the height of a digit from this width.
	nh := nw * fh / fw
	// Digit too high?
	if nh > h {
		// Fit digits based on height.
		nh = h
		nw = fw / fh * nh
	}
	// Calculate dot size.
	m.dotSize = int((nh / fh))
	if m.dotSize < 1 {
		m.dotSize = 1
	}
	// Save everything, making the actual width smaller by 1 dot to account
	// for spacing between digits.
	m.numWidth = int(nw) - m.dotSize
	m.numHeight = int(nh)
}

func (m *Image) drawHorizLine(fromX, toX, y int, color color.Color) {
	for x := fromX; x <= toX; x++ {
		m.Set(x, y, color)
	}
}

func (m *Image) drawCircle(x, y, radius int, color color.Color) {
	f := 1 - radius
	dfx := 1
	dfy := -2 * radius
	xo := 0
	yo := radius

	m.Set(x, y+radius, color)
	m.Set(x, y-radius, color)
	m.drawHorizLine(x-radius, x+radius, y, color)

	for xo < yo {
		if f >= 0 {
			yo--
			dfy += 2
			f += dfy
		}
		xo++
		dfx += 2
		f += dfx
		m.drawHorizLine(x-xo, x+xo, y+yo, color)
		m.drawHorizLine(x-xo, x+xo, y-yo, color)
		m.drawHorizLine(x-yo, x+yo, y+xo, color)
		m.drawHorizLine(x-yo, x+yo, y-xo, color)
	}
}

func (m *Image) colorDistortCircle(xc, yc, r int) {
	d := 2*r
	r2 := r*r

	dg := int(m.rng.Intn(120)) - 60
	if dg < 0 {
		dg -= 10
	} else {
		dg += 10
	}

	dr := int(m.rng.Intn(120)) - 60
	if dr < 0 {
		dr -= 10
	} else {
		dr += 10
	}

	db := int(m.rng.Intn(120)) - 60
	if db < 0 {
		db -= 10
	} else {
		db += 10
	}

	for y := 0; y <= d; y++ {
		for x := 0; x <= d; x++ {
			x2 := (x-r)*(x-r)
			y2 := (y-r)*(y-r)

			if x2+y2 <= r2 {
				xp := xc - r + x
				yp := yc - r + y
				r,g,b,a := m.At(xp, yp).RGBA()
				if a != 0x00 { //background
					m.Set(xp, yp, color.RGBA{
						uint8(int(r)+dr),
						uint8(int(g)+dg),
						uint8(int(b)+db),
						0xFF,
					})
				}
			}
		}
	}
}

func (m *Image) fillWithCircles(n, maxradius int) {
	maxx := m.Bounds().Max.X
	maxy := m.Bounds().Max.Y
	for i := 0; i < n; i++ {
		color := color.RGBA{
			uint8(m.rng.Intn(129)),
			uint8(m.rng.Intn(129)),
			uint8(m.rng.Intn(129)),
			0xFF,
		}
		r := m.rng.Int(1, maxradius)
		m.drawCircle(m.rng.Int(r, maxx-r), m.rng.Int(r, maxy-r), r, color)
	}
}

func (m *Image) fillWithDistortCircles(n, maxradius int) {
	maxx := m.Bounds().Max.X
	maxy := m.Bounds().Max.Y
	for i := 0; i < n; i++ {
		r := m.rng.Int(maxradius/10, maxradius-1)
		m.colorDistortCircle(m.rng.Int(0, maxx), m.rng.Int(0, maxy), r)
	}
}

func (m *Image) strikeThrough(color color.Color) {
	maxx := m.Bounds().Max.X
	maxy := m.Bounds().Max.Y
	y := m.rng.Int(maxy/3, maxy-maxy/3)
	amplitude := m.rng.Float(5, 20)
	period := m.rng.Float(80, 180)
	dx := 2.0 * math.Pi / period
	for x := 0; x < maxx; x++ {
		xo := amplitude * math.Cos(float64(y)*dx)
		yo := amplitude * math.Sin(float64(x)*dx)
		for yn := 0; yn < m.dotSize; yn++ {
			r := m.rng.Int(0, m.dotSize)
			m.drawCircle(x+int(xo), y+int(yo)+(yn*m.dotSize), r/2, color)
		}
	}
}

func (m *Image) drawDigit(digit []byte, x, y int, color color.Color) {
	skf := m.rng.Float(-maxSkew, maxSkew)
	xs := float64(x)
	r := m.dotSize / 2
	y += m.rng.Int(-r, r)
	for yo := 0; yo < fontHeight; yo++ {
		for xo := 0; xo < fontWidth; xo++ {
			if digit[yo*fontWidth+xo] != blackChar {
				continue
			}
			m.drawCircle(x+xo*m.dotSize, y+yo*m.dotSize, r, color)
		}
		xs += skf
		x = int(xs)
	}
}

func (m *Image) distort(amplude float64, period float64) {
	w := m.Bounds().Max.X
	h := m.Bounds().Max.Y
	r32,g32,b32,_ := background.RGBA()
	r:=uint8(r32)
	g:=uint8(g32)
	b:=uint8(b32)
	dbg := 32.0/float32(w);

	oldm := m
	newm := image.NewRGBA(image.Rect(0, 0, w, h))

	dx := 2.0 * math.Pi / period
	for x := 0; x < w; x++ {
		dbgx:=uint8(float32(x)*dbg)
		bg:=color.RGBA{r-dbgx, g-dbgx, b-dbgx, 0xFF}
		for y := 0; y < h; y++ {
			xo := amplude * math.Sin(float64(y)*dx)
			yo := amplude * math.Cos(float64(x)*dx)
			c := oldm.At(x+int(xo), y+int(yo))
			newm.Set(x, y, c)
			_,_,_, a := c.RGBA()
			if a == 0x00 {
				newm.Set(x, y, bg)
			}
		}
	}
	m.RGBA = newm
}

func (m *Image) randomBrightness(c color.RGBA, max uint8) color.RGBA {
	minc := min3(c.R, c.G, c.B)
	maxc := max3(c.R, c.G, c.B)
	if maxc > max {
		return c
	}
	n := m.rng.Intn(int(max-maxc)) - int(minc)
	return color.RGBA{
		uint8(int(c.R) + n),
		uint8(int(c.G) + n),
		uint8(int(c.B) + n),
		uint8(c.A),
	}
}

func min3(x, y, z uint8) (m uint8) {
	m = x
	if y < m {
		m = y
	}
	if z < m {
		m = z
	}
	return
}

func max3(x, y, z uint8) (m uint8) {
	m = x
	if y > m {
		m = y
	}
	if z > m {
		m = z
	}
	return
}
<|endoftext|>


<commit_before>package imgio

import (
	"errors"
	"image/draw"
	"io"
	"sync"
)

type Image struct {
	img draw.Image
	gen PointsSequenceGenerator
	prw PointReadWriter
	mux sync.RWMutex

	pointCursor uint64
	byteCursor  int
}

func NewImage(img draw.Image, gen PointsSequenceGenerator, prw PointReadWriter) *Image {
	return &Image{
		img: img,
		gen: gen,
		prw: prw,
	}
}

// Read implements io.Reader interface
func (f *Image) Read(p []byte) (n int, err error) {
	f.mux.RLock()
	defer f.mux.RUnlock()

	f.gen.Seek(f.pointCursor)

	if !f.gen.Valid() {
		return 0, io.EOF
	}
	if len(p) == 0 {
		return 0, nil
	}

	for {
		if !f.gen.Valid() {
			return n, io.EOF
		}
		if n >= len(p) {
			break
		}

		point := f.gen.Current()
		color := f.img.At(point.X, point.Y)
		buff, nBytesRead := f.prw.Read(f.byteCursor, color, point)

		if len(p)-n > nBytesRead {
			copy(p[n:], buff[:nBytesRead])
			n += nBytesRead

			f.gen.Next()
			f.pointCursor++
			f.byteCursor = 0
		} else {
			end := nBytesRead - len(p) + n
			copy(p[n:], buff[:end])
			n += end
			f.byteCursor = end + 1
			break
		}
	}

	return
}

// Write implements io.Writer interface
func (f *Image) Write(p []byte) (n int, err error) {
	if len(p) == 0 {
		return 0, nil
	}

	f.mux.Lock()
	defer f.mux.Unlock()

	f.gen.Seek(f.pointCursor)

	if !f.gen.Valid() {
		return 0, errors.New("overflow")
	}

	for {
		if !f.gen.Valid() {
			return n, errors.New("overflow")
		}
		if len(p) == 0 {
			return n, nil
		}

		point := f.gen.Current()
		srcColor := f.img.At(point.X, point.Y)
		color, writtenBytes := f.prw.Write(p, f.byteCursor, srcColor, point)
		f.img.Set(point.X, point.Y, color)
		f.gen.Next()
		f.pointCursor++
		n += writtenBytes
		p = p[writtenBytes:]
	}

	return
}

func (s *Image) Size() (size int64) {
	s.gen.Rewind()
	for s.gen.Valid() {
		point := s.gen.Current()
		size += s.prw.Size(point)
		s.gen.Next()
	}
	return
}
<commit_msg>Fix bytes cursor<commit_after>package imgio

import (
	"errors"
	"image/draw"
	"io"
	"sync"
)

type Image struct {
	img draw.Image
	gen PointsSequenceGenerator
	prw PointReadWriter
	mux sync.RWMutex

	pointCursor uint64
	byteCursor  int
}

func NewImage(img draw.Image, gen PointsSequenceGenerator, prw PointReadWriter) *Image {
	return &Image{
		img: img,
		gen: gen,
		prw: prw,
	}
}

// Read implements io.Reader interface
func (f *Image) Read(p []byte) (n int, err error) {
	f.mux.RLock()
	defer f.mux.RUnlock()

	f.gen.Seek(f.pointCursor)

	if !f.gen.Valid() {
		return 0, io.EOF
	}
	if len(p) == 0 {
		return 0, nil
	}

	for {
		if !f.gen.Valid() {
			return n, io.EOF
		}
		if n >= len(p) {
			return
		}

		point := f.gen.Current()
		color := f.img.At(point.X, point.Y)
		buff, nBytesRead := f.prw.Read(f.byteCursor, color, point)

		if len(p)-n > nBytesRead {
			copy(p[n:], buff[:nBytesRead])
			n += nBytesRead
			f.gen.Next()
			f.pointCursor++
			f.byteCursor = 0
		} else {
			end := nBytesRead - len(p) + n
			copy(p[n:], buff[:end])
			n += end
			f.byteCursor = end + 1
			return
		}
	}

	return
}

// Write implements io.Writer interface
func (f *Image) Write(p []byte) (n int, err error) {
	if len(p) == 0 {
		return 0, nil
	}

	f.mux.Lock()
	defer f.mux.Unlock()

	f.gen.Seek(f.pointCursor)

	for {
		if !f.gen.Valid() {
			return n, errors.New("overflow")
		}
		if len(p) == 0 {
			return n, nil
		}

		point := f.gen.Current()
		srcColor := f.img.At(point.X, point.Y)
		color, writtenBytes := f.prw.Write(p, f.byteCursor, srcColor, point)
		f.img.Set(point.X, point.Y, color)
		f.gen.Next()
		f.pointCursor++
		n += writtenBytes
		p = p[writtenBytes:]
		if f.prw.Size(point) > int64(writtenBytes) {
			f.byteCursor = writtenBytes + 1
		}
	}

	return
}

func (s *Image) Size() (size int64) {
	s.gen.Rewind()
	for s.gen.Valid() {
		point := s.gen.Current()
		size += s.prw.Size(point)
		s.gen.Next()
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2018 AlexStocks(https://github.com/AlexStocks).
// All rights reserved.  Use of this source code is
// governed by Apache License 2.0.

// Package gxfilter provides a interface for service filter
package gxfilter

import (
	"context"
	"fmt"
	"strings"
	"sync/atomic"
	"time"
)

import (
	"github.com/AlexStocks/goext/database/registry"
	jerrors "github.com/juju/errors"
)

/////////////////////////////////////
// ServiceArray
/////////////////////////////////////

var (
	ErrServiceArrayEmpty = jerrors.New("ServiceArray empty")
)

// ServiceHash is a user service hash interface to select a service provider.
// @context is trasfered from ServiceArray::Select
type ServiceHash func(context.Context, *ServiceArray) (*gxregistry.Service, error)

type ServiceArray struct {
	Arr    []*gxregistry.Service
	Active time.Time
	idx    int64 // the last selected service provider index
}

func NewServiceArray(Arr []*gxregistry.Service) *ServiceArray {
	return &ServiceArray{
		Arr:    Arr,
		Active: time.Now(),
	}
}

func (s ServiceArray) String() string {
	var builder strings.Builder
	builder.WriteString(fmt.Sprintf("Active:%s, idx:%d, Arr len:%d, Arr:{", s.Active, s.idx, len(s.Arr)))
	for i := range s.Arr {
		builder.WriteString(fmt.Sprintf("%d:%s, ", i, s.Arr[i]))
	}
	builder.WriteString("}")

	return builder.String()
}

func (s *ServiceArray) Select(ctx context.Context, hash ServiceHash) (*gxregistry.Service, error) {
	arrSize := len(s.Arr)
	if arrSize == 0 {
		return nil, ErrServiceArrayEmpty
	}

	if hash != nil {
		return hash(ctx, s)
	}

	// default use roundrobin hash alg
	idx := atomic.AddInt64(&s.idx, 1) % int64(arrSize)

	return s.Arr[idx], nil
}

func (s *ServiceArray) Add(service *gxregistry.Service, ttl time.Duration) {
	s.Arr = append(s.Arr, service)
	s.Active = time.Now().Add(ttl)
}

func (s *ServiceArray) Del(service *gxregistry.Service, ttl time.Duration) {
	for i, svc := range s.Arr {
		if svc.Equal(service) {
			s.Arr = append(s.Arr[:i], s.Arr[i+1:]...)
			s.Active = time.Now().Add(ttl)
			// log.Debug("i:%d, new services:%+v", i, services)
			break
		}
	}
}
<commit_msg>Mod: trasfer -> transfer<commit_after>// Copyright 2018 AlexStocks(https://github.com/AlexStocks).
// All rights reserved.  Use of this source code is
// governed by Apache License 2.0.

// Package gxfilter provides a interface for service filter
package gxfilter

import (
	"context"
	"fmt"
	"strings"
	"sync/atomic"
	"time"
)

import (
	"github.com/AlexStocks/goext/database/registry"
	jerrors "github.com/juju/errors"
)

/////////////////////////////////////
// ServiceArray
/////////////////////////////////////

var (
	ErrServiceArrayEmpty = jerrors.New("ServiceArray empty")
)

// ServiceHash is a user service hash interface to select a service provider.
// @context is transfered from ServiceArray::Select
type ServiceHash func(context.Context, *ServiceArray) (*gxregistry.Service, error)

type ServiceArray struct {
	Arr    []*gxregistry.Service
	Active time.Time
	idx    int64 // the last selected service provider index
}

func NewServiceArray(Arr []*gxregistry.Service) *ServiceArray {
	return &ServiceArray{
		Arr:    Arr,
		Active: time.Now(),
	}
}

func (s ServiceArray) String() string {
	var builder strings.Builder
	builder.WriteString(fmt.Sprintf("Active:%s, idx:%d, Arr len:%d, Arr:{", s.Active, s.idx, len(s.Arr)))
	for i := range s.Arr {
		builder.WriteString(fmt.Sprintf("%d:%s, ", i, s.Arr[i]))
	}
	builder.WriteString("}")

	return builder.String()
}

func (s *ServiceArray) Select(ctx context.Context, hash ServiceHash) (*gxregistry.Service, error) {
	arrSize := len(s.Arr)
	if arrSize == 0 {
		return nil, ErrServiceArrayEmpty
	}

	if hash != nil {
		return hash(ctx, s)
	}

	// default use roundrobin hash alg
	idx := atomic.AddInt64(&s.idx, 1) % int64(arrSize)

	return s.Arr[idx], nil
}

func (s *ServiceArray) Add(service *gxregistry.Service, ttl time.Duration) {
	s.Arr = append(s.Arr, service)
	s.Active = time.Now().Add(ttl)
}

func (s *ServiceArray) Del(service *gxregistry.Service, ttl time.Duration) {
	for i, svc := range s.Arr {
		if svc.Equal(service) {
			s.Arr = append(s.Arr[:i], s.Arr[i+1:]...)
			s.Active = time.Now().Add(ttl)
			// log.Debug("i:%d, new services:%+v", i, services)
			break
		}
	}
}
<|endoftext|>


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

package retrieval

import (
	"testing"

	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"

	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/storage/local"
)

func TestTargetSetRecreatesTargetGroupsEveryRun(t *testing.T) {
	scrapeConfig := &config.ScrapeConfig{}

	sOne := `
job_name: "foo"
dns_sd_configs:
- names:
  - "srv.name.one.example.org"
`
	if err := yaml.Unmarshal([]byte(sOne), scrapeConfig); err != nil {
		t.Fatalf("Unable to load YAML config sOne: %s", err)
	}

	// Not properly setting it up, but that seems okay
	mss := &local.MemorySeriesStorage{}

	ts := newTargetSet(scrapeConfig, mss)

	ts.runProviders(context.Background(), providersFromConfig(scrapeConfig))

	verifyPresence(t, ts.tgroups, "dns/0/srv.name.one.example.org", true)

	sTwo := `
job_name: "foo"
dns_sd_configs:
- names:
  - "srv.name.two.example.org"
`
	if err := yaml.Unmarshal([]byte(sTwo), scrapeConfig); err != nil {
		t.Fatalf("Unable to load YAML config sTwo: %s", err)
	}

	ts.runProviders(context.Background(), providersFromConfig(scrapeConfig))

	verifyPresence(t, ts.tgroups, "dns/0/srv.name.one.example.org", false)
	verifyPresence(t, ts.tgroups, "dns/0/srv.name.two.example.org", true)

}

func verifyPresence(t *testing.T, tgroups map[string][]*Target, name string, present bool) {
	if _, ok := tgroups[name]; ok != present {
		msg := ""
		if !present {
			msg = "not "
		}
		t.Fatalf("'%s' should %sbe present in TargetSet.tgroups: %s", name, msg, tgroups)
	}

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

package retrieval

import (
	"testing"

	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"

	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/storage/local"
)

func TestTargetSetRecreatesTargetGroupsEveryRun(t *testing.T) {

	verifyPresence := func(tgroups map[string][]*Target, name string, present bool) {
		if _, ok := tgroups[name]; ok != present {
			msg := ""
			if !present {
				msg = "not "
			}
			t.Fatalf("'%s' should %sbe present in TargetSet.tgroups: %s", name, msg, tgroups)
		}

	}

	scrapeConfig := &config.ScrapeConfig{}

	sOne := `
job_name: "foo"
dns_sd_configs:
- names:
  - "srv.name.one.example.org"
`
	if err := yaml.Unmarshal([]byte(sOne), scrapeConfig); err != nil {
		t.Fatalf("Unable to load YAML config sOne: %s", err)
	}

	// Not properly setting it up, but that seems okay
	mss := &local.MemorySeriesStorage{}

	ts := newTargetSet(scrapeConfig, mss)

	ts.runProviders(context.Background(), providersFromConfig(scrapeConfig))

	verifyPresence(ts.tgroups, "dns/0/srv.name.one.example.org", true)

	sTwo := `
job_name: "foo"
dns_sd_configs:
- names:
  - "srv.name.two.example.org"
`
	if err := yaml.Unmarshal([]byte(sTwo), scrapeConfig); err != nil {
		t.Fatalf("Unable to load YAML config sTwo: %s", err)
	}

	ts.runProviders(context.Background(), providersFromConfig(scrapeConfig))

	verifyPresence(ts.tgroups, "dns/0/srv.name.one.example.org", false)
	verifyPresence(ts.tgroups, "dns/0/srv.name.two.example.org", true)
}
<|endoftext|>