<commit_before>package mackerelplugin

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Metrics struct {
	Name    string  `json:"name"`
	Label   string  `json:"label"`
	Diff    bool    `json:"diff"`
	Type    string  `json:"type"`
	Stacked bool    `json:"stacked"`
	Scale   float64 `json:"scale"`
}

type Graphs struct {
	Label   string    `json:"label"`
	Unit    string    `json:"unit"`
	Metrics []Metrics `json:"metrics"`
}

type Plugin interface {
	FetchMetrics() (map[string]interface{}, error)
	GraphDefinition() map[string]Graphs
}

type MackerelPlugin struct {
	Plugin
	Tempfile string
}

func NewMackerelPlugin(plugin Plugin) MackerelPlugin {
	mp := MackerelPlugin{plugin, "/tmp/mackerel-plugin-default"}
	return mp
}

func (h *MackerelPlugin) printValue(w io.Writer, key string, value interface{}, now time.Time) {
	switch value.(type) {
	case uint32:
		fmt.Fprintf(w, "%s\t%d\t%d\n", key, value.(uint32), now.Unix())
	case uint64:
		fmt.Fprintf(w, "%s\t%d\t%d\n", key, value.(uint64), now.Unix())
	case float64:
		if math.IsNaN(value.(float64)) || math.IsInf(value.(float64), 0) {
			log.Printf("Invalid value: key = %s, value = %f\n", key, value)
		} else {
			fmt.Fprintf(w, "%s\t%f\t%d\n", key, value.(float64), now.Unix())
		}
	}
}

func (h *MackerelPlugin) fetchLastValues() (map[string]interface{}, time.Time, error) {
	lastTime := time.Now()

	f, err := os.Open(h.Tempfilename())
	if err != nil {
		if os.IsNotExist(err) {
			return nil, lastTime, nil
		}
		return nil, lastTime, err
	}
	defer f.Close()

	stat := make(map[string]interface{})
	decoder := json.NewDecoder(f)
	err = decoder.Decode(&stat)
	switch stat["_lastTime"].(type) {
	case float64:
		lastTime = time.Unix(int64(stat["_lastTime"].(float64)), 0)
	case int64:
		lastTime = time.Unix(stat["_lastTime"].(int64), 0)
	}
	if err != nil {
		return stat, lastTime, err
	}
	return stat, lastTime, nil
}

func (h *MackerelPlugin) saveValues(values map[string]interface{}, now time.Time) error {
	f, err := os.Create(h.Tempfilename())
	if err != nil {
		return err
	}
	defer f.Close()

	values["_lastTime"] = now.Unix()
	encoder := json.NewEncoder(f)
	err = encoder.Encode(values)
	if err != nil {
		return err
	}

	return nil
}

func (h *MackerelPlugin) calcDiff(value float64, now time.Time, lastValue float64, lastTime time.Time) (float64, error) {
	diffTime := now.Unix() - lastTime.Unix()
	if diffTime > 600 {
		return 0, errors.New("Too long duration")
	}

	diff := (value - lastValue) * 60 / float64(diffTime)

	if lastValue <= value {
		return diff, nil
	}
	return 0.0, errors.New("Counter seems to be reseted.")
}

func (h *MackerelPlugin) calcDiffUint32(value uint32, now time.Time, lastValue uint32, lastTime time.Time, lastDiff float64) (float64, error) {
	diffTime := now.Unix() - lastTime.Unix()
	if diffTime > 600 {
		return 0, errors.New("Too long duration")
	}

	diff := float64((value-lastValue)*60) / float64(diffTime)

	if lastValue <= value || diff < lastDiff*10 {
		return diff, nil
	}
	return 0.0, errors.New("Counter seems to be reseted.")

}

func (h *MackerelPlugin) calcDiffUint64(value uint64, now time.Time, lastValue uint64, lastTime time.Time, lastDiff float64) (float64, error) {
	diffTime := now.Unix() - lastTime.Unix()
	if diffTime > 600 {
		return 0, errors.New("Too long duration")
	}

	diff := float64((value-lastValue)*60) / float64(diffTime)

	if lastValue <= value || diff < lastDiff*10 {
		return diff, nil
	}
	return 0.0, errors.New("Counter seems to be reseted.")
}

func (h *MackerelPlugin) Tempfilename() string {
	return h.Tempfile
}

func (h *MackerelPlugin) formatValues(prefix string, metric Metrics, stat *map[string]interface{}, lastStat *map[string]interface{}, now time.Time, lastTime time.Time) {
	value, ok := (*stat)[metric.Name]
	if !ok || value == nil {
		return
	}

	switch value.(type) {
	case string:
		switch metric.Type {
		case "uint32":
			value, _ = strconv.ParseUint(value.(string), 10, 32)
		case "uint64":
			value, _ = strconv.ParseUint(value.(string), 10, 64)
		default:
			value, _ = strconv.ParseFloat(value.(string), 64)
		}
	}

	if metric.Diff {
		_, ok := (*lastStat)[metric.Name]
		if ok {
			var lastDiff float64
			if (*lastStat)[".last_diff."+metric.Name] != nil {
				lastDiff = toFloat64((*lastStat)[".last_diff."+metric.Name])
			}
			var err error
			switch metric.Type {
			case "uint32":
				value, err = h.calcDiffUint32(toUint32(value), now, toUint32((*lastStat)[metric.Name]), lastTime, lastDiff)
			case "uint64":
				value, err = h.calcDiffUint64(toUint64(value), now, toUint64((*lastStat)[metric.Name]), lastTime, lastDiff)
			default:
				value, err = h.calcDiff(toFloat64(value), now, toFloat64((*lastStat)[metric.Name]), lastTime)
			}
			if err != nil {
				log.Println("OutputValues: ", err)
				return
			} else {
				(*stat)[".last_diff."+metric.Name] = value
			}
		} else {
			log.Printf("%s does not exist at last fetch\n", metric.Name)
			return
		}
	}

	if metric.Scale != 0 {
		switch metric.Type {
		case "uint32":
			value = toUint32(value) * uint32(metric.Scale)
		case "uint64":
			value = toUint64(value) * uint64(metric.Scale)
		default:
			value = toFloat64(value) * metric.Scale
		}
	}

	if len(prefix) > 0 {
		h.printValue(os.Stdout, prefix+"."+metric.Name, value, now)
	} else {
		h.printValue(os.Stdout, metric.Name, value, now)
	}
}

func (h *MackerelPlugin) formatValuesWithWildcard(prefix string, metric Metrics, stat *map[string]interface{}, lastStat *map[string]interface{}, now time.Time, lastTime time.Time) {
	regexpStr := `\A` + prefix + "." + metric.Name
	regexpStr = strings.Replace(regexpStr, ".", "\\.", -1)
	regexpStr = strings.Replace(regexpStr, "*", "[-a-zA-Z0-9_]+", -1)
	regexpStr = strings.Replace(regexpStr, "#", "[-a-zA-Z0-9_]+", -1)
	re, err := regexp.Compile(regexpStr)
	if err != nil {
		log.Fatalln("Failed to compile regexp: ", err)
	}
	for k, _ := range *stat {
		if re.MatchString(k) {
			metricEach := metric
			metricEach.Name = k
			h.formatValues("", metricEach, stat, lastStat, now, lastTime)
		}
	}
}

func (h *MackerelPlugin) Run() {
	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		h.OutputDefinitions()
	} else {
		h.OutputValues()
	}
}

func (h *MackerelPlugin) OutputValues() {
	now := time.Now()
	stat, err := h.FetchMetrics()
	if err != nil {
		log.Fatalln("OutputValues: ", err)
	}

	lastStat, lastTime, err := h.fetchLastValues()
	if err != nil {
		log.Println("fetchLastValues (ignore):", err)
	}

	for key, graph := range h.GraphDefinition() {
		for _, metric := range graph.Metrics {
			if strings.ContainsAny(key+metric.Name, "*#") {
				h.formatValuesWithWildcard(key, metric, &stat, &lastStat, now, lastTime)
			} else {
				h.formatValues(key, metric, &stat, &lastStat, now, lastTime)
			}
		}
	}

	err = h.saveValues(stat, now)
	if err != nil {
		log.Fatalf("saveValues: ", err)
	}
}

type GraphDef struct {
	Graphs map[string]Graphs `json:"graphs"`
}

func (h *MackerelPlugin) OutputDefinitions() {
	fmt.Println("# mackerel-agent-plugin")
	var graphs GraphDef
	graphs.Graphs = h.GraphDefinition()

	b, err := json.Marshal(graphs)
	if err != nil {
		log.Fatalln("OutputDefinitions: ", err)
	}
	fmt.Println(string(b))
}

func toUint32(value interface{}) uint32 {
	var ret uint32
	switch value.(type) {
	case uint32:
		ret = value.(uint32)
	case uint64:
		ret = uint32(value.(uint64))
	case float64:
		ret = uint32(value.(float64))
	case string:
		v, err := strconv.ParseUint(value.(string), 10, 32)
		if err != nil {
			ret = uint32(v)
		}
	}
	return ret
}

func toUint64(value interface{}) uint64 {
	var ret uint64
	switch value.(type) {
	case uint32:
		ret = uint64(value.(uint32))
	case uint64:
		ret = value.(uint64)
	case float64:
		ret = uint64(value.(float64))
	case string:
		ret, _ = strconv.ParseUint(value.(string), 10, 64)
	}
	return ret
}

func toFloat64(value interface{}) float64 {
	var ret float64
	switch value.(type) {
	case uint32:
		ret = float64(value.(uint32))
	case uint64:
		ret = float64(value.(uint64))
	case float64:
		ret = value.(float64)
	case string:
		ret, _ = strconv.ParseFloat(value.(string), 64)
	}
	return ret
}
<commit_msg>Fix toUint32 when given a string value<commit_after>package mackerelplugin

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Metrics struct {
	Name    string  `json:"name"`
	Label   string  `json:"label"`
	Diff    bool    `json:"diff"`
	Type    string  `json:"type"`
	Stacked bool    `json:"stacked"`
	Scale   float64 `json:"scale"`
}

type Graphs struct {
	Label   string    `json:"label"`
	Unit    string    `json:"unit"`
	Metrics []Metrics `json:"metrics"`
}

type Plugin interface {
	FetchMetrics() (map[string]interface{}, error)
	GraphDefinition() map[string]Graphs
}

type MackerelPlugin struct {
	Plugin
	Tempfile string
}

func NewMackerelPlugin(plugin Plugin) MackerelPlugin {
	mp := MackerelPlugin{plugin, "/tmp/mackerel-plugin-default"}
	return mp
}

func (h *MackerelPlugin) printValue(w io.Writer, key string, value interface{}, now time.Time) {
	switch value.(type) {
	case uint32:
		fmt.Fprintf(w, "%s\t%d\t%d\n", key, value.(uint32), now.Unix())
	case uint64:
		fmt.Fprintf(w, "%s\t%d\t%d\n", key, value.(uint64), now.Unix())
	case float64:
		if math.IsNaN(value.(float64)) || math.IsInf(value.(float64), 0) {
			log.Printf("Invalid value: key = %s, value = %f\n", key, value)
		} else {
			fmt.Fprintf(w, "%s\t%f\t%d\n", key, value.(float64), now.Unix())
		}
	}
}

func (h *MackerelPlugin) fetchLastValues() (map[string]interface{}, time.Time, error) {
	lastTime := time.Now()

	f, err := os.Open(h.Tempfilename())
	if err != nil {
		if os.IsNotExist(err) {
			return nil, lastTime, nil
		}
		return nil, lastTime, err
	}
	defer f.Close()

	stat := make(map[string]interface{})
	decoder := json.NewDecoder(f)
	err = decoder.Decode(&stat)
	switch stat["_lastTime"].(type) {
	case float64:
		lastTime = time.Unix(int64(stat["_lastTime"].(float64)), 0)
	case int64:
		lastTime = time.Unix(stat["_lastTime"].(int64), 0)
	}
	if err != nil {
		return stat, lastTime, err
	}
	return stat, lastTime, nil
}

func (h *MackerelPlugin) saveValues(values map[string]interface{}, now time.Time) error {
	f, err := os.Create(h.Tempfilename())
	if err != nil {
		return err
	}
	defer f.Close()

	values["_lastTime"] = now.Unix()
	encoder := json.NewEncoder(f)
	err = encoder.Encode(values)
	if err != nil {
		return err
	}

	return nil
}

func (h *MackerelPlugin) calcDiff(value float64, now time.Time, lastValue float64, lastTime time.Time) (float64, error) {
	diffTime := now.Unix() - lastTime.Unix()
	if diffTime > 600 {
		return 0, errors.New("Too long duration")
	}

	diff := (value - lastValue) * 60 / float64(diffTime)

	if lastValue <= value {
		return diff, nil
	}
	return 0.0, errors.New("Counter seems to be reseted.")
}

func (h *MackerelPlugin) calcDiffUint32(value uint32, now time.Time, lastValue uint32, lastTime time.Time, lastDiff float64) (float64, error) {
	diffTime := now.Unix() - lastTime.Unix()
	if diffTime > 600 {
		return 0, errors.New("Too long duration")
	}

	diff := float64((value-lastValue)*60) / float64(diffTime)

	if lastValue <= value || diff < lastDiff*10 {
		return diff, nil
	}
	return 0.0, errors.New("Counter seems to be reseted.")

}

func (h *MackerelPlugin) calcDiffUint64(value uint64, now time.Time, lastValue uint64, lastTime time.Time, lastDiff float64) (float64, error) {
	diffTime := now.Unix() - lastTime.Unix()
	if diffTime > 600 {
		return 0, errors.New("Too long duration")
	}

	diff := float64((value-lastValue)*60) / float64(diffTime)

	if lastValue <= value || diff < lastDiff*10 {
		return diff, nil
	}
	return 0.0, errors.New("Counter seems to be reseted.")
}

func (h *MackerelPlugin) Tempfilename() string {
	return h.Tempfile
}

func (h *MackerelPlugin) formatValues(prefix string, metric Metrics, stat *map[string]interface{}, lastStat *map[string]interface{}, now time.Time, lastTime time.Time) {
	value, ok := (*stat)[metric.Name]
	if !ok || value == nil {
		return
	}

	switch value.(type) {
	case string:
		switch metric.Type {
		case "uint32":
			value, _ = strconv.ParseUint(value.(string), 10, 32)
		case "uint64":
			value, _ = strconv.ParseUint(value.(string), 10, 64)
		default:
			value, _ = strconv.ParseFloat(value.(string), 64)
		}
	}

	if metric.Diff {
		_, ok := (*lastStat)[metric.Name]
		if ok {
			var lastDiff float64
			if (*lastStat)[".last_diff."+metric.Name] != nil {
				lastDiff = toFloat64((*lastStat)[".last_diff."+metric.Name])
			}
			var err error
			switch metric.Type {
			case "uint32":
				value, err = h.calcDiffUint32(toUint32(value), now, toUint32((*lastStat)[metric.Name]), lastTime, lastDiff)
			case "uint64":
				value, err = h.calcDiffUint64(toUint64(value), now, toUint64((*lastStat)[metric.Name]), lastTime, lastDiff)
			default:
				value, err = h.calcDiff(toFloat64(value), now, toFloat64((*lastStat)[metric.Name]), lastTime)
			}
			if err != nil {
				log.Println("OutputValues: ", err)
				return
			} else {
				(*stat)[".last_diff."+metric.Name] = value
			}
		} else {
			log.Printf("%s does not exist at last fetch\n", metric.Name)
			return
		}
	}

	if metric.Scale != 0 {
		switch metric.Type {
		case "uint32":
			value = toUint32(value) * uint32(metric.Scale)
		case "uint64":
			value = toUint64(value) * uint64(metric.Scale)
		default:
			value = toFloat64(value) * metric.Scale
		}
	}

	if len(prefix) > 0 {
		h.printValue(os.Stdout, prefix+"."+metric.Name, value, now)
	} else {
		h.printValue(os.Stdout, metric.Name, value, now)
	}
}

func (h *MackerelPlugin) formatValuesWithWildcard(prefix string, metric Metrics, stat *map[string]interface{}, lastStat *map[string]interface{}, now time.Time, lastTime time.Time) {
	regexpStr := `\A` + prefix + "." + metric.Name
	regexpStr = strings.Replace(regexpStr, ".", "\\.", -1)
	regexpStr = strings.Replace(regexpStr, "*", "[-a-zA-Z0-9_]+", -1)
	regexpStr = strings.Replace(regexpStr, "#", "[-a-zA-Z0-9_]+", -1)
	re, err := regexp.Compile(regexpStr)
	if err != nil {
		log.Fatalln("Failed to compile regexp: ", err)
	}
	for k, _ := range *stat {
		if re.MatchString(k) {
			metricEach := metric
			metricEach.Name = k
			h.formatValues("", metricEach, stat, lastStat, now, lastTime)
		}
	}
}

func (h *MackerelPlugin) Run() {
	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		h.OutputDefinitions()
	} else {
		h.OutputValues()
	}
}

func (h *MackerelPlugin) OutputValues() {
	now := time.Now()
	stat, err := h.FetchMetrics()
	if err != nil {
		log.Fatalln("OutputValues: ", err)
	}

	lastStat, lastTime, err := h.fetchLastValues()
	if err != nil {
		log.Println("fetchLastValues (ignore):", err)
	}

	for key, graph := range h.GraphDefinition() {
		for _, metric := range graph.Metrics {
			if strings.ContainsAny(key+metric.Name, "*#") {
				h.formatValuesWithWildcard(key, metric, &stat, &lastStat, now, lastTime)
			} else {
				h.formatValues(key, metric, &stat, &lastStat, now, lastTime)
			}
		}
	}

	err = h.saveValues(stat, now)
	if err != nil {
		log.Fatalf("saveValues: ", err)
	}
}

type GraphDef struct {
	Graphs map[string]Graphs `json:"graphs"`
}

func (h *MackerelPlugin) OutputDefinitions() {
	fmt.Println("# mackerel-agent-plugin")
	var graphs GraphDef
	graphs.Graphs = h.GraphDefinition()

	b, err := json.Marshal(graphs)
	if err != nil {
		log.Fatalln("OutputDefinitions: ", err)
	}
	fmt.Println(string(b))
}

func toUint32(value interface{}) uint32 {
	var ret uint32
	switch value.(type) {
	case uint32:
		ret = value.(uint32)
	case uint64:
		ret = uint32(value.(uint64))
	case float64:
		ret = uint32(value.(float64))
	case string:
		v, err := strconv.ParseUint(value.(string), 10, 32)
		if err == nil {
			ret = uint32(v)
		}
	}
	return ret
}

func toUint64(value interface{}) uint64 {
	var ret uint64
	switch value.(type) {
	case uint32:
		ret = uint64(value.(uint32))
	case uint64:
		ret = value.(uint64)
	case float64:
		ret = uint64(value.(float64))
	case string:
		ret, _ = strconv.ParseUint(value.(string), 10, 64)
	}
	return ret
}

func toFloat64(value interface{}) float64 {
	var ret float64
	switch value.(type) {
	case uint32:
		ret = float64(value.(uint32))
	case uint64:
		ret = float64(value.(uint64))
	case float64:
		ret = value.(float64)
	case string:
		ret, _ = strconv.ParseFloat(value.(string), 64)
	}
	return ret
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"log"
	"net"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/miekg/dns"
)

type Service struct {
	Name    string
	Image   string
	Ip      net.IP
	Ttl     int
	Aliases []string
}

func NewService() (s *Service) {
	s = &Service{Ttl: -1}
	return
}

type ServiceListProvider interface {
	AddService(string, Service)
	RemoveService(string) error
	GetService(string) (Service, error)
	GetAllServices() map[string]Service
}

type DNSServer struct {
	config   *Config
	server   *dns.Server
	services map[string]*Service
	lock     *sync.RWMutex
}

func NewDNSServer(c *Config) *DNSServer {
	s := &DNSServer{
		config:   c,
		services: make(map[string]*Service),
		lock:     &sync.RWMutex{},
	}

	mux := dns.NewServeMux()
	mux.HandleFunc(c.domain[len(c.domain)-1]+".", s.handleRequest)
	mux.HandleFunc("arpa.", s.handleRequest)
	mux.HandleFunc(".", s.handleOrForward)

	s.server = &dns.Server{Addr: c.dnsAddr, Net: "udp", Handler: mux}

	return s
}

func (s *DNSServer) Start() error {
	return s.server.ListenAndServe()
}

func (s *DNSServer) Stop() {
	s.server.Shutdown()
}

func (s *DNSServer) AddService(id string, service Service) {
	defer s.lock.Unlock()
	s.lock.Lock()

	id = s.getExpandedID(id)
	s.services[id] = &service

	if s.config.verbose {
		log.Println("Added service:", id, service)
	}
}

func (s *DNSServer) RemoveService(id string) error {
	defer s.lock.Unlock()
	s.lock.Lock()

	id = s.getExpandedID(id)
	if _, ok := s.services[id]; !ok {
		return errors.New("No such service: " + id)
	}

	delete(s.services, id)

	if s.config.verbose {
		log.Println("Stopped service:", id)
	}

	return nil
}

func (s *DNSServer) GetService(id string) (Service, error) {
	defer s.lock.RUnlock()
	s.lock.RLock()

	id = s.getExpandedID(id)
	if s, ok := s.services[id]; ok {
		return *s, nil
	}
	// Check for a pa
	return *new(Service), errors.New("No such service: " + id)
}

func (s *DNSServer) GetAllServices() map[string]Service {
	defer s.lock.RUnlock()
	s.lock.RLock()

	list := make(map[string]Service, len(s.services))
	for id, service := range s.services {
		list[id] = *service
	}

	return list
}

func (s *DNSServer) listDomains(service *Service) chan string {
	c := make(chan string)

	go func() {

		if service.Image == "" {
			c <- service.Name + "." + s.config.domain.String() + "."
		} else {
			domain := service.Image + "." + s.config.domain.String() + "."

			c <- domain
			c <- service.Name + "." + domain
		}

		for _, alias := range service.Aliases {
			c <- alias + "."
		}

		close(c)
	}()

	return c
}

func (s *DNSServer) handleOrForward(w dns.ResponseWriter, r *dns.Msg) {
	m := s.doHandle(w, r)

	// We could build a response, let's return it
	if len(m.Answer) != 0 {
		w.WriteMsg(m)
		return
	}

	// Otherwise just forward the request to another server
	c := new(dns.Client)
	if in, _, err := c.Exchange(r, s.config.nameserver); err != nil {
		log.Print(err)
		w.WriteMsg(new(dns.Msg))
	} else {
		w.WriteMsg(in)
	}
}

func (s *DNSServer) handleRequest(w dns.ResponseWriter, r *dns.Msg) {
	m := s.doHandle(w, r)

	// We could NOT build a response
	if len(m.Answer) == 0 && len(m.Ns) == 0 {
		m.Ns = s.createSOA()
		m.SetRcode(r, dns.RcodeNameError)
	}

	w.WriteMsg(m)
}

func (s *DNSServer) doHandle(w dns.ResponseWriter, r *dns.Msg) *dns.Msg {
	m := new(dns.Msg)
	m.SetReply(r)

	// Send empty response for empty requests
	if len(r.Question) == 0 {
		m.Ns = s.createSOA()
		return m
	}

	switch r.Question[0].Qtype {
	case dns.TypePTR:
		s.handlePTRRequest(r, m)
	case dns.TypeA:
		s.handleARequest(r, m)
	case dns.TypeSOA:
		m.Answer = s.createSOA()
	default:
		m.Ns = s.createSOA()
		m.SetRcode(r, dns.RcodeNotImplemented)
	}

	return m
}

func (s *DNSServer) handlePTRRequest(r *dns.Msg, m *dns.Msg) {
	m.Answer = make([]dns.RR, 0, 2)
	query := r.Question[0].Name

	if query[len(query)-1] == '.' {
		query = query[:len(query)-1]
	}

	for service := range s.queryIp(query) {
		var ttl int
		if service.Ttl != -1 {
			ttl = service.Ttl
		} else {
			ttl = s.config.ttl
		}

		for domain := range s.listDomains(service) {
			rr := new(dns.PTR)
			rr.Hdr = dns.RR_Header{
				Name:   r.Question[0].Name,
				Rrtype: dns.TypePTR,
				Class:  dns.ClassINET,
				Ttl:    uint32(ttl),
			}
			rr.Ptr = domain

			m.Answer = append(m.Answer, rr)
		}
	}
}

func (s *DNSServer) handleARequest(r *dns.Msg, m *dns.Msg) {
	m.Answer = make([]dns.RR, 0, 2)
	query := r.Question[0].Name

	if query[len(query)-1] == '.' {
		query = query[:len(query)-1]
	}

	for service := range s.queryServices(query) {
		rr := new(dns.A)

		var ttl int
		if service.Ttl != -1 {
			ttl = service.Ttl
		} else {
			ttl = s.config.ttl
		}

		rr.Hdr = dns.RR_Header{
			Name:   r.Question[0].Name,
			Rrtype: dns.TypeA,
			Class:  dns.ClassINET,
			Ttl:    uint32(ttl),
		}
		rr.A = service.Ip
		m.Answer = append(m.Answer, rr)
	}
}

func (s *DNSServer) queryIp(query string) chan *Service {
	c := make(chan *Service, 3)
	reversedIp := strings.TrimSuffix(query, ".in-addr.arpa")
	ip := strings.Join(reverse(strings.Split(reversedIp, ".")), ".")

	go func() {
		defer s.lock.RUnlock()
		s.lock.RLock()

		for _, service := range s.services {
			if service.Ip.String() == ip {
				c <- service
			}
		}

		close(c)
	}()

	return c
}

func (s *DNSServer) queryServices(query string) chan *Service {
	c := make(chan *Service, 3)

	go func() {
		query := strings.Split(strings.ToLower(query), ".")

		defer s.lock.RUnlock()
		s.lock.RLock()

		for _, service := range s.services {
			// create the name for this service, skip empty strings
			test := []string{}
			// todo: add some cache to avoid calculating this every time
			if len(service.Name) > 0 {
				test = append(test, strings.Split(service.Name, ".")...)
			}

			if len(service.Image) > 0 {
				test = append(test, strings.Split(service.Image, ".")...)
			}

			test = append(test, s.config.domain...)

			if isPrefixQuery(query, test) {
				c <- service
			}

			// check aliases
			for _, alias := range service.Aliases {
				if isPrefixQuery(query, strings.Split(alias, ".")) {
					c <- service
				}
			}
		}

		close(c)

	}()

	return c

}

// Checks for a partial match for container SHA and outputs it if found.
func (s *DNSServer) getExpandedID(in string) (out string) {
	out = in

	// Hard to make a judgement on small image names.
	if len(in) < 4 {
		return
	}

	if isHex, _ := regexp.MatchString("^[0-9a-f]+$", in); !isHex {
		return
	}

	for id := range s.services {
		if len(id) == 64 {
			if isHex, _ := regexp.MatchString("^[0-9a-f]+$", id); isHex {
				if strings.HasPrefix(id, in) {
					out = id
					return
				}
			}
		}
	}
	return
}

// Ttl is used from config so that not-found result responses are not cached
// for a long time. The other defaults left as is(skydns source) because they
// do not have an use case in this situation.
func (s *DNSServer) createSOA() []dns.RR {
	dom := dns.Fqdn(s.config.domain.String() + ".")
	soa := &dns.SOA{
		Hdr: dns.RR_Header{
			Name:   dom,
			Rrtype: dns.TypeSOA,
			Class:  dns.ClassINET,
			Ttl:    uint32(s.config.ttl)},
		Ns:      "dnsdock." + dom,
		Mbox:    "dnsdock.dnsdock." + dom,
		Serial:  uint32(time.Now().Truncate(time.Hour).Unix()),
		Refresh: 28800,
		Retry:   7200,
		Expire:  604800,
		Minttl:  uint32(s.config.ttl),
	}
	return []dns.RR{soa}
}

// isPrefixQuery is used to determine whether "query" is a potential prefix
// query for "name". It allows for wildcards (*) in the query. However is makes
// one exception to accomodate the desired behavior we wish from dnsdock,
// namely, the query may be longer than "name" and still be a valid prefix
// query for "name".
// Examples:
//   foo.bar.baz.qux is a valid query for bar.baz.qux (longer prefix is okay)
//   foo.*.baz.qux   is a valid query for bar.baz.qux (wildcards okay)
//   *.baz.qux       is a valid query for baz.baz.qux (wildcard prefix okay)
func isPrefixQuery(query, name []string) bool {
	for i, j := len(query)-1, len(name)-1; i >= 0 && j >= 0; i, j = i-1, j-1 {
		if query[i] != name[j] && query[i] != "*" {
			return false
		}
	}
	return true
}

func reverse(input []string) []string {
	if len(input) == 0 {
		return input
	}

	return append(reverse(input[1:]), input[0])
}
<commit_msg>Added handling of MX lookups<commit_after>package main

import (
	"errors"
	"log"
	"net"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/miekg/dns"
)

type Service struct {
	Name    string
	Image   string
	Ip      net.IP
	Ttl     int
	Aliases []string
}

func NewService() (s *Service) {
	s = &Service{Ttl: -1}
	return
}

type ServiceListProvider interface {
	AddService(string, Service)
	RemoveService(string) error
	GetService(string) (Service, error)
	GetAllServices() map[string]Service
}

type DNSServer struct {
	config   *Config
	server   *dns.Server
	services map[string]*Service
	lock     *sync.RWMutex
}

func NewDNSServer(c *Config) *DNSServer {
	s := &DNSServer{
		config:   c,
		services: make(map[string]*Service),
		lock:     &sync.RWMutex{},
	}

	mux := dns.NewServeMux()
	mux.HandleFunc(c.domain[len(c.domain)-1]+".", s.handleRequest)
	mux.HandleFunc("arpa.", s.handleRequest)
	mux.HandleFunc(".", s.handleOrForward)

	s.server = &dns.Server{Addr: c.dnsAddr, Net: "udp", Handler: mux}

	return s
}

func (s *DNSServer) Start() error {
	return s.server.ListenAndServe()
}

func (s *DNSServer) Stop() {
	s.server.Shutdown()
}

func (s *DNSServer) AddService(id string, service Service) {
	defer s.lock.Unlock()
	s.lock.Lock()

	id = s.getExpandedID(id)
	s.services[id] = &service

	if s.config.verbose {
		log.Println("Added service:", id, service)
	}
}

func (s *DNSServer) RemoveService(id string) error {
	defer s.lock.Unlock()
	s.lock.Lock()

	id = s.getExpandedID(id)
	if _, ok := s.services[id]; !ok {
		return errors.New("No such service: " + id)
	}

	delete(s.services, id)

	if s.config.verbose {
		log.Println("Stopped service:", id)
	}

	return nil
}

func (s *DNSServer) GetService(id string) (Service, error) {
	defer s.lock.RUnlock()
	s.lock.RLock()

	id = s.getExpandedID(id)
	if s, ok := s.services[id]; ok {
		return *s, nil
	}
	// Check for a pa
	return *new(Service), errors.New("No such service: " + id)
}

func (s *DNSServer) GetAllServices() map[string]Service {
	defer s.lock.RUnlock()
	s.lock.RLock()

	list := make(map[string]Service, len(s.services))
	for id, service := range s.services {
		list[id] = *service
	}

	return list
}

func (s *DNSServer) listDomains(service *Service) chan string {
	c := make(chan string)

	go func() {

		if service.Image == "" {
			c <- service.Name + "." + s.config.domain.String() + "."
		} else {
			domain := service.Image + "." + s.config.domain.String() + "."

			c <- domain
			c <- service.Name + "." + domain
		}

		for _, alias := range service.Aliases {
			c <- alias + "."
		}

		close(c)
	}()

	return c
}

func (s *DNSServer) handleOrForward(w dns.ResponseWriter, r *dns.Msg) {
	m := s.doHandle(w, r)

	// We could build a response, let's return it
	if len(m.Answer) != 0 {
		w.WriteMsg(m)
		return
	}

	// Otherwise just forward the request to another server
	c := new(dns.Client)
	if in, _, err := c.Exchange(r, s.config.nameserver); err != nil {
		log.Print(err)
		w.WriteMsg(new(dns.Msg))
	} else {
		w.WriteMsg(in)
	}
}

func (s *DNSServer) handleRequest(w dns.ResponseWriter, r *dns.Msg) {
	m := s.doHandle(w, r)

	// We could NOT build a response
	if len(m.Answer) == 0 && len(m.Ns) == 0 {
		m.Ns = s.createSOA()
		m.SetRcode(r, dns.RcodeNameError)
	}

	w.WriteMsg(m)
}

func (s *DNSServer) doHandle(w dns.ResponseWriter, r *dns.Msg) *dns.Msg {
	m := new(dns.Msg)
	m.SetReply(r)

	// Send empty response for empty requests
	if len(r.Question) == 0 {
		m.Ns = s.createSOA()
		return m
	}

	switch r.Question[0].Qtype {
	case dns.TypePTR:
		s.handlePTRRequest(r, m)
	case dns.TypeMX:
		s.handleMXRequest(r, m)
	case dns.TypeA:
		s.handleARequest(r, m)
	case dns.TypeSOA:
		m.Answer = s.createSOA()
	default:
		m.Ns = s.createSOA()
		m.SetRcode(r, dns.RcodeNotImplemented)
	}

	return m
}

func (s *DNSServer) handleMXRequest(r *dns.Msg, m *dns.Msg) {
	m.Answer = make([]dns.RR, 0, 2)
	query := r.Question[0].Name

	if query[len(query)-1] == '.' {
		query = query[:len(query)-1]
	}

	for service := range s.queryIp(query) {
		rr := new(dns.MX)

		var ttl int
		if service.Ttl != -1 {
			ttl = service.Ttl
		} else {
			ttl = s.config.ttl
		}

		rr.Hdr = dns.RR_Header{
			Name:   r.Question[0].Name,
			Rrtype: dns.TypeMX,
			Class:  dns.ClassINET,
			Ttl:    uint32(ttl),
		}
		rr.Mx = r.Question[0].Name
		m.Answer = append(m.Answer, rr)
	}
}

func (s *DNSServer) handlePTRRequest(r *dns.Msg, m *dns.Msg) {
	m.Answer = make([]dns.RR, 0, 2)
	query := r.Question[0].Name

	if query[len(query)-1] == '.' {
		query = query[:len(query)-1]
	}

	for service := range s.queryIp(query) {
		var ttl int
		if service.Ttl != -1 {
			ttl = service.Ttl
		} else {
			ttl = s.config.ttl
		}

		for domain := range s.listDomains(service) {
			rr := new(dns.PTR)
			rr.Hdr = dns.RR_Header{
				Name:   r.Question[0].Name,
				Rrtype: dns.TypePTR,
				Class:  dns.ClassINET,
				Ttl:    uint32(ttl),
			}
			rr.Ptr = domain

			m.Answer = append(m.Answer, rr)
		}
	}
}

func (s *DNSServer) handleARequest(r *dns.Msg, m *dns.Msg) {
	m.Answer = make([]dns.RR, 0, 2)
	query := r.Question[0].Name

	if query[len(query)-1] == '.' {
		query = query[:len(query)-1]
	}

	for service := range s.queryServices(query) {
		rr := new(dns.A)

		var ttl int
		if service.Ttl != -1 {
			ttl = service.Ttl
		} else {
			ttl = s.config.ttl
		}

		rr.Hdr = dns.RR_Header{
			Name:   r.Question[0].Name,
			Rrtype: dns.TypeA,
			Class:  dns.ClassINET,
			Ttl:    uint32(ttl),
		}
		rr.A = service.Ip
		m.Answer = append(m.Answer, rr)
	}
}

func (s *DNSServer) queryIp(query string) chan *Service {
	c := make(chan *Service, 3)
	reversedIp := strings.TrimSuffix(query, ".in-addr.arpa")
	ip := strings.Join(reverse(strings.Split(reversedIp, ".")), ".")

	go func() {
		defer s.lock.RUnlock()
		s.lock.RLock()

		for _, service := range s.services {
			if service.Ip.String() == ip {
				c <- service
			}
		}

		close(c)
	}()

	return c
}

func (s *DNSServer) queryServices(query string) chan *Service {
	c := make(chan *Service, 3)

	go func() {
		query := strings.Split(strings.ToLower(query), ".")

		defer s.lock.RUnlock()
		s.lock.RLock()

		for _, service := range s.services {
			// create the name for this service, skip empty strings
			test := []string{}
			// todo: add some cache to avoid calculating this every time
			if len(service.Name) > 0 {
				test = append(test, strings.Split(service.Name, ".")...)
			}

			if len(service.Image) > 0 {
				test = append(test, strings.Split(service.Image, ".")...)
			}

			test = append(test, s.config.domain...)

			if isPrefixQuery(query, test) {
				c <- service
			}

			// check aliases
			for _, alias := range service.Aliases {
				if isPrefixQuery(query, strings.Split(alias, ".")) {
					c <- service
				}
			}
		}

		close(c)

	}()

	return c

}

// Checks for a partial match for container SHA and outputs it if found.
func (s *DNSServer) getExpandedID(in string) (out string) {
	out = in

	// Hard to make a judgement on small image names.
	if len(in) < 4 {
		return
	}

	if isHex, _ := regexp.MatchString("^[0-9a-f]+$", in); !isHex {
		return
	}

	for id := range s.services {
		if len(id) == 64 {
			if isHex, _ := regexp.MatchString("^[0-9a-f]+$", id); isHex {
				if strings.HasPrefix(id, in) {
					out = id
					return
				}
			}
		}
	}
	return
}

// Ttl is used from config so that not-found result responses are not cached
// for a long time. The other defaults left as is(skydns source) because they
// do not have an use case in this situation.
func (s *DNSServer) createSOA() []dns.RR {
	dom := dns.Fqdn(s.config.domain.String() + ".")
	soa := &dns.SOA{
		Hdr: dns.RR_Header{
			Name:   dom,
			Rrtype: dns.TypeSOA,
			Class:  dns.ClassINET,
			Ttl:    uint32(s.config.ttl)},
		Ns:      "dnsdock." + dom,
		Mbox:    "dnsdock.dnsdock." + dom,
		Serial:  uint32(time.Now().Truncate(time.Hour).Unix()),
		Refresh: 28800,
		Retry:   7200,
		Expire:  604800,
		Minttl:  uint32(s.config.ttl),
	}
	return []dns.RR{soa}
}

// isPrefixQuery is used to determine whether "query" is a potential prefix
// query for "name". It allows for wildcards (*) in the query. However is makes
// one exception to accomodate the desired behavior we wish from dnsdock,
// namely, the query may be longer than "name" and still be a valid prefix
// query for "name".
// Examples:
//   foo.bar.baz.qux is a valid query for bar.baz.qux (longer prefix is okay)
//   foo.*.baz.qux   is a valid query for bar.baz.qux (wildcards okay)
//   *.baz.qux       is a valid query for baz.baz.qux (wildcard prefix okay)
func isPrefixQuery(query, name []string) bool {
	for i, j := len(query)-1, len(name)-1; i >= 0 && j >= 0; i, j = i-1, j-1 {
		if query[i] != name[j] && query[i] != "*" {
			return false
		}
	}
	return true
}

func reverse(input []string) []string {
	if len(input) == 0 {
		return input
	}

	return append(reverse(input[1:]), input[0])
}
<|endoftext|>


<commit_before>package resreporter

import (
	"bufio"
	"fmt"
	"log"
	"os"
)

const meminfo = "/proc/meminfo"

var Memory = memory{}

type memory struct{}

func (memory) Used() (used uint64, total uint64) {
	fd, err := os.Open(meminfo)
	if err != nil {
		//XXX Should this use metafora's logger somehow?
		log.Printf("Error reading free memory via "+meminfo+": %v", err)

		// Effectively disable the balancer since an error happened
		return 0, 100
	}
	defer fd.Close()

	s := bufio.NewScanner(fd)
	foundFree, foundCache, foundBuf := false, false, false
	var cache uint64
	var buffered uint64
	var free uint64
	for s.Scan() {
		if total > 0 && foundFree && foundCache && foundBuf {
			break
		}
		if total == 0 {
			if n, _ := fmt.Sscanf(s.Text(), "MemTotal:%d", &total); n == 1 {
				continue
			}
		}
		if foundFree {
			if n, _ := fmt.Sscanf(s.Text(), "MemFree:%d", &free); n == 1 {
				continue
			}
		}
		if !foundCache {
			if n, _ := fmt.Sscanf(s.Text(), "Cached:%d", &cache); n == 1 {
				foundCache = true
				continue
			}
		}
		if !foundBuf {
			if n, _ := fmt.Sscanf(s.Text(), "Buffers:%d", &buffered); n == 1 {
				foundBuf = true
				continue
			}
		}
	}
	if err := s.Err(); err != nil {
		//XXX Should this use metafora's logger somehow?
		log.Printf("Error reading free memory via "+meminfo+": %v", err)

		// Effectively disable the balancer since an error happened
		return 0, 100
	}

	return total - (free + buffered + cache), total
}

func (memory) String() string { return "kB" }
<commit_msg>Fix the worst logging in all of metafora<commit_after>package resreporter

import (
	"bufio"
	"fmt"
	"os"

	"github.com/lytics/metafora"
)

const meminfo = "/proc/meminfo"

var Memory = memory{}

type memory struct{}

func (memory) Used() (used uint64, total uint64) {
	fd, err := os.Open(meminfo)
	if err != nil {
		metafora.Errorf("Error reading free memory via "+meminfo+": %v", err)

		// Effectively disable the balancer since an error happened
		return 0, 100
	}
	defer fd.Close()

	s := bufio.NewScanner(fd)
	foundFree, foundCache, foundBuf := false, false, false
	var cache uint64
	var buffered uint64
	var free uint64
	for s.Scan() {
		if total > 0 && foundFree && foundCache && foundBuf {
			break
		}
		if total == 0 {
			if n, _ := fmt.Sscanf(s.Text(), "MemTotal:%d", &total); n == 1 {
				continue
			}
		}
		if foundFree {
			if n, _ := fmt.Sscanf(s.Text(), "MemFree:%d", &free); n == 1 {
				continue
			}
		}
		if !foundCache {
			if n, _ := fmt.Sscanf(s.Text(), "Cached:%d", &cache); n == 1 {
				foundCache = true
				continue
			}
		}
		if !foundBuf {
			if n, _ := fmt.Sscanf(s.Text(), "Buffers:%d", &buffered); n == 1 {
				foundBuf = true
				continue
			}
		}
	}
	if err := s.Err(); err != nil {
		metafora.Errorf("Error reading free memory via "+meminfo+": %v", err)

		// Effectively disable the balancer since an error happened
		return 0, 100
	}

	return total - (free + buffered + cache), total
}

func (memory) String() string { return "kB" }
<|endoftext|>


<commit_before>package goat

import (
	"bencode"
	"fmt"
	"strconv"
)

// Tracker announce request
func TrackerAnnounce(user UserRecord, query map[string]string, resChan chan []byte) {
	// Store announce information in struct
	announce := MapToAnnounceLog(query)

	// Request to store announce
	go announce.Save()

	Static.LogChan <- fmt.Sprintf("announce: [ip: %s, port:%d]", announce.Ip, announce.Port)
	Static.LogChan <- fmt.Sprintf("announce: [info_hash: %s]", announce.InfoHash)
	Static.LogChan <- fmt.Sprintf("announce: [up: %d, down: %d, left: %d]", announce.Uploaded, announce.Downloaded, announce.Left)

	// Only report announce when needed
	if announce.Event != "" {
		Static.LogChan <- fmt.Sprintf("announce: [event: %s]", announce.Event)
	}

	// Check for a matching file via info_hash
	file := new(FileRecord).Load(announce.InfoHash, "info_hash")
	if file == (FileRecord{}) {
		// Torrent is not currently registered
		resChan <- TrackerError("Unregistered torrent")

		// Create an entry in file table for this hash, but mark it as unverified
		file.InfoHash = announce.InfoHash
		file.Verified = false
		file.Completed = 0
		go file.Save()
		return
	}

	// Ensure file is verified, meaning we will permit tracking of it
	if !file.Verified {
		resChan <- TrackerError("Unverified torrent")
		return
	}

	// Check existing record for this user with this file
	fileUser := new(FileUserRecord).Load(file.Id, user.Id)

	// New user, starting torrent
	if announce.Event == "started" && fileUser == (FileUserRecord{}) {
		// Create new relationship
		fileUser.FileId = file.Id
		fileUser.UserId = user.Id
		fileUser.Active = true
		fileUser.Announced = 1

		// If announce reports 0 left, but no existing record, user is probably the initial seeder
		if announce.Left == 0 {
			fileUser.Completed = true
		} else {
			fileUser.Completed = false
		}

		// Track the initial uploaded, download, and left values
		// NOTE: clients report absolute values, so delta should NEVER be calculated for these
		fileUser.Uploaded = announce.Uploaded
		fileUser.Downloaded = announce.Downloaded
		fileUser.Left = announce.Left
	} else {
		// Else, pre-existing record, so update
		// Event "stopped", mark as inactive
		// NOTE: likely only reported by clients which are actively seeding, NOT when stopped during leeching
		if announce.Event == "stopped" {
			fileUser.Active = false
		} else {
			// Else, "started", "completed", or no status, mark as active
			fileUser.Active = true
		}

		// Check for completion
		// Could be from a peer stating completed, or a seed reporting 0 left
		if announce.Event == "completed" || announce.Left == 0 {
			fileUser.Completed = true

			// If status completed, mark file as completed by another user
			if announce.Event == "completed" {
				file.Completed = file.Completed + 1
			}
		} else {
			fileUser.Completed = false
		}

		// Add an announce
		fileUser.Announced = fileUser.Announced + 1

		// Store latest statistics, but do so in a sane way (no removing upload/download, no adding left)
		// NOTE: clients report absolute values, so delta should NEVER be calculated for these
		// NOTE: It is also worth noting that if a client re-downloads a file they have previously downloaded,
		// but the FileUserRecord relationship is not cleared, they will essentially get a "free" download, with
		// no extra download penalty to their share ratio
		// For the time being, this behavior will be expected and acceptable
		if announce.Uploaded > fileUser.Uploaded {
			fileUser.Uploaded = announce.Uploaded
		}
		if announce.Downloaded > fileUser.Downloaded {
			fileUser.Downloaded = announce.Downloaded
		}
		if announce.Left < fileUser.Left {
			fileUser.Left = announce.Left
		}
	}

	// Check for numwant parameter, return up to that number of peers
	// Default is 50 per protocol
	numwant := 50
	if _, ok := query["numwant"]; ok {
		// Verify numwant is an integer
		num, err := strconv.Atoi(query["numwant"])
		if err == nil {
			numwant = num
		}
	}

	// Tracker announce response, with generated peerlist of length numwant, excluding this user
	resChan <- bencode.EncDictMap(map[string][]byte{
		"interval":     bencode.EncInt(RandRange(3200, 4000)),
		"min interval": bencode.EncInt(1800),
		"peers":        bencode.EncBytes(file.PeerList(query["ip"], numwant)),
	})

	// Update records AFTER user gets their response
	// NOTE: not goroutines because these occur after the client gets a response anyway, and because these
	// can create some interesting race conditions if not updated synchronously
	file.Save()
	fileUser.Save()
	user.Save()
}

// Report a bencoded []byte response as specified by input string
func TrackerError(err string) []byte {
	return bencode.EncDictMap(map[string][]byte{
		"failure reason": bencode.EncString(err),
		"interval":       bencode.EncInt(RandRange(3200, 4000)),
		"min interval":   bencode.EncInt(1800),
	})
}
<commit_msg>Update struct saving routines, now that race conditions are gone<commit_after>package goat

import (
	"bencode"
	"fmt"
	"strconv"
)

// Tracker announce request
func TrackerAnnounce(user UserRecord, query map[string]string, resChan chan []byte) {
	// Store announce information in struct
	announce := MapToAnnounceLog(query)

	// Request to store announce
	go announce.Save()

	Static.LogChan <- fmt.Sprintf("announce: [ip: %s, port:%d]", announce.Ip, announce.Port)
	Static.LogChan <- fmt.Sprintf("announce: [info_hash: %s]", announce.InfoHash)
	Static.LogChan <- fmt.Sprintf("announce: [up: %d, down: %d, left: %d]", announce.Uploaded, announce.Downloaded, announce.Left)

	// Only report announce when needed
	if announce.Event != "" {
		Static.LogChan <- fmt.Sprintf("announce: [event: %s]", announce.Event)
	}

	// Check for a matching file via info_hash
	file := new(FileRecord).Load(announce.InfoHash, "info_hash")
	if file == (FileRecord{}) {
		// Torrent is not currently registered
		resChan <- TrackerError("Unregistered torrent")

		// Create an entry in file table for this hash, but mark it as unverified
		file.InfoHash = announce.InfoHash
		file.Verified = false
		file.Completed = 0
		go file.Save()
		return
	}

	// Ensure file is verified, meaning we will permit tracking of it
	if !file.Verified {
		resChan <- TrackerError("Unverified torrent")
		return
	}

	// Check existing record for this user with this file
	fileUser := new(FileUserRecord).Load(file.Id, user.Id)

	// New user, starting torrent
	if announce.Event == "started" && fileUser == (FileUserRecord{}) {
		// Create new relationship
		fileUser.FileId = file.Id
		fileUser.UserId = user.Id
		fileUser.Active = true
		fileUser.Announced = 1

		// If announce reports 0 left, but no existing record, user is probably the initial seeder
		if announce.Left == 0 {
			fileUser.Completed = true
		} else {
			fileUser.Completed = false
		}

		// Track the initial uploaded, download, and left values
		// NOTE: clients report absolute values, so delta should NEVER be calculated for these
		fileUser.Uploaded = announce.Uploaded
		fileUser.Downloaded = announce.Downloaded
		fileUser.Left = announce.Left
	} else {
		// Else, pre-existing record, so update
		// Event "stopped", mark as inactive
		// NOTE: likely only reported by clients which are actively seeding, NOT when stopped during leeching
		if announce.Event == "stopped" {
			fileUser.Active = false
		} else {
			// Else, "started", "completed", or no status, mark as active
			fileUser.Active = true
		}

		// Check for completion
		// Could be from a peer stating completed, or a seed reporting 0 left
		if announce.Event == "completed" || announce.Left == 0 {
			fileUser.Completed = true

			// If status completed, mark file as completed by another user
			if announce.Event == "completed" {
				file.Completed = file.Completed + 1
				go file.Save()
			}
		} else {
			fileUser.Completed = false
		}

		// Add an announce
		fileUser.Announced = fileUser.Announced + 1

		// Store latest statistics, but do so in a sane way (no removing upload/download, no adding left)
		// NOTE: clients report absolute values, so delta should NEVER be calculated for these
		// NOTE: It is also worth noting that if a client re-downloads a file they have previously downloaded,
		// but the FileUserRecord relationship is not cleared, they will essentially get a "free" download, with
		// no extra download penalty to their share ratio
		// For the time being, this behavior will be expected and acceptable
		if announce.Uploaded > fileUser.Uploaded {
			fileUser.Uploaded = announce.Uploaded
		}
		if announce.Downloaded > fileUser.Downloaded {
			fileUser.Downloaded = announce.Downloaded
		}
		if announce.Left < fileUser.Left {
			fileUser.Left = announce.Left
		}
	}

	// Update file/user relationship record
	go fileUser.Save()

	// Check for numwant parameter, return up to that number of peers
	// Default is 50 per protocol
	numwant := 50
	if _, ok := query["numwant"]; ok {
		// Verify numwant is an integer
		num, err := strconv.Atoi(query["numwant"])
		if err == nil {
			numwant = num
		}
	}

	// Tracker announce response, with generated peerlist of length numwant, excluding this user
	resChan <- bencode.EncDictMap(map[string][]byte{
		"interval":     bencode.EncInt(RandRange(3200, 4000)),
		"min interval": bencode.EncInt(1800),
		"peers":        bencode.EncBytes(file.PeerList(query["ip"], numwant)),
	})
}

// Report a bencoded []byte response as specified by input string
func TrackerError(err string) []byte {
	return bencode.EncDictMap(map[string][]byte{
		"failure reason": bencode.EncString(err),
		"interval":       bencode.EncInt(RandRange(3200, 4000)),
		"min interval":   bencode.EncInt(1800),
	})
}
<|endoftext|>


<commit_before>// 5 october 2014
package main

import (
	"net/url"
	"os"
)

type HelpSource interface {
	Name() string
	Books() []Topic
	Orphans() []Topic
//	SharedAssets() []Asset
	Lookup(url *url.URL) Topic
}

type Topic interface {
	Name() string
	Prepare() (string, error)
	Parent() Topic
	Children() []Topic
	Source() HelpSource
	Less(t Topic) bool
}

type TopicSorter []Topic

func (t TopicSorter) Len() int { return len(t) }
func (t TopicSorter) Less(i, j int) bool { return t[i].Less(t[j]) }
func (t TopicSorter) Swap(i, j int) { t[i], t[j] = t[j], t[i] }

var Library []Topic

func LoadMSHILibrary() {
	m, err := OpenMSHI(os.Args[1])
	if err != nil {
		// TODO
		panic(err)
	}
	for _, b := range m.Books() {
		Library = append(Library, b)
	}
}

func LoadLibraries() {
	LoadMSHILibrary()
}
<commit_msg>Paved the way for more documentation types.<commit_after>// 5 october 2014
package main

import (
	"net/url"
	"os"
	"path/filepath"
)

type HelpSource interface {
	Name() string
	Books() []Topic
	Orphans() []Topic
//	SharedAssets() []Asset
	Lookup(url *url.URL) Topic
}

type Topic interface {
	Name() string
	Prepare() (string, error)
	Parent() Topic
	Children() []Topic
	Source() HelpSource
	Less(t Topic) bool
}

type TopicSorter []Topic

func (t TopicSorter) Len() int { return len(t) }
func (t TopicSorter) Less(i, j int) bool { return t[i].Less(t[j]) }
func (t TopicSorter) Swap(i, j int) { t[i], t[j] = t[j], t[i] }

var Library []Topic

func LoadMSHILibrary(dir string) {
	m, err := OpenMSHI(filepath.Join(dir, "mshi"))
	if err != nil {
		// TODO
		panic(err)
	}
	for _, b := range m.Books() {
		Library = append(Library, b)
	}
}

func LoadLibraries() {
	LoadMSHILibrary(os.Args[1])
}
<|endoftext|>


<commit_before>package libvirt

import (
	"github.com/alexzorin/libvirt-go"
	"github.com/mistifyio/mistify-agent/rpc"
	"github.com/mistifyio/mistify-agent/client"
	"github.com/mistifyio/mistify-agent/log"
	"syscall"
	"net/http"
	"fmt"
)

var StateNames = map[int]string{
	libvirt.VIR_DOMAIN_NOSTATE: "No State",
	libvirt.VIR_DOMAIN_RUNNING: "Running",
	libvirt.VIR_DOMAIN_BLOCKED: "Blocked",
	libvirt.VIR_DOMAIN_PAUSED: "Paused",
	libvirt.VIR_DOMAIN_SHUTDOWN: "Shutdown",
	libvirt.VIR_DOMAIN_CRASHED: "Crashed",
	libvirt.VIR_DOMAIN_PMSUSPENDED: "Suspended",
	libvirt.VIR_DOMAIN_SHUTOFF: "Shutoff",
}

type (
	Libvirt struct {
		uri         string
		connections chan *libvirt.VirConnection
		max         int
	}

	Domain struct {
		*libvirt.VirDomain
		State int
	}
)

func NewLibvirt(uri string, max int) (*Libvirt, error) {
	lv := &Libvirt{
		uri:         uri,
		max:         max,
		connections: make(chan *libvirt.VirConnection, max),
	}

	for i := 0; i < max; i++ {
		v, err := libvirt.NewVirConnection(uri)
		if err != nil {
			return nil, err
		}
		lv.connections <- &v
	}

	return lv, nil
}

func (lv *Libvirt) RunHTTP(port uint) error {
	server, err := rpc.NewServer(int(port))
	if err != nil {
		return err
	}

	server.RegisterService(lv)
	return server.ListenAndServe()
}

func (lv *Libvirt) getConnection() *libvirt.VirConnection {
	conn := <-lv.connections
	defer func() {
		lv.connections <- conn
	}()

	return conn
}

func (lv *Libvirt) LookupDomainByName(name string) (*libvirt.VirDomain, error) {
	conn := lv.getConnection()

	domain, err := conn.LookupDomainByName(name)
	if err != nil {
		return nil, err
	}

	return &domain, err
}

func (lv *Libvirt) NewDomain(guest *client.Guest) (*libvirt.VirDomain, error) {
	conn := lv.getConnection()

	domain, err := conn.DomainDefineXML(fmt.Sprintf(`<domain type="test"><name>%s</name><memory unit="MiB">%d</memory><os><type>hvm</type></os></domain>`, guest.Id, guest.Memory))
	if err != nil {
		return nil, err
	}

	return &domain, err
}

func GetState(domain *libvirt.VirDomain) (int, error) {
	state, err := domain.GetState()
	if err != nil {
		return libvirt.VIR_DOMAIN_NOSTATE, err
	}

	return state[0], nil
}

func (lv *Libvirt) DomainWrapper(fn func(*libvirt.VirDomain, int) error) func(*http.Request, *rpc.GuestRequest, *rpc.GuestResponse) error {
	return func(r *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
		if request.Guest == nil || request.Guest.Id == "" {
			return syscall.EINVAL
		}

		domain, err := lv.LookupDomainByName(request.Guest.Id)
		defer domain.Free()
		if err != nil {
			return err
		}

		state, err := GetState(domain)
		if err != nil {
			return err
		}

		err = fn(domain, state)
		if err != nil {
			return err
		}

		*response = rpc.GuestResponse{
			Guest: request.Guest,
		}

		state, err = GetState(domain)
		if err != nil {
			return err
		}

		response.Guest.State = StateNames[state]

		return nil
	}
}

// Guest Actions

func (lv *Libvirt) Restart(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	return lv.Reboot(http, request, response)
}

func (lv *Libvirt) Poweroff(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	return lv.Shutdown(http, request, response)
}

func (lv *Libvirt) Delete(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Delete %s\n", request.Guest.Id)

	domain, err := lv.LookupDomainByName(request.Guest.Id)
	defer domain.Free()
	if err != nil {
		return err
	}

	err = domain.Destroy()
	if err != nil {
		return err
	}

	*response = rpc.GuestResponse{
		Guest: nil,
	}

	return nil
}

func (lv *Libvirt) Create(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Create %s\n", request.Guest.Id)

	domain, err := lv.NewDomain(request.Guest)
	defer domain.Free()
	if err != nil {
		return err
	}

	err = domain.Create()
	if err != nil {
		return err
	}

	*response = rpc.GuestResponse{
		Guest: request.Guest,
	}

	state, err := GetState(domain)
	if err != nil {
		return err
	}

	response.Guest.State = StateNames[state]

	return nil
}

func (lv *Libvirt) Run(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Run %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {
		switch state {

		case libvirt.VIR_DOMAIN_RUNNING:
			// nothing to do

		case libvirt.VIR_DOMAIN_SHUTDOWN, libvirt.VIR_DOMAIN_SHUTOFF, libvirt.VIR_DOMAIN_BLOCKED, libvirt.VIR_DOMAIN_NOSTATE:
			return domain.Create()

		case libvirt.VIR_DOMAIN_PAUSED, libvirt.VIR_DOMAIN_PMSUSPENDED:
			return domain.Resume()
		}

		return nil
	})(http, request, response)
}

func (lv *Libvirt) Reboot(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Reboot %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {
		return domain.Reboot(0)
	})(http, request, response)
}

func (lv *Libvirt) Shutdown(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Shutdown %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {

		switch state {
		case libvirt.VIR_DOMAIN_SHUTDOWN, libvirt.VIR_DOMAIN_SHUTOFF:
			// nothing to do

		default:
			return domain.Shutdown()
		}

		return nil
	})(http, request, response)
}

func (lv *Libvirt) Status(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Status %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {
		// DomainWrapper gets the state already, no need to do anything here
		return nil
	})(http, request, response)
}

/*
func (lv *Libvirt) CpuMetrics(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
}

func (lv *Libvirt) DiskMetrics(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
}

func (lv *Libvirt) NicMetrics(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
}
*/
<commit_msg>MIST-138: just return same guest from request<commit_after>package libvirt

import (
	"github.com/alexzorin/libvirt-go"
	"github.com/mistifyio/mistify-agent/rpc"
	"github.com/mistifyio/mistify-agent/client"
	"github.com/mistifyio/mistify-agent/log"
	"syscall"
	"net/http"
	"fmt"
)

var StateNames = map[int]string{
	libvirt.VIR_DOMAIN_NOSTATE: "No State",
	libvirt.VIR_DOMAIN_RUNNING: "Running",
	libvirt.VIR_DOMAIN_BLOCKED: "Blocked",
	libvirt.VIR_DOMAIN_PAUSED: "Paused",
	libvirt.VIR_DOMAIN_SHUTDOWN: "Shutdown",
	libvirt.VIR_DOMAIN_CRASHED: "Crashed",
	libvirt.VIR_DOMAIN_PMSUSPENDED: "Suspended",
	libvirt.VIR_DOMAIN_SHUTOFF: "Shutoff",
}

type (
	Libvirt struct {
		uri         string
		connections chan *libvirt.VirConnection
		max         int
	}

	Domain struct {
		*libvirt.VirDomain
		State int
	}
)

func NewLibvirt(uri string, max int) (*Libvirt, error) {
	lv := &Libvirt{
		uri:         uri,
		max:         max,
		connections: make(chan *libvirt.VirConnection, max),
	}

	for i := 0; i < max; i++ {
		v, err := libvirt.NewVirConnection(uri)
		if err != nil {
			return nil, err
		}
		lv.connections <- &v
	}

	return lv, nil
}

func (lv *Libvirt) RunHTTP(port uint) error {
	server, err := rpc.NewServer(int(port))
	if err != nil {
		return err
	}

	server.RegisterService(lv)
	return server.ListenAndServe()
}

func (lv *Libvirt) getConnection() *libvirt.VirConnection {
	conn := <-lv.connections
	defer func() {
		lv.connections <- conn
	}()

	return conn
}

func (lv *Libvirt) LookupDomainByName(name string) (*libvirt.VirDomain, error) {
	conn := lv.getConnection()

	domain, err := conn.LookupDomainByName(name)
	if err != nil {
		return nil, err
	}

	return &domain, err
}

func (lv *Libvirt) NewDomain(guest *client.Guest) (*libvirt.VirDomain, error) {
	conn := lv.getConnection()

	domain, err := conn.DomainDefineXML(fmt.Sprintf(`<domain type="test"><name>%s</name><memory unit="MiB">%d</memory><os><type>hvm</type></os></domain>`, guest.Id, guest.Memory))
	if err != nil {
		return nil, err
	}

	return &domain, err
}

func GetState(domain *libvirt.VirDomain) (int, error) {
	state, err := domain.GetState()
	if err != nil {
		return libvirt.VIR_DOMAIN_NOSTATE, err
	}

	return state[0], nil
}

func (lv *Libvirt) DomainWrapper(fn func(*libvirt.VirDomain, int) error) func(*http.Request, *rpc.GuestRequest, *rpc.GuestResponse) error {
	return func(r *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
		if request.Guest == nil || request.Guest.Id == "" {
			return syscall.EINVAL
		}

		domain, err := lv.LookupDomainByName(request.Guest.Id)
		defer domain.Free()
		if err != nil {
			return err
		}

		state, err := GetState(domain)
		if err != nil {
			return err
		}

		err = fn(domain, state)
		if err != nil {
			return err
		}

		*response = rpc.GuestResponse{
			Guest: request.Guest,
		}

		state, err = GetState(domain)
		if err != nil {
			return err
		}

		response.Guest.State = StateNames[state]

		return nil
	}
}

// Guest Actions

func (lv *Libvirt) Restart(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	return lv.Reboot(http, request, response)
}

func (lv *Libvirt) Poweroff(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	return lv.Shutdown(http, request, response)
}

func (lv *Libvirt) Delete(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Delete %s\n", request.Guest.Id)

	domain, err := lv.LookupDomainByName(request.Guest.Id)
	defer domain.Free()
	if err != nil {
		return err
	}

	err = domain.Destroy()
	if err != nil {
		return err
	}

	*response = rpc.GuestResponse{
		Guest: request.Guest,
	}

	return nil
}

func (lv *Libvirt) Create(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Create %s\n", request.Guest.Id)

	domain, err := lv.NewDomain(request.Guest)
	defer domain.Free()
	if err != nil {
		return err
	}

	err = domain.Create()
	if err != nil {
		return err
	}

	*response = rpc.GuestResponse{
		Guest: request.Guest,
	}

	state, err := GetState(domain)
	if err != nil {
		return err
	}

	response.Guest.State = StateNames[state]

	return nil
}

func (lv *Libvirt) Run(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Run %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {
		switch state {

		case libvirt.VIR_DOMAIN_RUNNING:
			// nothing to do

		case libvirt.VIR_DOMAIN_SHUTDOWN, libvirt.VIR_DOMAIN_SHUTOFF, libvirt.VIR_DOMAIN_BLOCKED, libvirt.VIR_DOMAIN_NOSTATE:
			return domain.Create()

		case libvirt.VIR_DOMAIN_PAUSED, libvirt.VIR_DOMAIN_PMSUSPENDED:
			return domain.Resume()
		}

		return nil
	})(http, request, response)
}

func (lv *Libvirt) Reboot(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Reboot %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {
		return domain.Reboot(0)
	})(http, request, response)
}

func (lv *Libvirt) Shutdown(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Shutdown %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {

		switch state {
		case libvirt.VIR_DOMAIN_SHUTDOWN, libvirt.VIR_DOMAIN_SHUTOFF:
			// nothing to do

		default:
			return domain.Shutdown()
		}

		return nil
	})(http, request, response)
}

func (lv *Libvirt) Status(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Status %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {
		// DomainWrapper gets the state already, no need to do anything here
		return nil
	})(http, request, response)
}

/*
func (lv *Libvirt) CpuMetrics(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
}

func (lv *Libvirt) DiskMetrics(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
}

func (lv *Libvirt) NicMetrics(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
}
*/
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"os/user"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/BurntSushi/toml"
)

func defaultConfigDir() string {
	return filepath.Join(xdgConfigDir(), "remote-do")
}

func xdgConfigDir() string {
	if d := os.Getenv("XDG_CONFIG_HOME"); d != "" {
		return d
	}
	var h string
	if u, err := user.Current(); err == nil {
		h = u.HomeDir
	} else {
		h = os.Getenv("HOME")
	}
	return filepath.Join(h, ".config")
}

var (
	configDir      = flag.String("configdir", defaultConfigDir(), "path to configuration directory")
	selectedRemote = flag.String("remote", "", "remote server to use (defaults to default in config)")
	namePre        = flag.String("name", "", "prefix to give session file")
)

func usage() {
	fmt.Fprintln(os.Stderr, "remote-do [args] cmd args...")
	flag.PrintDefaults()
	os.Exit(1)
}

func main() {
	rand.Seed(time.Now().UnixNano())
	log.SetPrefix("remote-do: ")
	log.SetFlags(0)
	flag.Usage = usage
	flag.Parse()
	if flag.NArg() == 0 {
		flag.Usage()
	}
	config, err := loadConfig(*configDir)
	if err != nil {
		log.Fatal(err)
	}
	remoteName := *selectedRemote
	if remoteName == "" {
		remoteName = config.DefaultRemote
	}
	if remoteName == "" {
		log.Fatal("no selected or default remote")
	}
	remote, ok := config.Remotes[remoteName]
	if !ok {
		log.Fatalf("selected remote %q has no config", remoteName)
	}
	if hasRelativeRemotePaths(remote) {
		remHome, err := getRemoteHome(remote)
		if err != nil {
			log.Fatalf("unable to get remote home directory: %v", err)
		}
		remote = resolveRelative(remote, remHome)
	}
	pre := ""
	if *namePre != "" {
		pre = *namePre + "-"
	}
	session, f, err := mkSessionFile(pre)
	if err != nil {
		log.Fatal(err)
	}

	remCacheDir := path.Join(remote.TopDir, "cache")
	remSessionDir := path.Join(remote.TopDir, session)
	remTempDir := remSessionDir + "-tmp"

	fmt.Println("clearing data and copying files...")
	_ = cexec("ssh", remote.Host, fmt.Sprintf(
		"rm -rf '%s'; [[ -d '%s' ]] && cp -R '%s' '%s'",
		remSessionDir, remCacheDir, remCacheDir, remSessionDir))
	args := []string{"-az", "--delete", "."}
	if _, err := os.Stat(".remote-do-ignore"); err == nil {
		args = []string{"-az", "--delete", "--exclude-from=.remote-do-ignore", "."}
	}
	err = cexec("rsync", append(args, remote.Host+":"+remSessionDir+"/")...)
	if err != nil {
		log.Fatalf("failed to copy files: %v", err)
	}
	err = cexec("ssh", remote.Host,
		fmt.Sprintf("rsync -a --delete '%s/' '%s'", remSessionDir, remCacheDir))
	if err != nil {
		log.Fatalf("failed to update remote cache: %v", err)
	}

	fmt.Println("running commands...")
	doRun, err := execTmpl(doRunTmpl, struct {
		Cmd []string
	}{
		Cmd: flag.Args(),
	})
	if err != nil {
		log.Fatalf("unable to execute doRun template: %v", err)
	}
	if !strings.HasPrefix(remote.JunestHome, "/") {
		remote.JunestHome = "$HOME/" + remote.JunestHome
	}
	if !strings.HasPrefix(remote.JunestRepo, "/") {
		remote.JunestRepo = "$HOME/" + remote.JunestRepo
	}
	scriptBody, err := execTmpl(scriptBodyTmpl, struct {
		TempDir    string
		JunestCmd  []string
		SessionDir string
	}{
		TempDir:    remTempDir,
		JunestCmd:  []string{"JUNEST_HOME=" + remote.JunestHome, path.Join(remote.JunestRepo, "bin", "junest")},
		SessionDir: remSessionDir,
	})
	if err != nil {
		log.Fatalf("unable to execute scriptBody template: %v", err)
	}
	cmdRun, err := execTmpl(cmdRunTmpl, struct {
		DoRun      string
		ScriptBody string
		SessionDir string
	}{
		DoRun:      doRun,
		ScriptBody: scriptBody,
		SessionDir: remSessionDir,
	})
	if err != nil {
		log.Fatalf("unable to execute cmdRun template: %v", err)
	}

	if err := cexec("ssh", remote.Host, cmdRun); err != nil {
		log.Fatalf("error running command: %v", err)
	}

	err = sessionBodyTmpl.Execute(f, struct {
		Host       string
		SessionDir string
		TempDir    string
		Cmd        []string
	}{
		Host:       remote.Host,
		SessionDir: remSessionDir,
		TempDir:    remTempDir,
		Cmd:        flag.Args(),
	})
	if err != nil {
		log.Fatalf("unable to execute sessionBody template: %v", err)
	}

	f.Close()
	os.Chmod(session, 0755)
	fmt.Println(session)
}

func getRemoteHome(r Remote) (string, error) {
	out, err := exec.Command("ssh", r.Host, "echo", "$HOME").CombinedOutput()
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(out)), nil
}

func cexec(name string, args ...string) error {
	cmd := exec.Command(name, args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func execTmpl(t *template.Template, data interface{}) (string, error) {
	var buf bytes.Buffer
	err := t.Execute(&buf, data)
	return buf.String(), err
}

const doRun = `#!/usr/bin/env bash
export TERM=xterm
export TMPDIR=/runtmp
cd "${0%/*}"
{{range $i, $e := .Cmd}}"{{$e}}" {{end}};
`

const scriptBody = `#!/usr/bin/env bash
mkdir -p "{{.TempDir}}"
env {{range $i, $e := .JunestCmd}}"{{$e}}" {{end}} -p "-b {{.TempDir}}:/runtmp" {{.SessionDir}}/remote-inner.bash
`

const cmdRun = `\
echo '{{.DoRun}}'>'{{.SessionDir}}/remote-inner.bash' && \
echo '{{.ScriptBody}}'>'{{.SessionDir}}/remote-outer.bash' && \
chmod a+x '{{.SessionDir}}/remote-inner.bash' && \
chmod a+x '{{.SessionDir}}/remote-outer.bash' && \
'{{.SessionDir}}/remote-outer.bash'
`

const sessionBody = `#!/usr/bin/env bash
if [[ $1 == clean ]]; then
	# clean will delete files for us that are in the home directory.
	# This lets us safely use root to do so without worrying about
	# deleting unwanted things.
	# This is required because of the permissions docker files are
	# written with.
	ssh "{{.Host}}" "rm -rf '{{.SessionDir}}' '{{.TempDir}}'"
	rm -f "$0"
elif [[ $1 == cmd ]]; then
	echo {{range $i, $e := .Cmd}}{{$e}} {{end}}
else
	echo "usage: $0 cmd|clean" >&2
	exit 1
fi
`

var (
	doRunTmpl       = template.Must(template.New("doRun").Parse(doRun))
	scriptBodyTmpl  = template.Must(template.New("scriptBody").Parse(scriptBody))
	cmdRunTmpl      = template.Must(template.New("cmdRun").Parse(cmdRun))
	sessionBodyTmpl = template.Must(template.New("sessionBody").Parse(sessionBody))
)

func mkSessionFile(namePre string) (string, *os.File, error) {
Start:
	session := namePre + "session." + strconv.FormatInt(rand.Int63(), 36)[:5]
	f, err := os.OpenFile(session, os.O_EXCL|os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		if os.IsExist(err) {
			goto Start
		}
	}
	return session, f, err
}

const configTmpl = `
# Example config file (formatted using TOML).
# You should set default-remote and it's config.
#
#default-remote = default
#
#[remotes.default]
#host = www.example.com
#top-dir = /path/to/dir # or path/to/dir to be relative to your home dir
#junest-repo = /path/to/dir
#junest-home = /path/to/dir
`

func resolveRelative(r Remote, home string) Remote {
	for _, sp := range [...]*string{&r.TopDir, &r.JunestRepo, &r.JunestHome} {
		if !strings.HasPrefix(*sp, "/") {
			*sp = path.Join(home, *sp)
		}
	}
	return r
}

func hasRelativeRemotePaths(r Remote) bool {
	for _, s := range [...]string{r.TopDir, r.JunestRepo, r.JunestHome} {
		if !strings.HasPrefix(s, "/") {
			return true
		}
	}
	return false
}

type Remote struct {
	Host       string `toml:"host"`
	TopDir     string `toml:"top-dir"`
	JunestRepo string `toml:"junest-repo"`
	JunestHome string `toml:"junest-home"`
}

type Config struct {
	DefaultRemote string            `toml:"default-remote"`
	Remotes       map[string]Remote `toml:"remotes"`
}

func loadConfig(d string) (*Config, error) {
	if fi, err := os.Stat(d); err != nil || !fi.IsDir() {
		os.MkdirAll(d, 0755)
	}
	cp := filepath.Join(d, "config.toml")
	if _, err := os.Stat(cp); err != nil {
		ioutil.WriteFile(cp, []byte(configTmpl), 0644)
		return nil, fmt.Errorf("missing config, created skeleton at %s", cp)
	}

	var config Config
	_, err := toml.DecodeFile(cp, &config)
	return &config, err
}
<commit_msg>remote-do: make scripts more flexible<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"os/user"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/BurntSushi/toml"
)

func defaultConfigDir() string {
	return filepath.Join(xdgConfigDir(), "remote-do")
}

func xdgConfigDir() string {
	if d := os.Getenv("XDG_CONFIG_HOME"); d != "" {
		return d
	}
	var h string
	if u, err := user.Current(); err == nil {
		h = u.HomeDir
	} else {
		h = os.Getenv("HOME")
	}
	return filepath.Join(h, ".config")
}

var (
	configDir      = flag.String("configdir", defaultConfigDir(), "path to configuration directory")
	selectedRemote = flag.String("remote", "", "remote server to use (defaults to default in config)")
	namePre        = flag.String("name", "", "prefix to give session file")
)

func usage() {
	fmt.Fprintln(os.Stderr, "remote-do [args] cmd args...")
	flag.PrintDefaults()
	os.Exit(1)
}

func main() {
	rand.Seed(time.Now().UnixNano())
	log.SetPrefix("remote-do: ")
	log.SetFlags(0)
	flag.Usage = usage
	flag.Parse()
	if flag.NArg() == 0 {
		flag.Usage()
	}
	config, err := loadConfig(*configDir)
	if err != nil {
		log.Fatal(err)
	}
	remoteName := *selectedRemote
	if remoteName == "" {
		remoteName = config.DefaultRemote
	}
	if remoteName == "" {
		log.Fatal("no selected or default remote")
	}
	remote, ok := config.Remotes[remoteName]
	if !ok {
		log.Fatalf("selected remote %q has no config", remoteName)
	}
	if hasRelativeRemotePaths(remote) {
		remHome, err := getRemoteHome(remote)
		if err != nil {
			log.Fatalf("unable to get remote home directory: %v", err)
		}
		remote = resolveRelative(remote, remHome)
	}
	pre := ""
	if *namePre != "" {
		pre = *namePre + "-"
	}
	session, f, err := mkSessionFile(pre)
	if err != nil {
		log.Fatal(err)
	}

	remCacheDir := path.Join(remote.TopDir, "cache")
	remSessionDir := path.Join(remote.TopDir, session)
	remTempDir := remSessionDir + "-tmp"

	fmt.Println("clearing data and copying files...")
	_ = cexec("ssh", remote.Host, fmt.Sprintf(
		"rm -rf '%s'; [[ -d '%s' ]] && cp -R '%s' '%s'",
		remSessionDir, remCacheDir, remCacheDir, remSessionDir))
	args := []string{"-az", "--delete", "."}
	if _, err := os.Stat(".remote-do-ignore"); err == nil {
		args = []string{"-az", "--delete", "--exclude-from=.remote-do-ignore", "."}
	}
	err = cexec("rsync", append(args, remote.Host+":"+remSessionDir+"/")...)
	if err != nil {
		log.Fatalf("failed to copy files: %v", err)
	}
	err = cexec("ssh", remote.Host,
		fmt.Sprintf("rsync -a --delete '%s/' '%s'", remSessionDir, remCacheDir))
	if err != nil {
		log.Fatalf("failed to update remote cache: %v", err)
	}

	fmt.Println("running commands...")
	if !strings.HasPrefix(remote.JunestHome, "/") {
		remote.JunestHome = "$HOME/" + remote.JunestHome
	}
	if !strings.HasPrefix(remote.JunestRepo, "/") {
		remote.JunestRepo = "$HOME/" + remote.JunestRepo
	}
	outerBody, err := execTmpl(outerBodyTmpl, struct {
		TempDir    string
		JunestCmd  []string
		SessionDir string
	}{
		TempDir:    remTempDir,
		JunestCmd:  []string{"JUNEST_HOME=" + remote.JunestHome, path.Join(remote.JunestRepo, "bin", "junest")},
		SessionDir: remSessionDir,
	})
	if err != nil {
		log.Fatalf("unable to execute outerBody template: %v", err)
	}
	cmdRun, err := execTmpl(cmdRunTmpl, struct {
		InnerBody  string
		OuterBody  string
		SessionDir string
		Cmd        []string
	}{
		InnerBody:  innerBody,
		OuterBody:  outerBody,
		SessionDir: remSessionDir,
		Cmd:        flag.Args(),
	})
	if err != nil {
		log.Fatalf("unable to execute cmdRun template: %v", err)
	}

	if err := cexec("ssh", remote.Host, cmdRun); err != nil {
		log.Fatalf("error running command: %v", err)
	}

	err = sessionBodyTmpl.Execute(f, struct {
		Host       string
		SessionDir string
		TempDir    string
		Cmd        []string
	}{
		Host:       remote.Host,
		SessionDir: remSessionDir,
		TempDir:    remTempDir,
		Cmd:        flag.Args(),
	})
	if err != nil {
		log.Fatalf("unable to execute sessionBody template: %v", err)
	}

	f.Close()
	os.Chmod(session, 0755)
	fmt.Println(session)
}

func getRemoteHome(r Remote) (string, error) {
	out, err := exec.Command("ssh", r.Host, "echo", "$HOME").CombinedOutput()
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(out)), nil
}

func cexec(name string, args ...string) error {
	cmd := exec.Command(name, args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func execTmpl(t *template.Template, data interface{}) (string, error) {
	var buf bytes.Buffer
	err := t.Execute(&buf, data)
	return buf.String(), err
}

const innerBody = `#!/usr/bin/env bash
export TMPDIR=/runtmp
cd "${0%/*}"
cmd=$1
shift
"$cmd" "$@"
`

const outerBody = `#!/usr/bin/env bash
mkdir -p "{{.TempDir}}"
env {{range $i, $e := .JunestCmd}}"{{$e}}" {{end}} -p "-b {{.TempDir}}:/runtmp" \
	"{{.SessionDir}}/remote-inner.bash" "$@"
`

const cmdRun = `\
echo '{{.OuterBody}}'>'{{.SessionDir}}/remote-outer.bash' && \
echo '{{.InnerBody}}'>'{{.SessionDir}}/remote-inner.bash' && \
chmod a+x '{{.SessionDir}}/remote-inner.bash' '{{.SessionDir}}/remote-outer.bash' && \
'{{.SessionDir}}/remote-outer.bash' {{range $i, $e := .Cmd}}'{{$e}}' {{end}}
`

const sessionBody = `#!/usr/bin/env bash
if [[ $1 == clean ]]; then
	# clean will delete files for us that are in the home directory.
	# This lets us safely use root to do so without worrying about
	# deleting unwanted things.
	# This is required because of the permissions docker files are
	# written with.
	ssh "{{.Host}}" "rm -rf '{{.SessionDir}}' '{{.TempDir}}'"
	rm -f "$0"
elif [[ $1 == cmd ]]; then
	echo {{range $i, $e := .Cmd}}{{$e}} {{end}}
else
	echo "usage: $0 cmd|clean" >&2
	exit 1
fi
`

var (
	outerBodyTmpl   = template.Must(template.New("outerBody").Parse(outerBody))
	cmdRunTmpl      = template.Must(template.New("cmdRun").Parse(cmdRun))
	sessionBodyTmpl = template.Must(template.New("sessionBody").Parse(sessionBody))
)

func mkSessionFile(namePre string) (string, *os.File, error) {
Start:
	session := namePre + "session." + strconv.FormatInt(rand.Int63(), 36)[:5]
	f, err := os.OpenFile(session, os.O_EXCL|os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		if os.IsExist(err) {
			goto Start
		}
	}
	return session, f, err
}

const configTmpl = `
# Example config file (formatted using TOML).
# You should set default-remote and it's config.
#
#default-remote = default
#
#[remotes.default]
#host = www.example.com
#top-dir = /path/to/dir # or path/to/dir to be relative to your home dir
#junest-repo = /path/to/dir
#junest-home = /path/to/dir
`

func resolveRelative(r Remote, home string) Remote {
	for _, sp := range [...]*string{&r.TopDir, &r.JunestRepo, &r.JunestHome} {
		if !strings.HasPrefix(*sp, "/") {
			*sp = path.Join(home, *sp)
		}
	}
	return r
}

func hasRelativeRemotePaths(r Remote) bool {
	for _, s := range [...]string{r.TopDir, r.JunestRepo, r.JunestHome} {
		if !strings.HasPrefix(s, "/") {
			return true
		}
	}
	return false
}

type Remote struct {
	Host       string `toml:"host"`
	TopDir     string `toml:"top-dir"`
	JunestRepo string `toml:"junest-repo"`
	JunestHome string `toml:"junest-home"`
}

type Config struct {
	DefaultRemote string            `toml:"default-remote"`
	Remotes       map[string]Remote `toml:"remotes"`
}

func loadConfig(d string) (*Config, error) {
	if fi, err := os.Stat(d); err != nil || !fi.IsDir() {
		os.MkdirAll(d, 0755)
	}
	cp := filepath.Join(d, "config.toml")
	if _, err := os.Stat(cp); err != nil {
		ioutil.WriteFile(cp, []byte(configTmpl), 0644)
		return nil, fmt.Errorf("missing config, created skeleton at %s", cp)
	}

	var config Config
	_, err := toml.DecodeFile(cp, &config)
	return &config, err
}
<|endoftext|>


<commit_before>package zygo
func init() { GITLASTTAG = "1.5.7"; GITLASTCOMMIT = "7c879666142a93b23a32ce2f23e4b3c5bbd94e65" }
<commit_msg>v1.5.9 atg one-shot-parser<commit_after>package zygo
func init() { GITLASTTAG = "v1.5.9"; GITLASTCOMMIT = "0d9b31cfe24eaec6b4f794d956c9c5de22806c59" }
<|endoftext|>


<commit_before>// Methods for harvesting an OAI-PMH service
package oai

import (
	"net/http"
	"encoding/xml"
	"io/ioutil"
	"strings"
)

// Determine the resumption token in this OAIResponse
func (resp *OAIResponse) ResumptionToken() (hasResumptionToken bool, resumptionToken string) {
	hasResumptionToken = false
	resumptionToken = ""
	if resp == nil { return }

	// First attempt to obtain a resumption token from a ListIdentifiers response
	resumptionToken =  resp.ListIdentifiers.ResumptionToken

	// Then attempt to obtain a resumption token from a ListRecords response
	if resumptionToken == "" {
		resumptionToken =  resp.ListRecords.ResumptionToken
	}

	// If a non-empty resumption token turned up it can safely inferred that...
	if resumptionToken != "" { hasResumptionToken = true }

	return
}

// String representation of the OAI Request
func (req *OAIRequest) String() string {
	qs := []string{req.BaseUrl, "?set=", req.Set, "&metadataPrefix=", req.MetadataPrefix, "&verb=", req.Verb}

	if req.ResumptionToken != "" {
		qs = append(qs, "&resumptionToken=")
		qs = append(qs, req.ResumptionToken)
	}

	if req.Identifier != "" {
		qs = append(qs, "&identifier=")
		qs = append(qs, req.Identifier)
	}

	if req.From != "" {
		qs = append(qs, "&from=")
		qs = append(qs, req.From)
	}

	if req.Until != "" {
		qs = append(qs, "&until=")
		qs = append(qs, req.Until)
	}
	return strings.Join(qs, "")
}

// Perform an HTTP GET request using the OAI Requests fields
// and return an OAI Response reference
func (req *OAIRequest) Perform() (oaiResponse *OAIResponse) {
	// Perform the GET request
	resp, err := http.Get(req.String())
	if err != nil { panic(err) }

	// Make sure the response body object will be closed after
	// reading all the content body's data
	defer resp.Body.Close()

	// Read all the data
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil { panic(err) }

	// Unmarshall all the data
	err = xml.Unmarshal(body, &oaiResponse)
	if err != nil { panic(err) }

	return
}

// Perform a harvest of a complete OAI set, or simply one request
// call the batchCallback function argument with the OAI responses
func (req *OAIRequest) Harvest(batchCallback func(*OAIResponse)) {
	// Use Perform to get the OAI response
	oaiResponse := req.Perform()

	// Execute the callback function with the response
	batchCallback(oaiResponse)

	// Check for a resumptionToken
	hasResumptionToken, resumptionToken := oaiResponse.ResumptionToken()

	// Harvest further if there is a resumption token
	if hasResumptionToken == true {
		req.Set = ""
		req.MetadataPrefix = ""
		req.From = ""
		req.ResumptionToken = resumptionToken
		req.Harvest(batchCallback)
	}
}

// Reads OAI PMH response XML from a file
func FromFile(filename string) (oaiResponse *OAIResponse) {
	bytes, err := ioutil.ReadFile(filename)
	if err != nil { panic(err) }

	// Unmarshall all the data
	err = xml.Unmarshal(bytes, &oaiResponse)
	if err != nil { panic(err) }

	return
}

<commit_msg>added formatters<commit_after>// Methods for harvesting an OAI-PMH service
package oai

import (
	"net/http"
	"encoding/xml"
	"io/ioutil"
	"strings"
	"fmt"
)

// Determine the resumption token in this OAIResponse
func (resp *OAIResponse) ResumptionToken() (hasResumptionToken bool, resumptionToken string) {
	hasResumptionToken = false
	resumptionToken = ""
	if resp == nil { return }

	// First attempt to obtain a resumption token from a ListIdentifiers response
	resumptionToken =  resp.ListIdentifiers.ResumptionToken

	// Then attempt to obtain a resumption token from a ListRecords response
	if resumptionToken == "" {
		resumptionToken =  resp.ListRecords.ResumptionToken
	}

	// If a non-empty resumption token turned up it can safely inferred that...
	if resumptionToken != "" { hasResumptionToken = true }

	return
}

// String representation of the OAI Request
func (req *OAIRequest) String() string {
	qs := []string{req.BaseUrl, "?set=", req.Set, "&metadataPrefix=", req.MetadataPrefix, "&verb=", req.Verb}

	if req.ResumptionToken != "" {
		qs = append(qs, "&resumptionToken=")
		qs = append(qs, req.ResumptionToken)
	}

	if req.Identifier != "" {
		qs = append(qs, "&identifier=")
		qs = append(qs, req.Identifier)
	}

	if req.From != "" {
		qs = append(qs, "&from=")
		qs = append(qs, req.From)
	}

	if req.Until != "" {
		qs = append(qs, "&until=")
		qs = append(qs, req.Until)
	}
	return strings.Join(qs, "")
}

// Perform an HTTP GET request using the OAI Requests fields
// and return an OAI Response reference
func (req *OAIRequest) Perform() (oaiResponse *OAIResponse) {
	// Perform the GET request
	resp, err := http.Get(req.String())
	if err != nil { panic(err) }

	// Make sure the response body object will be closed after
	// reading all the content body's data
	defer resp.Body.Close()

	// Read all the data
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil { panic(err) }

	// Unmarshall all the data
	err = xml.Unmarshal(body, &oaiResponse)
	if err != nil { panic(err) }

	return
}

// Perform a harvest of a complete OAI set, or simply one request
// call the batchCallback function argument with the OAI responses
func (req *OAIRequest) Harvest(batchCallback func(*OAIResponse)) {
	// Use Perform to get the OAI response
	oaiResponse := req.Perform()

	// Execute the callback function with the response
	batchCallback(oaiResponse)

	// Check for a resumptionToken
	hasResumptionToken, resumptionToken := oaiResponse.ResumptionToken()

	// Harvest further if there is a resumption token
	if hasResumptionToken == true {
		req.Set = ""
		req.MetadataPrefix = ""
		req.From = ""
		req.ResumptionToken = resumptionToken
		req.Harvest(batchCallback)
	}
}

// Reads OAI PMH response XML from a file
func FromFile(filename string) (oaiResponse *OAIResponse) {
	bytes, err := ioutil.ReadFile(filename)
	if err != nil { panic(err) }

	// Unmarshall all the data
	err = xml.Unmarshal(bytes, &oaiResponse)
	if err != nil { panic(err) }

	return
}

// Formatter for Metadata content
func (md Metadata) GoString() string { return fmt.Sprintf("%s", md.Body); }

// Formatter for SetDescription content
func (sd SetDescription) GoString() string { return fmt.Sprintf("%s", sd.Body); }

// Formatter for SetDescription content
func (ab About) GoString() string { return fmt.Sprintf("%s", ab.Body); }

<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/takama/daemon"
	"log"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"
)

// #include <wiringPi.h>
// #cgo LDFLAGS: -lwiringPi
import "C"

const (
	// CPU temperature target, degrees C
	defaultTempTarget = 50.
	hysteresis        = float32(1.)

	pwmClockDivisor = 100

	/* Minimum duty cycle is the point below which the fan does
	/* not spin. This depends on both your fan and the switching
	/* transistor used. */
	defaultPwmDutyMin = 1
	pwmDutyMax        = 10

	// how often to update
	delaySeconds = 30

	// GPIO-1/BCM "18"/Pin 12 on a Raspberry PI 3
	defaultPin = 1

	// name of the service
	name        = "fancontrol"
	description = "cooling fan speed control based on CPU temperature"

	// port which daemon should be listen
	port = ":9977"
)

var stdlog, errlog *log.Logger

func fanControl(pwmDutyMin int, pin int, tempTarget float32) {
	cPin := C.int(pin)

	C.wiringPiSetup()

	// Power on "test". Allows the user to verify that their fan is working.
	C.pinMode(cPin, C.OUTPUT)
	C.digitalWrite(cPin, C.HIGH)
	time.Sleep(5 * time.Second)
	C.digitalWrite(cPin, C.LOW)

	C.pwmSetMode(C.PWM_MODE_MS)
	C.pinMode(cPin, C.PWM_OUTPUT)
	C.pwmSetRange(pwmDutyMax)
	C.pwmSetClock(pwmClockDivisor)
	C.pwmWrite(cPin, C.int(pwmDutyMin))
	temp := float32(0.)
	go cpuTempMonitor(func(cpuTemp float32) {
		if isCPUTempValid(cpuTemp) {
			temp = cpuTemp
		}
	})
	pwmDuty := 0

	delay := time.NewTicker(delaySeconds * time.Second)

	for {
		if temp > (tempTarget + hysteresis) {
			pwmDuty = iMax(iMin(pwmDutyMax, pwmDuty+1), pwmDutyMin)
		} else if temp < (tempTarget - hysteresis) {
			pwmDuty = iMax(pwmDuty-1, 0)
			if pwmDuty < pwmDutyMin {
				pwmDuty = 0
			}
		}
		//log.Println(temp, " ", pwmDuty)
		C.pwmWrite(cPin, C.int(pwmDuty))
		<-delay.C
	}

	// Default to "ON".
	C.pinMode(cPin, C.OUTPUT)
	C.digitalWrite(cPin, C.HIGH)
}

// Service has embedded daemon
type Service struct {
	daemon.Daemon
}

// Manage by daemon commands or run the daemon
func (service *Service) Manage() (string, error) {

	tempTarget := flag.Float64("temp", defaultTempTarget, "Target CPU Temperature, degrees C")
	pwmDutyMin := flag.Int("minduty", defaultPwmDutyMin, "Minimum PWM duty cycle")
	pin := flag.Int("pin", defaultPin, "PWM pin (wiringPi numbering)")
	flag.Parse()

	usage := "Usage: " + name + " install | remove | start | stop | status"
	// if received any kind of command, do it
	if flag.NArg() > 0 {
		command := os.Args[flag.NFlag()+1]
		switch command {
		case "install":
			return service.Install()
		case "remove":
			return service.Remove()
		case "start":
			return service.Start()
		case "stop":
			return service.Stop()
		case "status":
			return service.Status()
		default:
			return usage, nil
		}
	}

	go fanControl(*pwmDutyMin, *pin, float32(*tempTarget))

	// Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt, os.Kill, syscall.SIGTERM)

	// Set up listener for defined host and port
	listener, err := net.Listen("tcp", port)
	if err != nil {
		return "Possibly was a problem with the port binding", err
	}

	// set up channel on which to send accepted connections
	listen := make(chan net.Conn, 100)
	go acceptConnection(listener, listen)

	// loop work cycle with accept connections or interrupt
	// by system signal
	for {
		select {
		case conn := <-listen:
			go handleClient(conn)
		case killSignal := <-interrupt:
			stdlog.Println("Got signal:", killSignal)
			stdlog.Println("Stoping listening on ", listener.Addr())
			listener.Close()
			if killSignal == os.Interrupt {
				return "Daemon was interrupted by system signal", nil
			}
			return "Daemon was killed", nil
		}
	}
}

// Accept a client connection and collect it in a channel
func acceptConnection(listener net.Listener, listen chan net.Conn) {
	for {
		conn, err := listener.Accept()
		if err != nil {
			continue
		}
		listen <- conn
	}
}

func handleClient(client net.Conn) {
	defer client.Close()
	for {
		buf := make([]byte, 4096)
		numbytes, err := client.Read(buf)
		if numbytes == 0 || err != nil {
			return
		}
		client.Write(buf[:numbytes])
	}
}

func init() {
	stdlog = log.New(os.Stdout, "", 0)
	errlog = log.New(os.Stderr, "", 0)
}

func main() {
	srv, err := daemon.New(name, description, []string{}...)
	if err != nil {
		errlog.Println("Error: ", err)
		os.Exit(1)
	}
	service := &Service{srv}
	status, err := service.Manage()
	if err != nil {
		errlog.Println(status, "\nError: ", err)
		os.Exit(1)
	}
	fmt.Println(status)
}
<commit_msg>Remove socket listener from fancontrol. Replace with http server. Serve status in JSON format.<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/takama/daemon"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"
)

// #include <wiringPi.h>
// #cgo LDFLAGS: -lwiringPi
import "C"

const (
	// CPU temperature target, degrees C
	defaultTempTarget = 50.
	hysteresis        = float32(1.)

	pwmClockDivisor = 100

	/* Minimum duty cycle is the point below which the fan does
	/* not spin. This depends on both your fan and the switching
	/* transistor used. */
	defaultPwmDutyMin = 1
	pwmDutyMax        = 10

	// how often to update
	delaySeconds = 30

	// GPIO-1/BCM "18"/Pin 12 on a Raspberry PI 3
	defaultPin = 1

	// name of the service
	name        = "fancontrol"
	description = "cooling fan speed control based on CPU temperature"

	// Address on which daemon should be listen.
	addr = ":9977"
)

type FanControl struct {
	TempTarget     float32
	TempCurrent    float32
	PWMDutyMin     int
	PWMDutyMax     int
	PWMDutyCurrent int
	PWMPin         int
}

var myFanControl FanControl

var stdlog, errlog *log.Logger

func fanControl() {
	cPin := C.int(myFanControl.PWMPin)

	C.wiringPiSetup()

	// Power on "test". Allows the user to verify that their fan is working.
	C.pinMode(cPin, C.OUTPUT)
	C.digitalWrite(cPin, C.HIGH)
	time.Sleep(5 * time.Second)
	C.digitalWrite(cPin, C.LOW)

	C.pwmSetMode(C.PWM_MODE_MS)
	C.pinMode(cPin, C.PWM_OUTPUT)
	C.pwmSetRange(C.uint(myFanControl.PWMDutyMax))
	C.pwmSetClock(pwmClockDivisor)
	C.pwmWrite(cPin, C.int(myFanControl.PWMDutyMin))
	myFanControl.TempCurrent = 0
	go cpuTempMonitor(func(cpuTemp float32) {
		if isCPUTempValid(cpuTemp) {
			myFanControl.TempCurrent = cpuTemp
		}
	})

	myFanControl.PWMDutyCurrent = 0

	delay := time.NewTicker(delaySeconds * time.Second)

	for {
		if myFanControl.TempCurrent > (myFanControl.TempTarget + hysteresis) {
			myFanControl.PWMDutyCurrent = iMax(iMin(myFanControl.PWMDutyMax, myFanControl.PWMDutyCurrent+1), myFanControl.PWMDutyMin)
		} else if myFanControl.TempCurrent < (myFanControl.TempTarget - hysteresis) {
			myFanControl.PWMDutyCurrent = iMax(myFanControl.PWMDutyCurrent-1, 0)
			if myFanControl.PWMDutyCurrent < myFanControl.PWMDutyMin {
				myFanControl.PWMDutyCurrent = 0
			}
		}
		//log.Println(myFanControl.TempCurrent, " ", myFanControl.PWMDutyCurrent)
		C.pwmWrite(cPin, C.int(myFanControl.PWMDutyCurrent))
		<-delay.C
	}

	// Default to "ON".
	C.pinMode(cPin, C.OUTPUT)
	C.digitalWrite(cPin, C.HIGH)
}

// Service has embedded daemon
type Service struct {
	daemon.Daemon
}

// Manage by daemon commands or run the daemon
func (service *Service) Manage() (string, error) {

	tempTarget := flag.Float64("temp", defaultTempTarget, "Target CPU Temperature, degrees C")
	pwmDutyMin := flag.Int("minduty", defaultPwmDutyMin, "Minimum PWM duty cycle")
	pin := flag.Int("pin", defaultPin, "PWM pin (wiringPi numbering)")
	flag.Parse()

	usage := "Usage: " + name + " install | remove | start | stop | status"
	// if received any kind of command, do it
	if flag.NArg() > 0 {
		command := os.Args[flag.NFlag()+1]
		switch command {
		case "install":
			return service.Install()
		case "remove":
			return service.Remove()
		case "start":
			return service.Start()
		case "stop":
			return service.Stop()
		case "status":
			return service.Status()
		default:
			return usage, nil
		}
	}

	myFanControl.TempTarget = float32(*tempTarget)
	myFanControl.PWMDutyMin = *pwmDutyMin
	myFanControl.PWMDutyMax = pwmDutyMax
	myFanControl.PWMPin = *pin

	go fanControl()

	// Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt, os.Kill, syscall.SIGTERM)

	http.HandleFunc("/", handleStatusRequest)
	http.ListenAndServe(addr, nil)

	// interrupt by system signal
	for {
		killSignal := <-interrupt
		stdlog.Println("Got signal:", killSignal)
		if killSignal == os.Interrupt {
			return "Daemon was interrupted by system signal", nil
		}
		return "Daemon was killed", nil
	}
}

func handleStatusRequest(w http.ResponseWriter, r *http.Request) {
	statusJSON, _ := json.Marshal(&myFanControl)
	w.Write(statusJSON)
}

func init() {
	stdlog = log.New(os.Stdout, "", 0)
	errlog = log.New(os.Stderr, "", 0)
}

func main() {
	srv, err := daemon.New(name, description, []string{}...)
	if err != nil {
		errlog.Println("Error: ", err)
		os.Exit(1)
	}
	service := &Service{srv}
	status, err := service.Manage()
	if err != nil {
		errlog.Println(status, "\nError: ", err)
		os.Exit(1)
	}
	fmt.Println(status)
}
<|endoftext|>