<commit_before>//: ----------------------------------------------------------------------------
//: Copyright (C) 2017 Verizon.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    stats.go
//: details: exposes flow status
//: author:  Mehrdad Arshad Rad
//: date:    02/01/2017
//:
//: Licensed under the Apache License, Version 2.0 (the "License");
//: you may not use this file except in compliance with the License.
//: You may obtain a copy of the License at
//:
//:     http://www.apache.org/licenses/LICENSE-2.0
//:
//: Unless required by applicable law or agreed to in writing, software
//: distributed under the License is distributed on an "AS IS" BASIS,
//: WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//: See the License for the specific language governing permissions and
//: limitations under the License.
//: ----------------------------------------------------------------------------

package main

import (
	"encoding/json"
	"net"
	"net/http"
	"runtime"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var startTime = time.Now().Unix()

type rest struct {
	StartTime int64
	IPFIX     *IPFIXStats
	SFlow     *SFlowStats
	NetflowV5 *NetflowV5Stats
	NetflowV9 *NetflowV9Stats
}

func statsSysHandler(w http.ResponseWriter, r *http.Request) {
	var mem runtime.MemStats

	runtime.ReadMemStats(&mem)
	var data = &struct {
		MemAlloc        uint64
		MemTotalAlloc   uint64
		MemHeapAlloc    uint64
		MemHeapSys      uint64
		MemHeapReleased uint64
		MCacheInuse     uint64
		GCSys           uint64
		GCNext          uint64
		GCLast          string
		NumLogicalCPU   int
		NumGoroutine    int
		MaxProcs        int
		GoVersion       string
		StartTime       int64
	}{
		mem.Alloc,
		mem.TotalAlloc,
		mem.HeapAlloc,
		mem.HeapSys,
		mem.HeapReleased,
		mem.MCacheInuse,
		mem.GCSys,
		mem.NextGC,
		time.Unix(0, int64(mem.LastGC)).String(),
		runtime.NumCPU(),
		runtime.NumGoroutine(),
		runtime.GOMAXPROCS(-1),
		runtime.Version(),
		startTime,
	}

	j, err := json.Marshal(data)
	if err != nil {
		logger.Println(err)
	}

	if _, err = w.Write(j); err != nil {
		logger.Println(err)
	}
}

func statsFlowHandler(protos []proto) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		rd := &rest{StartTime: startTime}
		for _, p := range protos {
			switch p.(type) {
			case *IPFIX:
				ipfix, _ := p.(*IPFIX)
				rd.IPFIX = ipfix.status()
			case *SFlow:
				sflow, _ := p.(*SFlow)
				rd.SFlow = sflow.status()
			case *NetflowV5:
				netflowv5, _ := p.(*NetflowV5)
				rd.NetflowV5 = netflowv5.status()
			case *NetflowV9:
				netflowv9, _ := p.(*NetflowV9)
				rd.NetflowV9 = netflowv9.status()
			}
		}

		j, err := json.Marshal(rd)
		if err != nil {
			logger.Println(err)
		}

		if _, err = w.Write(j); err != nil {
			logger.Println(err)
		}
	}
}

func statsExpose(protos []proto) {
	if opts.StatsFormat != "prometheus" {
		statsRest(protos)
	} else {
		statsPrometheus(protos)
	}
}

func statsRest(protos []proto) {
	if !opts.StatsEnabled {
		return
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/sys", statsSysHandler)
	mux.HandleFunc("/flow", statsFlowHandler(protos))

	logger.Println("starting stats http server ...")

	addr := net.JoinHostPort(opts.StatsHTTPAddr, opts.StatsHTTPPort)
	err := http.ListenAndServe(addr, mux)
	if err != nil {
		logger.Fatal(err)
	}
}

func statsPrometheus(protos []proto) {
	for _, p := range protos {
		promCounterDecoded(p)
		promCounterMQError(p)
		promCounterUDP(p)
		promGaugeMessageQueue(p)
		promGaugeUDPQueue(p)
		promGaugeWorkers(p)
		promGaugeUDPMirrorQueue(p)
	}

	logger.Println("starting prometheus http server ...")

	addr := net.JoinHostPort(opts.StatsHTTPAddr, opts.StatsHTTPPort)
	http.Handle("/metrics", promhttp.Handler())
	http.ListenAndServe(addr, nil)
}

func promCounterDecoded(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_ipfix_decoded_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().DecodedCount)
			})
	case *SFlow:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_sflow_decoded_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().DecodedCount)
			})
	case *NetflowV5:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv5_decoded_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().DecodedCount)
			})
	case *NetflowV9:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv9_decoded_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().DecodedCount)
			})
	}
}

func promCounterMQError(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_ipfix_mq_error",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MQErrorCount)
			})
	case *SFlow:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_sflow_mq_error",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MQErrorCount)
			})
	case *NetflowV5:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv5_mq_error",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MQErrorCount)
			})
	case *NetflowV9:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv9_mq_error",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MQErrorCount)
			})
	}
}

func promCounterUDP(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_ipfix_udp_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPCount)
			})
	case *SFlow:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_sflow_udp_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPCount)
			})
	case *NetflowV5:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv5_udp_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPCount)
			})
	case *NetflowV9:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv9_udp_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPCount)
			})
	}
}

func promGaugeMessageQueue(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewGaugeFunc(prometheus.GaugeOpts{
			Name: "vflow_ipfix_message_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MessageQueue)
			})
	case *SFlow:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_sflow_message_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MessageQueue)
			})
	case *NetflowV5:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv5_message_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MessageQueue)
			})
	case *NetflowV9:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv9_message_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MessageQueue)
			})
	}
}

func promGaugeUDPQueue(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewGaugeFunc(prometheus.GaugeOpts{
			Name: "vflow_ipfix_udp_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPQueue)
			})
	case *SFlow:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_sflow_udp_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPQueue)
			})
	case *NetflowV5:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv5_udp_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPQueue)
			})
	case *NetflowV9:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv9_udp_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPQueue)
			})
	}
}

func promGaugeWorkers(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewGaugeFunc(prometheus.GaugeOpts{
			Name: "vflow_ipfix_workers",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().Workers)
			})
	case *SFlow:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_sflow_workers",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().Workers)
			})
	case *NetflowV5:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv5_workers",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().Workers)
			})
	case *NetflowV9:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv9_workers",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().Workers)
			})
	}
}

func promGaugeUDPMirrorQueue(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewGaugeFunc(prometheus.GaugeOpts{
			Name: "vflow_ipfix_udp_mirror_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPMirrorQueue)
			})
	}
}
<commit_msg>fix #148<commit_after>//: ----------------------------------------------------------------------------
//: Copyright (C) 2017 Verizon.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    stats.go
//: details: exposes flow status
//: author:  Mehrdad Arshad Rad
//: date:    02/01/2017
//:
//: Licensed under the Apache License, Version 2.0 (the "License");
//: you may not use this file except in compliance with the License.
//: You may obtain a copy of the License at
//:
//:     http://www.apache.org/licenses/LICENSE-2.0
//:
//: Unless required by applicable law or agreed to in writing, software
//: distributed under the License is distributed on an "AS IS" BASIS,
//: WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//: See the License for the specific language governing permissions and
//: limitations under the License.
//: ----------------------------------------------------------------------------

package main

import (
	"encoding/json"
	"net"
	"net/http"
	"runtime"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var startTime = time.Now().Unix()

type rest struct {
	StartTime int64
	IPFIX     *IPFIXStats
	SFlow     *SFlowStats
	NetflowV5 *NetflowV5Stats
	NetflowV9 *NetflowV9Stats
}

func statsSysHandler(w http.ResponseWriter, r *http.Request) {
	var mem runtime.MemStats

	runtime.ReadMemStats(&mem)
	var data = &struct {
		MemAlloc        uint64
		MemTotalAlloc   uint64
		MemHeapAlloc    uint64
		MemHeapSys      uint64
		MemHeapReleased uint64
		MCacheInuse     uint64
		GCSys           uint64
		GCNext          uint64
		GCLast          string
		NumLogicalCPU   int
		NumGoroutine    int
		MaxProcs        int
		GoVersion       string
		StartTime       int64
	}{
		mem.Alloc,
		mem.TotalAlloc,
		mem.HeapAlloc,
		mem.HeapSys,
		mem.HeapReleased,
		mem.MCacheInuse,
		mem.GCSys,
		mem.NextGC,
		time.Unix(0, int64(mem.LastGC)).String(),
		runtime.NumCPU(),
		runtime.NumGoroutine(),
		runtime.GOMAXPROCS(-1),
		runtime.Version(),
		startTime,
	}

	j, err := json.Marshal(data)
	if err != nil {
		logger.Println(err)
	}

	if _, err = w.Write(j); err != nil {
		logger.Println(err)
	}
}

func statsFlowHandler(protos []proto) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		rd := &rest{StartTime: startTime}
		for _, p := range protos {
			switch p.(type) {
			case *IPFIX:
				ipfix, _ := p.(*IPFIX)
				rd.IPFIX = ipfix.status()
			case *SFlow:
				sflow, _ := p.(*SFlow)
				rd.SFlow = sflow.status()
			case *NetflowV5:
				netflowv5, _ := p.(*NetflowV5)
				rd.NetflowV5 = netflowv5.status()
			case *NetflowV9:
				netflowv9, _ := p.(*NetflowV9)
				rd.NetflowV9 = netflowv9.status()
			}
		}

		j, err := json.Marshal(rd)
		if err != nil {
			logger.Println(err)
		}

		if _, err = w.Write(j); err != nil {
			logger.Println(err)
		}
	}
}

func statsExpose(protos []proto) {
	if !opts.StatsEnabled {
		return
	}

	if opts.StatsFormat != "prometheus" {
		statsRest(protos)
	} else {
		statsPrometheus(protos)
	}
}

func statsRest(protos []proto) {
	mux := http.NewServeMux()
	mux.HandleFunc("/sys", statsSysHandler)
	mux.HandleFunc("/flow", statsFlowHandler(protos))

	logger.Println("starting stats http server ...")

	addr := net.JoinHostPort(opts.StatsHTTPAddr, opts.StatsHTTPPort)
	err := http.ListenAndServe(addr, mux)
	if err != nil {
		logger.Fatal(err)
	}
}

func statsPrometheus(protos []proto) {
	for _, p := range protos {
		promCounterDecoded(p)
		promCounterMQError(p)
		promCounterUDP(p)
		promGaugeMessageQueue(p)
		promGaugeUDPQueue(p)
		promGaugeWorkers(p)
		promGaugeUDPMirrorQueue(p)
	}

	logger.Println("starting prometheus http server ...")

	addr := net.JoinHostPort(opts.StatsHTTPAddr, opts.StatsHTTPPort)
	http.Handle("/metrics", promhttp.Handler())
	http.ListenAndServe(addr, nil)
}

func promCounterDecoded(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_ipfix_decoded_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().DecodedCount)
			})
	case *SFlow:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_sflow_decoded_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().DecodedCount)
			})
	case *NetflowV5:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv5_decoded_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().DecodedCount)
			})
	case *NetflowV9:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv9_decoded_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().DecodedCount)
			})
	}
}

func promCounterMQError(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_ipfix_mq_error",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MQErrorCount)
			})
	case *SFlow:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_sflow_mq_error",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MQErrorCount)
			})
	case *NetflowV5:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv5_mq_error",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MQErrorCount)
			})
	case *NetflowV9:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv9_mq_error",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MQErrorCount)
			})
	}
}

func promCounterUDP(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_ipfix_udp_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPCount)
			})
	case *SFlow:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_sflow_udp_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPCount)
			})
	case *NetflowV5:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv5_udp_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPCount)
			})
	case *NetflowV9:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv9_udp_packets",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPCount)
			})
	}
}

func promGaugeMessageQueue(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewGaugeFunc(prometheus.GaugeOpts{
			Name: "vflow_ipfix_message_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MessageQueue)
			})
	case *SFlow:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_sflow_message_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MessageQueue)
			})
	case *NetflowV5:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv5_message_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MessageQueue)
			})
	case *NetflowV9:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv9_message_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().MessageQueue)
			})
	}
}

func promGaugeUDPQueue(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewGaugeFunc(prometheus.GaugeOpts{
			Name: "vflow_ipfix_udp_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPQueue)
			})
	case *SFlow:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_sflow_udp_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPQueue)
			})
	case *NetflowV5:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv5_udp_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPQueue)
			})
	case *NetflowV9:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv9_udp_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPQueue)
			})
	}
}

func promGaugeWorkers(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewGaugeFunc(prometheus.GaugeOpts{
			Name: "vflow_ipfix_workers",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().Workers)
			})
	case *SFlow:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_sflow_workers",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().Workers)
			})
	case *NetflowV5:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv5_workers",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().Workers)
			})
	case *NetflowV9:
		promauto.NewCounterFunc(prometheus.CounterOpts{
			Name: "vflow_netflowv9_workers",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().Workers)
			})
	}
}

func promGaugeUDPMirrorQueue(p interface{}) {
	switch flow := p.(type) {
	case *IPFIX:
		promauto.NewGaugeFunc(prometheus.GaugeOpts{
			Name: "vflow_ipfix_udp_mirror_queue",
			Help: "",
		},
			func() float64 {
				return float64(flow.status().UDPMirrorQueue)
			})
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net"
	"regexp"

	"github.com/czerwonk/dns-drain/changelog"
	"github.com/czerwonk/dns-drain/gcloud"
)

type DrainActionFunc func(Drainer) error

func drain() error {
	if len(*value) > 0 {
		return drainWithValue()
	}

	if len(*regexString) > 0 {
		return drainWithRegex()
	}

	return drainWithIpNet()
}

func drainWithRegex() error {
	regex, err := regexp.Compile(*regexString)
	if err != nil {
		return err
	}

	actionFunc := func(d Drainer) error {
		return d.DrainWithRegex(regex, *newValue)
	}
	return performDrain(actionFunc)
}

func drainWithValue() error {
	actionFunc := func(d Drainer) error {
		return d.DrainWithValue(*value, *newValue)
	}
	return performDrain(actionFunc)
}

func drainWithIpNet() error {
	ipNet, err := getNetFromIp()
	if err != nil {
		return err
	}

	var newIp net.IP
	if len(*newIpStr) > 0 {
		newIp = net.ParseIP(*newIpStr)
	}

	actionFunc := func(d Drainer) error {
		return d.DrainWithIpNet(ipNet, newIp)
	}
	return performDrain(actionFunc)
}

func getNetFromIp() (*net.IPNet, error) {
	_, ipNet, err := net.ParseCIDR(*ip)

	if err != nil {
		ipAddr := net.ParseIP(*ip)
		if len(ipAddr) == 0 {
			return nil, err
		}

		var e error
		if ipAddr.To4() != nil {
			_, ipNet, e = net.ParseCIDR(fmt.Sprintf("%s/32", ipAddr))
		} else {
			_, ipNet, e = net.ParseCIDR(fmt.Sprintf("%s/128", ipAddr))
		}

		if e == nil {
			err = nil
		}
	}

	return ipNet, err
}

func performDrain(actionFunc DrainActionFunc) error {
	logger, err := changelog.NewFileChangeLogger(*file)
	if err != nil {
		return err
	}
	defer flushAndCloseLogger(logger)

	c := gcloud.NewDrainer(*gcloudProject, *dry, zoneFilterRegex, skipFilterRegex, *typeFilter, logger)
	return actionFunc(c)
}

func flushAndCloseLogger(logger *changelog.FileChangeLogger) {
	err := logger.Flush()
	if err != nil {
		log.Printf("ERROR - %s\n", err)
	}

	err = logger.Close()
	if err != nil {
		log.Printf("ERROR - %s\n", err)
	}
}
<commit_msg>check if mode is ambigous<commit_after>package main

import (
	"errors"
	"fmt"
	"log"
	"net"
	"regexp"

	"github.com/czerwonk/dns-drain/changelog"
	"github.com/czerwonk/dns-drain/gcloud"
)

type DrainActionFunc func(Drainer) error

func drain() error {
	if modeIsAmbigous() {
		return errors.New("Mode is ambigous. Please use either ip, value or regex parameter.")
	}

	if len(*value) > 0 {
		return drainWithValue()
	}

	if len(*regexString) > 0 {
		return drainWithRegex()
	}

	return drainWithIpNet()
}

func modeIsAmbigous() bool {
	i := 0

	if len(*value) > 0 {
		i++
	}

	if len(*regexString) > 0 {
		i++
	}

	if len(*ip) > 0 {
		i++
	}

	return i != 1
}

func drainWithRegex() error {
	regex, err := regexp.Compile(*regexString)
	if err != nil {
		return err
	}

	actionFunc := func(d Drainer) error {
		return d.DrainWithRegex(regex, *newValue)
	}
	return performDrain(actionFunc)
}

func drainWithValue() error {
	actionFunc := func(d Drainer) error {
		return d.DrainWithValue(*value, *newValue)
	}
	return performDrain(actionFunc)
}

func drainWithIpNet() error {
	ipNet, err := getNetFromIp()
	if err != nil {
		return err
	}

	var newIp net.IP
	if len(*newIpStr) > 0 {
		newIp = net.ParseIP(*newIpStr)
	}

	actionFunc := func(d Drainer) error {
		return d.DrainWithIpNet(ipNet, newIp)
	}
	return performDrain(actionFunc)
}

func getNetFromIp() (*net.IPNet, error) {
	_, ipNet, err := net.ParseCIDR(*ip)

	if err != nil {
		ipAddr := net.ParseIP(*ip)
		if len(ipAddr) == 0 {
			return nil, err
		}

		var e error
		if ipAddr.To4() != nil {
			_, ipNet, e = net.ParseCIDR(fmt.Sprintf("%s/32", ipAddr))
		} else {
			_, ipNet, e = net.ParseCIDR(fmt.Sprintf("%s/128", ipAddr))
		}

		if e == nil {
			err = nil
		}
	}

	return ipNet, err
}

func performDrain(actionFunc DrainActionFunc) error {
	logger, err := changelog.NewFileChangeLogger(*file)
	if err != nil {
		return err
	}
	defer flushAndCloseLogger(logger)

	c := gcloud.NewDrainer(*gcloudProject, *dry, zoneFilterRegex, skipFilterRegex, *typeFilter, logger)
	return actionFunc(c)
}

func flushAndCloseLogger(logger *changelog.FileChangeLogger) {
	err := logger.Flush()
	if err != nil {
		log.Printf("ERROR - %s\n", err)
	}

	err = logger.Close()
	if err != nil {
		log.Printf("ERROR - %s\n", err)
	}
}
<|endoftext|>


<commit_before>// Package redsync provides a Redis-based distributed mutual exclusion lock implementation as described in the blog post http://antirez.com/news/77.
//
// Values containing the types defined in this package should not be copied.
package redsync

import (
	"crypto/rand"
	"encoding/base64"
	"errors"
	"net"
	"sync"
	"time"

	"github.com/fzzy/radix/redis"
)

const (
	DefaultExpiry = 8 * time.Second
	DefaultTries  = 16
	DefaultDelay  = 512 * time.Millisecond
	DefaultFactor = 0.01
)

var (
	ErrFailed = errors.New("failed to acquire lock")
)

type Locker interface {
	Lock() error
	Unlock()
}

// A Mutex is a mutual exclusion lock.
//
// Fields of a Mutex must not be changed after first use.
type Mutex struct {
	Name   string        // Resouce name
	Expiry time.Duration // Duration for which the lock is valid, DefaultExpiry if 0

	Tries int           // Number of attempts to acquire lock before admitting failure, DefaultTries if 0
	Delay time.Duration // Delay between two attempts to acquire lock, DefaultDelay if 0

	Factor float64 // Drift factor, DefaultFactor if 0

	value string
	until time.Time

	nodes []*redis.Client
	nodem sync.Mutex
}

var _ = Locker(&Mutex{})

// NewMutex returns a new Mutex on a named resource connected to the Redis instances at given addresses.
func NewMutex(name string, addrs []net.Addr) (*Mutex, error) {
	if len(addrs) == 0 {
		panic("redsync: addrs is empty")
	}

	nodes := []*redis.Client{}
	for _, addr := range addrs {
		node, _ := redis.Dial(addr.Network(), addr.String())
		nodes = append(nodes, node)
	}

	return &Mutex{
		Name:  name,
		nodes: nodes,
	}, nil
}

// Lock locks m.
// In case it returns an error on failure, you may retry to acquire the lock by calling this method again.
func (m *Mutex) Lock() error {
	m.nodem.Lock()
	defer m.nodem.Unlock()

	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		return err
	}
	value := base64.StdEncoding.EncodeToString(b)

	expiry := m.Expiry
	if expiry == 0 {
		expiry = DefaultExpiry
	}

	retries := m.Tries
	if retries == 0 {
		retries = DefaultTries
	}

	for i := 0; i < retries; i++ {
		n := 0
		start := time.Now()
		for _, node := range m.nodes {
			if node == nil {
				continue
			}

			reply := node.Cmd("set", m.Name, value, "nx", "px", int(expiry/time.Millisecond))
			if reply.Err != nil {
				continue
			}
			if reply.String() != "OK" {
				continue
			}
			n += 1
		}

		factor := m.Factor
		if factor == 0 {
			factor = DefaultFactor
		}

		until := time.Now().Add(m.Expiry - time.Now().Sub(start) - time.Duration(int64(float64(m.Expiry)*factor)) + 2*time.Millisecond)
		if n >= len(m.nodes)/2+1 && time.Now().Before(until) {
			m.value = value
			m.until = until
			return nil
		} else {
			for _, node := range m.nodes {
				if node == nil {
					continue
				}

				reply := node.Cmd("eval", `
					if redis.call("get", KEYS[1]) == ARGV[1] then
					    return redis.call("del", KEYS[1])
					else
					    return 0
					end
				`, 1, m.Name, value)
				if reply.Err != nil {
					continue
				}
			}
		}

		delay := m.Delay
		if delay == 0 {
			delay = DefaultDelay
		}
		time.Sleep(delay)
	}

	return ErrFailed
}

// Unlock unlocks m.
// It is a run-time error if m is not locked on entry to Unlock.
func (m *Mutex) Unlock() {
	m.nodem.Lock()
	defer m.nodem.Unlock()

	value := m.value
	if value == "" {
		panic("redsync: unlock of unlocked mutex")
	}

	m.value = ""
	m.until = time.Unix(0, 0)

	for _, node := range m.nodes {
		if node == nil {
			continue
		}

		node.Cmd("eval", `
			if redis.call("get", KEYS[1]) == ARGV[1] then
			    return redis.call("del", KEYS[1])
			else
			    return 0
			end
		`, 1, m.Name, value)
	}
}
<commit_msg>Make quorum tweakable; resolves #1<commit_after>// Package redsync provides a Redis-based distributed mutual exclusion lock implementation as described in the blog post http://antirez.com/news/77.
//
// Values containing the types defined in this package should not be copied.
package redsync

import (
	"crypto/rand"
	"encoding/base64"
	"errors"
	"net"
	"sync"
	"time"

	"github.com/fzzy/radix/redis"
)

const (
	DefaultExpiry = 8 * time.Second
	DefaultTries  = 16
	DefaultDelay  = 512 * time.Millisecond
	DefaultFactor = 0.01
)

var (
	ErrFailed = errors.New("failed to acquire lock")
)

type Locker interface {
	Lock() error
	Unlock()
}

// A Mutex is a mutual exclusion lock.
//
// Fields of a Mutex must not be changed after first use.
type Mutex struct {
	Name   string        // Resouce name
	Expiry time.Duration // Duration for which the lock is valid, DefaultExpiry if 0

	Tries int           // Number of attempts to acquire lock before admitting failure, DefaultTries if 0
	Delay time.Duration // Delay between two attempts to acquire lock, DefaultDelay if 0

	Factor float64 // Drift factor, DefaultFactor if 0

	Quorum int // Quorum for the lock, set to len(addrs)/2+1 by NewMutex()

	value string
	until time.Time

	nodes []*redis.Client
	nodem sync.Mutex
}

var _ = Locker(&Mutex{})

// NewMutex returns a new Mutex on a named resource connected to the Redis instances at given addresses.
func NewMutex(name string, addrs []net.Addr) (*Mutex, error) {
	if len(addrs) == 0 {
		panic("redsync: addrs is empty")
	}

	nodes := []*redis.Client{}
	for _, addr := range addrs {
		node, _ := redis.Dial(addr.Network(), addr.String())
		nodes = append(nodes, node)
	}

	return &Mutex{
		Name:   name,
		Quorum: len(addrs)/2 + 1,
		nodes:  nodes,
	}, nil
}

// Lock locks m.
// In case it returns an error on failure, you may retry to acquire the lock by calling this method again.
func (m *Mutex) Lock() error {
	m.nodem.Lock()
	defer m.nodem.Unlock()

	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		return err
	}
	value := base64.StdEncoding.EncodeToString(b)

	expiry := m.Expiry
	if expiry == 0 {
		expiry = DefaultExpiry
	}

	retries := m.Tries
	if retries == 0 {
		retries = DefaultTries
	}

	for i := 0; i < retries; i++ {
		n := 0
		start := time.Now()
		for _, node := range m.nodes {
			if node == nil {
				continue
			}

			reply := node.Cmd("set", m.Name, value, "nx", "px", int(expiry/time.Millisecond))
			if reply.Err != nil {
				continue
			}
			if reply.String() != "OK" {
				continue
			}
			n += 1
		}

		factor := m.Factor
		if factor == 0 {
			factor = DefaultFactor
		}

		until := time.Now().Add(m.Expiry - time.Now().Sub(start) - time.Duration(int64(float64(m.Expiry)*factor)) + 2*time.Millisecond)
		if n >= m.Quorum && time.Now().Before(until) {
			m.value = value
			m.until = until
			return nil
		} else {
			for _, node := range m.nodes {
				if node == nil {
					continue
				}

				reply := node.Cmd("eval", `
					if redis.call("get", KEYS[1]) == ARGV[1] then
					    return redis.call("del", KEYS[1])
					else
					    return 0
					end
				`, 1, m.Name, value)
				if reply.Err != nil {
					continue
				}
			}
		}

		delay := m.Delay
		if delay == 0 {
			delay = DefaultDelay
		}
		time.Sleep(delay)
	}

	return ErrFailed
}

// Unlock unlocks m.
// It is a run-time error if m is not locked on entry to Unlock.
func (m *Mutex) Unlock() {
	m.nodem.Lock()
	defer m.nodem.Unlock()

	value := m.value
	if value == "" {
		panic("redsync: unlock of unlocked mutex")
	}

	m.value = ""
	m.until = time.Unix(0, 0)

	for _, node := range m.nodes {
		if node == nil {
			continue
		}

		node.Cmd("eval", `
			if redis.call("get", KEYS[1]) == ARGV[1] then
			    return redis.call("del", KEYS[1])
			else
			    return 0
			end
		`, 1, m.Name, value)
	}
}
<|endoftext|>


<commit_before>package http

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

	"github.com/tonto/kit/http/middleware"

	"github.com/tonto/kit/http/respond"
)

type validator interface {
	Validate() error
}

// BaseService represents base http service
type BaseService struct {
	endpoints Endpoints
	mw        []middleware.Adapter
}

// Prefix returns service routing prefix (implements Service interface)
func (b *BaseService) Prefix() string {
	return "/"
}

// Endpoints returns all registered endpoints (implements Service interface)
func (b *BaseService) Endpoints() Endpoints {
	for _, e := range b.endpoints {
		if b.mw != nil {
			e.Handler = middleware.Adapt(e.Handler, b.mw...)
		}
	}
	return b.endpoints
}

// RegisterEndpoint is a helper method that registers service endpoint
func (b *BaseService) RegisterEndpoint(path string, h http.Handler, methods ...string) {
	if b.endpoints == nil {
		b.endpoints = make(map[string]*Endpoint)
	}
	b.endpoints[path] = &Endpoint{
		Methods: methods,
		Handler: h,
	}
}

// RegisterMiddleware is a helper method that registers provided middlewares
// for service wide usage, ie. provided middlewares are applied to all endpoints
func (b *BaseService) RegisterMiddleware(mw ...middleware.Adapter) {
	b.mw = mw
}

// HandlerFromMethod creates new handler from a given service method.
// Required request struct will be recognised and request body will be
// correctly unmarshaled to it.
func (b *BaseService) HandlerFromMethod(m interface{}) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		v := reflect.ValueOf(m)
		reqParamType := v.Type().In(2).Elem()
		req := reflect.New(reqParamType).Interface()

		err := json.NewDecoder(r.Body).Decode(req)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		defer r.Body.Close()

		if validator, ok := interface{}(req).(validator); ok {
			err = validator.Validate()
			if err != nil {
				respond.With(w, r, http.StatusBadRequest, err)
				return
			}
		}

		v.Call([]reflect.Value{
			reflect.ValueOf(w),
			reflect.ValueOf(r),
			reflect.ValueOf(req),
		})
	})
}
<commit_msg>json decoder writes err<commit_after>package http

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

	"github.com/tonto/kit/http/middleware"

	"github.com/tonto/kit/http/respond"
)

type validator interface {
	Validate() error
}

// BaseService represents base http service
type BaseService struct {
	endpoints Endpoints
	mw        []middleware.Adapter
}

// Prefix returns service routing prefix (implements Service interface)
func (b *BaseService) Prefix() string {
	return "/"
}

// Endpoints returns all registered endpoints (implements Service interface)
func (b *BaseService) Endpoints() Endpoints {
	for _, e := range b.endpoints {
		if b.mw != nil {
			e.Handler = middleware.Adapt(e.Handler, b.mw...)
		}
	}
	return b.endpoints
}

// RegisterEndpoint is a helper method that registers service endpoint
func (b *BaseService) RegisterEndpoint(path string, h http.Handler, methods ...string) {
	if b.endpoints == nil {
		b.endpoints = make(map[string]*Endpoint)
	}
	b.endpoints[path] = &Endpoint{
		Methods: methods,
		Handler: h,
	}
}

// RegisterMiddleware is a helper method that registers provided middlewares
// for service wide usage, ie. provided middlewares are applied to all endpoints
func (b *BaseService) RegisterMiddleware(mw ...middleware.Adapter) {
	b.mw = mw
}

// HandlerFromMethod creates new handler from a given service method.
// Required request struct will be recognised and request body will be
// correctly unmarshaled to it.
func (b *BaseService) HandlerFromMethod(m interface{}) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		v := reflect.ValueOf(m)
		reqParamType := v.Type().In(2).Elem()
		req := reflect.New(reqParamType).Interface()

		err := json.NewDecoder(r.Body).Decode(req)
		if err != nil {
			respond.With(w, r, http.StatusBadRequest, err)
			return
		}

		defer r.Body.Close()

		if validator, ok := interface{}(req).(validator); ok {
			err = validator.Validate()
			if err != nil {
				respond.With(w, r, http.StatusBadRequest, err)
				return
			}
		}

		v.Call([]reflect.Value{
			reflect.ValueOf(w),
			reflect.ValueOf(r),
			reflect.ValueOf(req),
		})
	})
}
<|endoftext|>


<commit_before>package http3

import (
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/lucas-clemente/quic-go"
	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/marten-seemann/qpack"
)

// allows mocking of quic.Listen and quic.ListenAddr
var (
	quicListen     = quic.Listen
	quicListenAddr = quic.ListenAddr
)

// Server is a HTTP2 server listening for QUIC connections.
type Server struct {
	*http.Server

	// By providing a quic.Config, it is possible to set parameters of the QUIC connection.
	// If nil, it uses reasonable default values.
	QuicConfig *quic.Config

	port uint32 // used atomically

	listenerMutex sync.Mutex
	listener      quic.Listener
	closed        bool

	supportedVersionsAsString string

	logger utils.Logger
}

// ListenAndServe listens on the UDP address s.Addr and calls s.Handler to handle HTTP/3 requests on incoming connections.
func (s *Server) ListenAndServe() error {
	if s.Server == nil {
		return errors.New("use of http3.Server without http.Server")
	}
	return s.serveImpl(s.TLSConfig, nil)
}

// ListenAndServeTLS listens on the UDP address s.Addr and calls s.Handler to handle HTTP/3 requests on incoming connections.
func (s *Server) ListenAndServeTLS(certFile, keyFile string) error {
	var err error
	certs := make([]tls.Certificate, 1)
	certs[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}
	// We currently only use the cert-related stuff from tls.Config,
	// so we don't need to make a full copy.
	config := &tls.Config{
		Certificates: certs,
	}
	return s.serveImpl(config, nil)
}

// Serve an existing UDP connection.
func (s *Server) Serve(conn net.PacketConn) error {
	return s.serveImpl(s.TLSConfig, conn)
}

func (s *Server) serveImpl(tlsConfig *tls.Config, conn net.PacketConn) error {
	if s.Server == nil {
		return errors.New("use of http3.Server without http.Server")
	}
	s.logger = utils.DefaultLogger.WithPrefix("server")
	s.listenerMutex.Lock()
	if s.closed {
		s.listenerMutex.Unlock()
		return errors.New("Server is already closed")
	}
	if s.listener != nil {
		s.listenerMutex.Unlock()
		return errors.New("ListenAndServe may only be called once")
	}

	var ln quic.Listener
	var err error
	if conn == nil {
		ln, err = quicListenAddr(s.Addr, tlsConfig, s.QuicConfig)
	} else {
		ln, err = quicListen(conn, tlsConfig, s.QuicConfig)
	}
	if err != nil {
		s.listenerMutex.Unlock()
		return err
	}
	s.listener = ln
	s.listenerMutex.Unlock()

	for {
		sess, err := ln.Accept()
		if err != nil {
			return err
		}
		go s.handleConn(sess)
	}
}

func (s *Server) handleConn(sess quic.Session) {
	// TODO: accept control streams
	decoder := qpack.NewDecoder(nil)

	for {
		str, err := sess.AcceptStream()
		if err != nil {
			s.logger.Debugf("Accepting stream failed: %s", err)
			return
		}
		// TODO: handle error
		go func() {
			if err := s.handleRequest(str, decoder); err != nil {
				s.logger.Debugf("Handling request failed: %s", err)
				str.CancelWrite(quic.ErrorCode(errorGeneralProtocolError))
				return
			}
			str.Close()
		}()
	}
}

// TODO: improve error handling.
// Most (but not all) of the errors occurring here are connection-level erros.
func (s *Server) handleRequest(str quic.Stream, decoder *qpack.Decoder) error {
	frame, err := parseNextFrame(str)
	if err != nil {
		str.CancelWrite(quic.ErrorCode(errorRequestCanceled))
		return err
	}
	hf, ok := frame.(*headersFrame)
	if !ok {
		str.CancelWrite(quic.ErrorCode(errorUnexpectedFrame))
		return errors.New("expected first frame to be a headers frame")
	}
	// TODO: check length
	headerBlock := make([]byte, hf.Length)
	if _, err := io.ReadFull(str, headerBlock); err != nil {
		str.CancelWrite(quic.ErrorCode(errorIncompleteRequest))
		return err
	}
	hfs, err := decoder.DecodeFull(headerBlock)
	if err != nil {
		// TODO: use the right error code
		str.CancelWrite(quic.ErrorCode(errorGeneralProtocolError))
		return err
	}
	req, err := requestFromHeaders(hfs)
	if err != nil {
		return err
	}
	req.Body = newRequestBody(str)

	if s.logger.Debug() {
		s.logger.Infof("%s %s%s, on stream %d", req.Method, req.Host, req.RequestURI, str.StreamID())
	} else {
		s.logger.Infof("%s %s%s", req.Method, req.Host, req.RequestURI)
	}

	req = req.WithContext(str.Context())
	responseWriter := newResponseWriter(str, s.logger)
	handler := s.Handler
	if handler == nil {
		handler = http.DefaultServeMux
	}

	var panicked, readEOF bool
	func() {
		defer func() {
			if p := recover(); p != nil {
				// Copied from net/http/server.go
				const size = 64 << 10
				buf := make([]byte, size)
				buf = buf[:runtime.Stack(buf, false)]
				s.logger.Errorf("http: panic serving: %v\n%s", p, buf)
				panicked = true
			}
		}()
		handler.ServeHTTP(responseWriter, req)
		// read the eof
		if _, err = str.Read([]byte{}); err == io.EOF {
			readEOF = true
		}
	}()

	if panicked {
		responseWriter.WriteHeader(500)
	} else {
		responseWriter.WriteHeader(200)
	}

	if !readEOF {
		str.CancelRead(quic.ErrorCode(errorEarlyResponse))
	}
	return nil
}

// Close the server immediately, aborting requests and sending CONNECTION_CLOSE frames to connected clients.
// Close in combination with ListenAndServe() (instead of Serve()) may race if it is called before a UDP socket is established.
func (s *Server) Close() error {
	s.listenerMutex.Lock()
	defer s.listenerMutex.Unlock()
	s.closed = true
	if s.listener != nil {
		err := s.listener.Close()
		s.listener = nil
		return err
	}
	return nil
}

// CloseGracefully shuts down the server gracefully. The server sends a GOAWAY frame first, then waits for either timeout to trigger, or for all running requests to complete.
// CloseGracefully in combination with ListenAndServe() (instead of Serve()) may race if it is called before a UDP socket is established.
func (s *Server) CloseGracefully(timeout time.Duration) error {
	// TODO: implement
	return nil
}

// SetQuicHeaders can be used to set the proper headers that announce that this server supports QUIC.
// The values that are set depend on the port information from s.Server.Addr, and currently look like this (if Addr has port 443):
//  Alt-Svc: quic=":443"; ma=2592000; v="33,32,31,30"
func (s *Server) SetQuicHeaders(hdr http.Header) error {
	port := atomic.LoadUint32(&s.port)

	if port == 0 {
		// Extract port from s.Server.Addr
		_, portStr, err := net.SplitHostPort(s.Server.Addr)
		if err != nil {
			return err
		}
		portInt, err := net.LookupPort("tcp", portStr)
		if err != nil {
			return err
		}
		port = uint32(portInt)
		atomic.StoreUint32(&s.port, port)
	}

	if s.supportedVersionsAsString == "" {
		var versions []string
		for _, v := range protocol.SupportedVersions {
			versions = append(versions, v.ToAltSvc())
		}
		s.supportedVersionsAsString = strings.Join(versions, ",")
	}

	hdr.Add("Alt-Svc", fmt.Sprintf(`quic=":%d"; ma=2592000; v="%s"`, port, s.supportedVersionsAsString))

	return nil
}

// ListenAndServeQUIC listens on the UDP network address addr and calls the
// handler for HTTP/3 requests on incoming connections. http.DefaultServeMux is
// used when handler is nil.
func ListenAndServeQUIC(addr, certFile, keyFile string, handler http.Handler) error {
	server := &Server{
		Server: &http.Server{
			Addr:    addr,
			Handler: handler,
		},
	}
	return server.ListenAndServeTLS(certFile, keyFile)
}

// ListenAndServe listens on the given network address for both, TLS and QUIC
// connetions in parallel. It returns if one of the two returns an error.
// http.DefaultServeMux is used when handler is nil.
// The correct Alt-Svc headers for QUIC are set.
func ListenAndServe(addr, certFile, keyFile string, handler http.Handler) error {
	// Load certs
	var err error
	certs := make([]tls.Certificate, 1)
	certs[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}
	// We currently only use the cert-related stuff from tls.Config,
	// so we don't need to make a full copy.
	config := &tls.Config{
		Certificates: certs,
	}

	// Open the listeners
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return err
	}
	udpConn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		return err
	}
	defer udpConn.Close()

	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		return err
	}
	tcpConn, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		return err
	}
	defer tcpConn.Close()

	tlsConn := tls.NewListener(tcpConn, config)
	defer tlsConn.Close()

	// Start the servers
	httpServer := &http.Server{
		Addr:      addr,
		TLSConfig: config,
	}

	quicServer := &Server{
		Server: httpServer,
	}

	if handler == nil {
		handler = http.DefaultServeMux
	}
	httpServer.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		quicServer.SetQuicHeaders(w.Header())
		handler.ServeHTTP(w, r)
	})

	hErr := make(chan error)
	qErr := make(chan error)
	go func() {
		hErr <- httpServer.Serve(tlsConn)
	}()
	go func() {
		qErr <- quicServer.Serve(udpConn)
	}()

	select {
	case err := <-hErr:
		quicServer.Close()
		return err
	case err := <-qErr:
		// Cannot close the HTTP server or wait for requests to complete properly :/
		return err
	}
}
<commit_msg>fix reading of the EOF in the HTTP/3 server<commit_after>package http3

import (
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/lucas-clemente/quic-go"
	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/marten-seemann/qpack"
)

// allows mocking of quic.Listen and quic.ListenAddr
var (
	quicListen     = quic.Listen
	quicListenAddr = quic.ListenAddr
)

// Server is a HTTP2 server listening for QUIC connections.
type Server struct {
	*http.Server

	// By providing a quic.Config, it is possible to set parameters of the QUIC connection.
	// If nil, it uses reasonable default values.
	QuicConfig *quic.Config

	port uint32 // used atomically

	listenerMutex sync.Mutex
	listener      quic.Listener
	closed        bool

	supportedVersionsAsString string

	logger utils.Logger
}

// ListenAndServe listens on the UDP address s.Addr and calls s.Handler to handle HTTP/3 requests on incoming connections.
func (s *Server) ListenAndServe() error {
	if s.Server == nil {
		return errors.New("use of http3.Server without http.Server")
	}
	return s.serveImpl(s.TLSConfig, nil)
}

// ListenAndServeTLS listens on the UDP address s.Addr and calls s.Handler to handle HTTP/3 requests on incoming connections.
func (s *Server) ListenAndServeTLS(certFile, keyFile string) error {
	var err error
	certs := make([]tls.Certificate, 1)
	certs[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}
	// We currently only use the cert-related stuff from tls.Config,
	// so we don't need to make a full copy.
	config := &tls.Config{
		Certificates: certs,
	}
	return s.serveImpl(config, nil)
}

// Serve an existing UDP connection.
func (s *Server) Serve(conn net.PacketConn) error {
	return s.serveImpl(s.TLSConfig, conn)
}

func (s *Server) serveImpl(tlsConfig *tls.Config, conn net.PacketConn) error {
	if s.Server == nil {
		return errors.New("use of http3.Server without http.Server")
	}
	s.logger = utils.DefaultLogger.WithPrefix("server")
	s.listenerMutex.Lock()
	if s.closed {
		s.listenerMutex.Unlock()
		return errors.New("Server is already closed")
	}
	if s.listener != nil {
		s.listenerMutex.Unlock()
		return errors.New("ListenAndServe may only be called once")
	}

	var ln quic.Listener
	var err error
	if conn == nil {
		ln, err = quicListenAddr(s.Addr, tlsConfig, s.QuicConfig)
	} else {
		ln, err = quicListen(conn, tlsConfig, s.QuicConfig)
	}
	if err != nil {
		s.listenerMutex.Unlock()
		return err
	}
	s.listener = ln
	s.listenerMutex.Unlock()

	for {
		sess, err := ln.Accept()
		if err != nil {
			return err
		}
		go s.handleConn(sess)
	}
}

func (s *Server) handleConn(sess quic.Session) {
	// TODO: accept control streams
	decoder := qpack.NewDecoder(nil)

	for {
		str, err := sess.AcceptStream()
		if err != nil {
			s.logger.Debugf("Accepting stream failed: %s", err)
			return
		}
		// TODO: handle error
		go func() {
			if err := s.handleRequest(str, decoder); err != nil {
				s.logger.Debugf("Handling request failed: %s", err)
				str.CancelWrite(quic.ErrorCode(errorGeneralProtocolError))
				return
			}
			str.Close()
		}()
	}
}

// TODO: improve error handling.
// Most (but not all) of the errors occurring here are connection-level erros.
func (s *Server) handleRequest(str quic.Stream, decoder *qpack.Decoder) error {
	frame, err := parseNextFrame(str)
	if err != nil {
		str.CancelWrite(quic.ErrorCode(errorRequestCanceled))
		return err
	}
	hf, ok := frame.(*headersFrame)
	if !ok {
		str.CancelWrite(quic.ErrorCode(errorUnexpectedFrame))
		return errors.New("expected first frame to be a headers frame")
	}
	// TODO: check length
	headerBlock := make([]byte, hf.Length)
	if _, err := io.ReadFull(str, headerBlock); err != nil {
		str.CancelWrite(quic.ErrorCode(errorIncompleteRequest))
		return err
	}
	hfs, err := decoder.DecodeFull(headerBlock)
	if err != nil {
		// TODO: use the right error code
		str.CancelWrite(quic.ErrorCode(errorGeneralProtocolError))
		return err
	}
	req, err := requestFromHeaders(hfs)
	if err != nil {
		return err
	}
	req.Body = newRequestBody(str)

	if s.logger.Debug() {
		s.logger.Infof("%s %s%s, on stream %d", req.Method, req.Host, req.RequestURI, str.StreamID())
	} else {
		s.logger.Infof("%s %s%s", req.Method, req.Host, req.RequestURI)
	}

	req = req.WithContext(str.Context())
	responseWriter := newResponseWriter(str, s.logger)
	handler := s.Handler
	if handler == nil {
		handler = http.DefaultServeMux
	}

	var panicked, readEOF bool
	func() {
		defer func() {
			if p := recover(); p != nil {
				// Copied from net/http/server.go
				const size = 64 << 10
				buf := make([]byte, size)
				buf = buf[:runtime.Stack(buf, false)]
				s.logger.Errorf("http: panic serving: %v\n%s", p, buf)
				panicked = true
			}
		}()
		handler.ServeHTTP(responseWriter, req)
		// read the eof
		if _, err = str.Read([]byte{0}); err == io.EOF {
			readEOF = true
		}
	}()

	if panicked {
		responseWriter.WriteHeader(500)
	} else {
		responseWriter.WriteHeader(200)
	}

	if !readEOF {
		str.CancelRead(quic.ErrorCode(errorEarlyResponse))
	}
	return nil
}

// Close the server immediately, aborting requests and sending CONNECTION_CLOSE frames to connected clients.
// Close in combination with ListenAndServe() (instead of Serve()) may race if it is called before a UDP socket is established.
func (s *Server) Close() error {
	s.listenerMutex.Lock()
	defer s.listenerMutex.Unlock()
	s.closed = true
	if s.listener != nil {
		err := s.listener.Close()
		s.listener = nil
		return err
	}
	return nil
}

// CloseGracefully shuts down the server gracefully. The server sends a GOAWAY frame first, then waits for either timeout to trigger, or for all running requests to complete.
// CloseGracefully in combination with ListenAndServe() (instead of Serve()) may race if it is called before a UDP socket is established.
func (s *Server) CloseGracefully(timeout time.Duration) error {
	// TODO: implement
	return nil
}

// SetQuicHeaders can be used to set the proper headers that announce that this server supports QUIC.
// The values that are set depend on the port information from s.Server.Addr, and currently look like this (if Addr has port 443):
//  Alt-Svc: quic=":443"; ma=2592000; v="33,32,31,30"
func (s *Server) SetQuicHeaders(hdr http.Header) error {
	port := atomic.LoadUint32(&s.port)

	if port == 0 {
		// Extract port from s.Server.Addr
		_, portStr, err := net.SplitHostPort(s.Server.Addr)
		if err != nil {
			return err
		}
		portInt, err := net.LookupPort("tcp", portStr)
		if err != nil {
			return err
		}
		port = uint32(portInt)
		atomic.StoreUint32(&s.port, port)
	}

	if s.supportedVersionsAsString == "" {
		var versions []string
		for _, v := range protocol.SupportedVersions {
			versions = append(versions, v.ToAltSvc())
		}
		s.supportedVersionsAsString = strings.Join(versions, ",")
	}

	hdr.Add("Alt-Svc", fmt.Sprintf(`quic=":%d"; ma=2592000; v="%s"`, port, s.supportedVersionsAsString))

	return nil
}

// ListenAndServeQUIC listens on the UDP network address addr and calls the
// handler for HTTP/3 requests on incoming connections. http.DefaultServeMux is
// used when handler is nil.
func ListenAndServeQUIC(addr, certFile, keyFile string, handler http.Handler) error {
	server := &Server{
		Server: &http.Server{
			Addr:    addr,
			Handler: handler,
		},
	}
	return server.ListenAndServeTLS(certFile, keyFile)
}

// ListenAndServe listens on the given network address for both, TLS and QUIC
// connetions in parallel. It returns if one of the two returns an error.
// http.DefaultServeMux is used when handler is nil.
// The correct Alt-Svc headers for QUIC are set.
func ListenAndServe(addr, certFile, keyFile string, handler http.Handler) error {
	// Load certs
	var err error
	certs := make([]tls.Certificate, 1)
	certs[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}
	// We currently only use the cert-related stuff from tls.Config,
	// so we don't need to make a full copy.
	config := &tls.Config{
		Certificates: certs,
	}

	// Open the listeners
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return err
	}
	udpConn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		return err
	}
	defer udpConn.Close()

	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		return err
	}
	tcpConn, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		return err
	}
	defer tcpConn.Close()

	tlsConn := tls.NewListener(tcpConn, config)
	defer tlsConn.Close()

	// Start the servers
	httpServer := &http.Server{
		Addr:      addr,
		TLSConfig: config,
	}

	quicServer := &Server{
		Server: httpServer,
	}

	if handler == nil {
		handler = http.DefaultServeMux
	}
	httpServer.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		quicServer.SetQuicHeaders(w.Header())
		handler.ServeHTTP(w, r)
	})

	hErr := make(chan error)
	qErr := make(chan error)
	go func() {
		hErr <- httpServer.Serve(tlsConn)
	}()
	go func() {
		qErr <- quicServer.Serve(udpConn)
	}()

	select {
	case err := <-hErr:
		quicServer.Close()
		return err
	case err := <-qErr:
		// Cannot close the HTTP server or wait for requests to complete properly :/
		return err
	}
}
<|endoftext|>


<commit_before>package ogo

import (
	"bytes"
	"compress/flate"
	"compress/gzip"
	"encoding/json"
	"html/template"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/Odinman/ogo/utils"
)

/* {{{ func (rc *RESTContext) SetHeader(k,v string)
 * set header
 */
func (rc *RESTContext) SetHeader(k, v string) {
	rc.Response.Header().Set(k, v)
}

/* }}} */

/* {{{ func (rc *RESTContext) SetStatus(status int)
 * 构建Status Code
 */
func (rc *RESTContext) SetStatus(status int) {
	//rc.Response.WriteHeader(status)
	rc.Status = status
}

/* }}} */

/* {{{ func (rc *RESTContext) Output(data interface{}) (err error)
 *
 */
func (rc *RESTContext) Output(data interface{}) (err error) {

	if rc.Status >= 200 && rc.Status < 300 && rc.Accept == ContentTypeHTML { //用户需要HTML
		tplFile := ""
		// tpl file
		if ti, ok := rc.Route.Options[KEY_TPL]; ok && ti.(string) != "" && utils.FileExists(ti.(string)) { //定义了tpl文件, 并且文件存在
			tplFile = ti.(string)
		} else if dt := filepath.Join(env.TplDir, rc.Request.URL.Path+".html"); utils.FileExists(dt) { //默认tpl文件, 为: tpldir+url.Path+".html"
			tplFile = dt
		}
		if tplFile != "" {
			if t, err := template.ParseFiles(tplFile); err == nil {
				return t.Execute(rc.Response, data)
			}
		}
	}

	// 以下仍然返回json
	rc.SetHeader("Content-Type", "application/json; charset=UTF-8")
	var content []byte
	if method := strings.ToLower(rc.Request.Method); method != "head" {
		if data != nil {
			if env.IndentJSON {
				content, _ = json.MarshalIndent(data, "", "  ")
			} else {
				content, _ = json.Marshal(data)
			}
		}
	}
	//write header & data
	_, err = rc.WriteBytes(content)

	return
}

/* }}} */

/* {{{ func (rc *RESTContext) HTTPOK(data []byte) (err error)
 * 属于request的错误
 */
func (rc *RESTContext) HTTPOK(data []byte) (err error) {
	rc.SetHeader("Content-Type", "text/html; charset=UTF-8")
	rc.Status = http.StatusOK

	// write data
	_, err = rc.WriteBytes(data)
	return
}

/* }}} */

/* {{{ func (rc *RESTContext) HTTPEmptyGif() (err error)
 * 属于request的错误
 */
func (rc *RESTContext) HTTPEmptyGif() (err error) {
	rc.SetHeader("Content-Type", "image/gif")
	rc.Status = http.StatusOK

	// write data
	_, err = rc.WriteBytes(EmptyGifBytes)
	return
}

/* }}} */

/* {{{ func (rc *RESTContext) HTTPBack() (err error)
 * 属于request的错误
 */
func (rc *RESTContext) HTTPBack() (err error) {
	rc.Status = http.StatusOK
	rc.SetHeader("Content-Type", "text/html; charset=UTF-8")
	rc.SetHeader("Cache-Control", "max-age=0")
	rc.SetHeader("Cache-Control", "no-cache")
	rc.SetHeader("Cache-Control", "must-revalidate")
	rc.SetHeader("Cache-Control", "private")
	rc.SetHeader("Expires", "Mon, 26 Jul 1997 05:00:00 GMT")
	rc.SetHeader("Pragma", "no-cache")

	// write data
	data := []byte(`<?xml version="1.0"?>
<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Strict//EN" "DTD/xhtml1-strict.dtd">

<meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
<meta http-equiv="Cache-Control" content="max-age=0" forua="true" />
<meta http-equiv="Cache-Control" content="no-cache" forua="true" />
<meta http-equiv="Cache-Control" content="must-revalidate" forua="true" />
<title></title>

Back

`)
	_, err = rc.WriteBytes(data)
	return
}

/* }}} */

/* {{{ func (rc *RESTContext) HTTPRedirect(url string) (err error)
 * 属于request的错误
 */
func (rc *RESTContext) HTTPRedirect(url string) (err error) {
	rc.Status = http.StatusFound //302
	rc.SetHeader("Content-Type", "text/html; charset=UTF-8")
	rc.SetHeader("Cache-Control", "max-age=0")
	rc.SetHeader("Cache-Control", "no-cache")
	rc.SetHeader("Cache-Control", "must-revalidate")
	rc.SetHeader("Cache-Control", "private")
	rc.SetHeader("Expires", "Mon, 26 Jul 1997 05:00:00 GMT")
	rc.SetHeader("Pragma", "no-cache")
	rc.SetHeader("Location", url)

	_, err = rc.WriteBytes([]byte{})
	return
}

/* }}} */

/* {{{ func (rc *RESTContext) HTTPError(status int) (err error)
 *
 */
func (rc *RESTContext) HTTPError(status int) (err error) {

	rc.SetStatus(status)

	// write data
	err = rc.Output(rc.NewRESTError(status, nil))

	return
}

/* }}} */

/* {{{ func (rc *RESTContext) WriteBytes(data []byte) (n int, e error)
 * 输出内容,如果需要压缩,统一在这里进行
 */
func (rc *RESTContext) WriteBytes(data []byte) (n int, e error) {
	if dLen := len(data); dLen > 0 { //有内容才需要
		if env.EnableGzip == true && rc.Request.Header.Get("Accept-Encoding") != "" {
			splitted := strings.SplitN(rc.Request.Header.Get("Accept-Encoding"), ",", -1)
			encodings := make([]string, len(splitted))

			for i, val := range splitted {
				encodings[i] = strings.TrimSpace(val)
			}
			for _, val := range encodings {
				if val == "gzip" {
					rc.SetHeader("Content-Encoding", "gzip")
					b := new(bytes.Buffer)
					w, _ := gzip.NewWriterLevel(b, gzip.BestSpeed)
					w.Write(data)
					w.Close()
					data = b.Bytes()
					break
				} else if val == "deflate" {
					rc.SetHeader("Content-Encoding", "deflate")
					b := new(bytes.Buffer)
					w, _ := flate.NewWriter(b, flate.BestSpeed)
					w.Write(data)
					w.Close()
					data = b.Bytes()
					break
				}
			}
		}
		rc.ContentLength = dLen
		rc.SetHeader("Content-Length", strconv.Itoa(rc.ContentLength))
	}
	if rc.Status == 0 {
		rc.Status = http.StatusOK
	}
	//在Write之前要WriteHeader
	rc.Response.WriteHeader(rc.Status)
	if len(data) > 0 {
		_, e = rc.Response.Write(data)
	}

	return
}

/* }}} */

/* {{{ func (rc *RESTContext) ServeBinary(mimetype string, data []byte)
 * 直接出二进制内容
 */
func (rc *RESTContext) ServeBinary(mimetype string, data []byte) {
	rc.SetHeader("Content-Type", mimetype)
	rc.WriteBytes(data)
}

/* }}} */
<commit_msg>update<commit_after>package ogo

import (
	"bytes"
	"compress/flate"
	"compress/gzip"
	"encoding/json"
	"html/template"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/Odinman/ogo/utils"
)

/* {{{ func (rc *RESTContext) SetHeader(k,v string)
 * set header
 */
func (rc *RESTContext) SetHeader(k, v string) {
	rc.Response.Header().Set(k, v)
}

/* }}} */

/* {{{ func (rc *RESTContext) SetStatus(status int)
 * 构建Status Code
 */
func (rc *RESTContext) SetStatus(status int) {
	//rc.Response.WriteHeader(status)
	rc.Status = status
}

/* }}} */

/* {{{ func (rc *RESTContext) Output(data interface{}) (err error)
 *
 */
func (rc *RESTContext) Output(data interface{}) (err error) {

	if (rc.Status == 0 || (rc.Status >= 200 && rc.Status < 300)) && rc.Accept == ContentTypeHTML { //用户需要HTML
		tplFile := ""
		// tpl file
		if ti, ok := rc.Route.Options[KEY_TPL]; ok && ti.(string) != "" && utils.FileExists(ti.(string)) { //定义了tpl文件, 并且文件存在
			tplFile = ti.(string)
		} else if dt := filepath.Join(env.TplDir, rc.Request.URL.Path+".html"); utils.FileExists(dt) { //默认tpl文件, 为: tpldir+url.Path+".html"
			tplFile = dt
		}
		if tplFile != "" {
			if t, err := template.ParseFiles(tplFile); err == nil {
				return t.Execute(rc.Response, data)
			}
		}
	}

	// 以下仍然返回json
	rc.SetHeader("Content-Type", "application/json; charset=UTF-8")
	var content []byte
	if method := strings.ToLower(rc.Request.Method); method != "head" {
		if data != nil {
			if env.IndentJSON {
				content, _ = json.MarshalIndent(data, "", "  ")
			} else {
				content, _ = json.Marshal(data)
			}
		}
	}
	//write header & data
	_, err = rc.WriteBytes(content)

	return
}

/* }}} */

/* {{{ func (rc *RESTContext) HTTPOK(data []byte) (err error)
 * 属于request的错误
 */
func (rc *RESTContext) HTTPOK(data []byte) (err error) {
	rc.SetHeader("Content-Type", "text/html; charset=UTF-8")
	rc.Status = http.StatusOK

	// write data
	_, err = rc.WriteBytes(data)
	return
}

/* }}} */

/* {{{ func (rc *RESTContext) HTTPEmptyGif() (err error)
 * 属于request的错误
 */
func (rc *RESTContext) HTTPEmptyGif() (err error) {
	rc.SetHeader("Content-Type", "image/gif")
	rc.Status = http.StatusOK

	// write data
	_, err = rc.WriteBytes(EmptyGifBytes)
	return
}

/* }}} */

/* {{{ func (rc *RESTContext) HTTPBack() (err error)
 * 属于request的错误
 */
func (rc *RESTContext) HTTPBack() (err error) {
	rc.Status = http.StatusOK
	rc.SetHeader("Content-Type", "text/html; charset=UTF-8")
	rc.SetHeader("Cache-Control", "max-age=0")
	rc.SetHeader("Cache-Control", "no-cache")
	rc.SetHeader("Cache-Control", "must-revalidate")
	rc.SetHeader("Cache-Control", "private")
	rc.SetHeader("Expires", "Mon, 26 Jul 1997 05:00:00 GMT")
	rc.SetHeader("Pragma", "no-cache")

	// write data
	data := []byte(`<?xml version="1.0"?>
<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Strict//EN" "DTD/xhtml1-strict.dtd">

<meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
<meta http-equiv="Cache-Control" content="max-age=0" forua="true" />
<meta http-equiv="Cache-Control" content="no-cache" forua="true" />
<meta http-equiv="Cache-Control" content="must-revalidate" forua="true" />
<title></title>

Back

`)
	_, err = rc.WriteBytes(data)
	return
}

/* }}} */

/* {{{ func (rc *RESTContext) HTTPRedirect(url string) (err error)
 * 属于request的错误
 */
func (rc *RESTContext) HTTPRedirect(url string) (err error) {
	rc.Status = http.StatusFound //302
	rc.SetHeader("Content-Type", "text/html; charset=UTF-8")
	rc.SetHeader("Cache-Control", "max-age=0")
	rc.SetHeader("Cache-Control", "no-cache")
	rc.SetHeader("Cache-Control", "must-revalidate")
	rc.SetHeader("Cache-Control", "private")
	rc.SetHeader("Expires", "Mon, 26 Jul 1997 05:00:00 GMT")
	rc.SetHeader("Pragma", "no-cache")
	rc.SetHeader("Location", url)

	_, err = rc.WriteBytes([]byte{})
	return
}

/* }}} */

/* {{{ func (rc *RESTContext) HTTPError(status int) (err error)
 *
 */
func (rc *RESTContext) HTTPError(status int) (err error) {

	rc.SetStatus(status)

	// write data
	err = rc.Output(rc.NewRESTError(status, nil))

	return
}

/* }}} */

/* {{{ func (rc *RESTContext) WriteBytes(data []byte) (n int, e error)
 * 输出内容,如果需要压缩,统一在这里进行
 */
func (rc *RESTContext) WriteBytes(data []byte) (n int, e error) {
	if dLen := len(data); dLen > 0 { //有内容才需要
		if env.EnableGzip == true && rc.Request.Header.Get("Accept-Encoding") != "" {
			splitted := strings.SplitN(rc.Request.Header.Get("Accept-Encoding"), ",", -1)
			encodings := make([]string, len(splitted))

			for i, val := range splitted {
				encodings[i] = strings.TrimSpace(val)
			}
			for _, val := range encodings {
				if val == "gzip" {
					rc.SetHeader("Content-Encoding", "gzip")
					b := new(bytes.Buffer)
					w, _ := gzip.NewWriterLevel(b, gzip.BestSpeed)
					w.Write(data)
					w.Close()
					data = b.Bytes()
					break
				} else if val == "deflate" {
					rc.SetHeader("Content-Encoding", "deflate")
					b := new(bytes.Buffer)
					w, _ := flate.NewWriter(b, flate.BestSpeed)
					w.Write(data)
					w.Close()
					data = b.Bytes()
					break
				}
			}
		}
		rc.ContentLength = dLen
		rc.SetHeader("Content-Length", strconv.Itoa(rc.ContentLength))
	}
	if rc.Status == 0 {
		rc.Status = http.StatusOK
	}
	//在Write之前要WriteHeader
	rc.Response.WriteHeader(rc.Status)
	if len(data) > 0 {
		_, e = rc.Response.Write(data)
	}

	return
}

/* }}} */

/* {{{ func (rc *RESTContext) ServeBinary(mimetype string, data []byte)
 * 直接出二进制内容
 */
func (rc *RESTContext) ServeBinary(mimetype string, data []byte) {
	rc.SetHeader("Content-Type", mimetype)
	rc.WriteBytes(data)
}

/* }}} */
<|endoftext|>


<commit_before>package itchio

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/go-errors/errors"
	"github.com/mitchellh/mapstructure"
)

// Get performs an HTTP GET request to the API
func (c *Client) Get(url string) (*http.Response, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	return c.Do(req)
}

// PostForm performs an HTTP POST request to the API, with url-encoded parameters
func (c *Client) PostForm(url string, data url.Values) (*http.Response, error) {
	req, err := http.NewRequest("POST", url, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	return c.Do(req)
}

// Do performs a request (any method). It takes care of JWT or API key
// authentication, sets the propre user agent, has built-in retry,
func (c *Client) Do(req *http.Request) (*http.Response, error) {
	if strings.HasPrefix(c.Key, "jwt:") {
		req.Header.Add("Authorization", strings.Split(c.Key, ":")[1])
	}
	req.Header.Set("User-Agent", c.UserAgent)

	var res *http.Response
	var err error

	retryPatterns := append(c.RetryPatterns, time.Millisecond)

	for _, sleepTime := range retryPatterns {
		res, err = c.HTTPClient.Do(req)
		if err != nil {
			if strings.Contains(err.Error(), "TLS handshake timeout") {
				time.Sleep(sleepTime + time.Duration(rand.Int()%1000)*time.Millisecond)
				continue
			}
			return nil, err
		}

		if res.StatusCode == 503 {
			// Rate limited, try again according to patterns.
			// following https://cloud.google.com/storage/docs/json_api/v1/how-tos/upload#exp-backoff to the letter
			res.Body.Close()
			time.Sleep(sleepTime + time.Duration(rand.Int()%1000)*time.Millisecond)
			continue
		}

		break
	}

	return res, err
}

// MakePath crafts an API url from our configured base URL
func (c *Client) MakePath(format string, a ...interface{}) string {
	base := strings.Trim(c.BaseURL, "/")
	subPath := strings.Trim(fmt.Sprintf(format, a...), "/")

	var key string
	if strings.HasPrefix(c.Key, "jwt:") {
		key = "jwt"
	} else {
		key = c.Key
	}
	return fmt.Sprintf("%s/%s/%s", base, key, subPath)
}

// ParseAPIResponse unmarshals an HTTP response into one of out response
// data structures
func ParseAPIResponse(dst interface{}, res *http.Response) error {
	if res == nil || res.Body == nil {
		return fmt.Errorf("No response from server")
	}

	bodyReader := res.Body
	defer bodyReader.Close()

	if res.StatusCode/100 != 2 {
		return fmt.Errorf("Server returned %s for %s", res.Status, res.Request.URL.String())
	}

	body, err := ioutil.ReadAll(bodyReader)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	intermediate := make(map[string]interface{})

	err = json.NewDecoder(bytes.NewReader(body)).Decode(&intermediate)
	if err != nil {
		msg := fmt.Sprintf("JSON decode error: %s\n\nBody: %s\n\n", err.Error(), string(body))
		return errors.New(msg)
	}

	if errorsField, ok := intermediate["errors"]; ok {
		if errorsList, ok := errorsField.([]string); ok {
			if len(errorsList) > 0 {
				// TODO: handle multiple errors
				return fmt.Errorf("itch.io API error: %s", strings.Join(errorsList, ","))
			}
		}
	}

	intermediate = camelifyMap(intermediate)

	decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
		TagName: "json",
		Result:  dst,
	})
	if err != nil {
		return errors.Wrap(err, 0)
	}

	err = decoder.Decode(intermediate)
	if err != nil {
		msg := fmt.Sprintf("mapstructure decode error: %s\n\nBody: %#v\n\n", err.Error(), intermediate)
		return errors.New(msg)
	}

	return nil
}

// FindBuildFile looks for an uploaded file of the right type
// in a list of file. Returns nil if it can't find one.
func FindBuildFile(fileType BuildFileType, files []*BuildFile) *BuildFile {
	for _, f := range files {
		if f.Type == fileType && f.State == BuildFileStateUploaded {
			return f
		}
	}

	return nil
}

// ItchfsURL returns the itchfs:/// url usable to download a given file
// from a given build
func (build Build) ItchfsURL(file *BuildFile, apiKey string) string {
	return ItchfsURL(build.ID, file.ID, apiKey)
}

// ItchfsURL returns the itchfs:/// url usable to download a given file
// from a given build
func ItchfsURL(buildID int64, fileID int64, apiKey string) string {
	values := url.Values{}
	values.Set("api_key", apiKey)
	return fmt.Sprintf("itchfs:///wharf/builds/%d/files/%d/download?%s",
		buildID, fileID, values.Encode())
}
<commit_msg>Proper error handling + weak json typing. Closes https://github.com/itchio/itch/issues/1549<commit_after>package itchio

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/go-errors/errors"
	"github.com/mitchellh/mapstructure"
)

// Get performs an HTTP GET request to the API
func (c *Client) Get(url string) (*http.Response, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	return c.Do(req)
}

// PostForm performs an HTTP POST request to the API, with url-encoded parameters
func (c *Client) PostForm(url string, data url.Values) (*http.Response, error) {
	req, err := http.NewRequest("POST", url, strings.NewReader(data.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	return c.Do(req)
}

// Do performs a request (any method). It takes care of JWT or API key
// authentication, sets the propre user agent, has built-in retry,
func (c *Client) Do(req *http.Request) (*http.Response, error) {
	if strings.HasPrefix(c.Key, "jwt:") {
		req.Header.Add("Authorization", strings.Split(c.Key, ":")[1])
	}
	req.Header.Set("User-Agent", c.UserAgent)

	var res *http.Response
	var err error

	retryPatterns := append(c.RetryPatterns, time.Millisecond)

	for _, sleepTime := range retryPatterns {
		res, err = c.HTTPClient.Do(req)
		if err != nil {
			if strings.Contains(err.Error(), "TLS handshake timeout") {
				time.Sleep(sleepTime + time.Duration(rand.Int()%1000)*time.Millisecond)
				continue
			}
			return nil, err
		}

		if res.StatusCode == 503 {
			// Rate limited, try again according to patterns.
			// following https://cloud.google.com/storage/docs/json_api/v1/how-tos/upload#exp-backoff to the letter
			res.Body.Close()
			time.Sleep(sleepTime + time.Duration(rand.Int()%1000)*time.Millisecond)
			continue
		}

		break
	}

	return res, err
}

// MakePath crafts an API url from our configured base URL
func (c *Client) MakePath(format string, a ...interface{}) string {
	base := strings.Trim(c.BaseURL, "/")
	subPath := strings.Trim(fmt.Sprintf(format, a...), "/")

	var key string
	if strings.HasPrefix(c.Key, "jwt:") {
		key = "jwt"
	} else {
		key = c.Key
	}
	return fmt.Sprintf("%s/%s/%s", base, key, subPath)
}

// ParseAPIResponse unmarshals an HTTP response into one of out response
// data structures
func ParseAPIResponse(dst interface{}, res *http.Response) error {
	if res == nil || res.Body == nil {
		return fmt.Errorf("No response from server")
	}

	bodyReader := res.Body
	defer bodyReader.Close()

	if res.StatusCode/100 != 2 {
		return fmt.Errorf("Server returned %s for %s", res.Status, res.Request.URL.String())
	}

	body, err := ioutil.ReadAll(bodyReader)
	if err != nil {
		return errors.Wrap(err, 0)
	}

	intermediate := make(map[string]interface{})

	err = json.NewDecoder(bytes.NewReader(body)).Decode(&intermediate)
	if err != nil {
		msg := fmt.Sprintf("JSON decode error: %s\n\nBody: %s\n\n", err.Error(), string(body))
		return errors.New(msg)
	}

	if errorsField, ok := intermediate["errors"]; ok {
		if errorsList, ok := errorsField.([]interface{}); ok {
			var messages []string
			for _, el := range errorsList {
				if errorMessage, ok := el.(string); ok {
					messages = append(messages, errorMessage)
				}
			}
			if len(messages) > 0 {
				return fmt.Errorf("itch.io API error: %s", strings.Join(messages, ","))
			}
		}
	}

	intermediate = camelifyMap(intermediate)

	decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
		TagName: "json",
		Result:  dst,
		// see https://github.com/itchio/itch/issues/1549
		WeaklyTypedInput: true,
	})
	if err != nil {
		return errors.Wrap(err, 0)
	}

	err = decoder.Decode(intermediate)
	if err != nil {
		msg := fmt.Sprintf("mapstructure decode error: %s\n\nBody: %#v\n\n", err.Error(), intermediate)
		return errors.New(msg)
	}

	return nil
}

// FindBuildFile looks for an uploaded file of the right type
// in a list of file. Returns nil if it can't find one.
func FindBuildFile(fileType BuildFileType, files []*BuildFile) *BuildFile {
	for _, f := range files {
		if f.Type == fileType && f.State == BuildFileStateUploaded {
			return f
		}
	}

	return nil
}

// ItchfsURL returns the itchfs:/// url usable to download a given file
// from a given build
func (build Build) ItchfsURL(file *BuildFile, apiKey string) string {
	return ItchfsURL(build.ID, file.ID, apiKey)
}

// ItchfsURL returns the itchfs:/// url usable to download a given file
// from a given build
func ItchfsURL(buildID int64, fileID int64, apiKey string) string {
	values := url.Values{}
	values.Set("api_key", apiKey)
	return fmt.Sprintf("itchfs:///wharf/builds/%d/files/%d/download?%s",
		buildID, fileID, values.Encode())
}
<|endoftext|>


<commit_before>// Copyright 2017 The LUCI 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 buildbucket

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/golang/protobuf/ptypes"
	"github.com/golang/protobuf/ptypes/timestamp"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"go.chromium.org/gae/service/datastore"
	"go.chromium.org/luci/buildbucket"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	bbv1 "go.chromium.org/luci/common/api/buildbucket/buildbucket/v1"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/tsmon/field"
	"go.chromium.org/luci/common/tsmon/metric"
	"go.chromium.org/luci/milo/common"
	"go.chromium.org/luci/milo/common/model"
	"go.chromium.org/luci/server/auth"
	"go.chromium.org/luci/server/router"
	"google.golang.org/genproto/protobuf/field_mask"
)

var (
	buildCounter = metric.NewCounter(
		"luci/milo/buildbucket_pubsub/builds",
		"The number of buildbucket builds received by Milo from PubSub",
		nil,
		field.String("bucket"),
		// True for luci build, False for non-luci (ie buildbot) build.
		field.Bool("luci"),
		// Status can be "COMPLETED", "SCHEDULED", or "STARTED"
		field.String("status"),
		// Action can be one of 3 options.
		//   * "Created" - This is the first time Milo heard about this build
		//   * "Modified" - Milo updated some information about this build vs. what
		//     it knew before.
		//   * "Rejected" - Milo was unable to accept this build.
		field.String("action"))
)

// PubSubHandler is a webhook that stores the builds coming in from pubsub.
func PubSubHandler(ctx *router.Context) {
	err := pubSubHandlerImpl(ctx.Context, ctx.Request)
	if err != nil {
		logging.Errorf(ctx.Context, "error while handling pubsub event")
		errors.Log(ctx.Context, err)
	}
	if transient.Tag.In(err) {
		// Transient errors are 500 so that PubSub retries them.
		ctx.Writer.WriteHeader(http.StatusInternalServerError)
		return
	}
	// No errors or non-transient errors are 200s so that PubSub does not retry
	// them.
	ctx.Writer.WriteHeader(http.StatusOK)
}

func mustTimestamp(ts *timestamp.Timestamp) time.Time {
	if t, err := ptypes.Timestamp(ts); err == nil {
		return t
	}
	return time.Time{}
}

var summaryBuildMask = &field_mask.FieldMask{
	Paths: []string{
		"id",
		"builder",
		"number",
		"create_time",
		"start_time",
		"end_time",
		"update_time",
		"status",
		"summary_markdown",
		"tags",
		"infra.swarming",
	},
}

// getSummary returns a model.BuildSummary representing a buildbucket build.
func getSummary(c context.Context, host string, project string, id int64) (*model.BuildSummary, error) {
	client, err := buildbucketClient(c, host, auth.AsProject, auth.WithProject(project))
	if err != nil {
		return nil, err
	}
	b, err := client.GetBuild(c, &buildbucketpb.GetBuildRequest{
		Id:     id,
		Fields: summaryBuildMask,
	})
	if err != nil {
		return nil, err
	}
	buildAddress := fmt.Sprintf("%d", b.Id)
	if b.Number != 0 {
		buildAddress = fmt.Sprintf("luci.%s.%s/%s/%d", b.Builder.Project, b.Builder.Bucket, b.Builder.Builder, b.Number)
	}

	// Note: The parent for buildbucket build summaries is currently a fake entity.
	// In the future, builds can be cached here, but we currently don't do that.
	buildKey := datastore.MakeKey(c, "buildbucket.Build", fmt.Sprintf("%s:%s", host, buildAddress))
	swarming := b.GetInfra().GetSwarming()

	return &model.BuildSummary{
		ProjectID: b.Builder.Project,
		BuildKey:  buildKey,
		BuilderID: BuilderID{*b.Builder}.String(),
		BuildID:   "buildbucket/" + buildAddress,
		BuildSet:  b.Buildsets(),
		ContextURI: []string{
			fmt.Sprintf("buildbucket://%s/build/%d", host, id),
			fmt.Sprintf("swarming://%s/task/%s", swarming.GetHostname(), swarming.GetTaskId()),
		},
		Created: mustTimestamp(b.CreateTime),
		Summary: model.Summary{
			Start:  mustTimestamp(b.StartTime),
			End:    mustTimestamp(b.EndTime),
			Status: statusMap[b.Status],
		},
		Version:      mustTimestamp(b.UpdateTime).UnixNano(),
		Experimental: b.GetInput().GetExperimental(),
	}, nil
}

// generateSummary takes a decoded buildbucket event and generates
// a model.BuildSummary from it.
//
// This is the portion of the summarization process which cannot fail (i.e. is
// pure-data).
func generateSummary(c context.Context, hostname string, build buildbucket.Build) (*model.BuildSummary, error) {
	bs, err := getSummary(c, hostname, build.Project, build.ID)
	if err != nil {
		logging.WithError(err).Errorf(c, "got error while getting summary")
	}

	buildset := build.Tags[bbv1.TagBuildSet]
	if buildset == nil {
		buildset = []string{}
	}
	ret := &model.BuildSummary{
		ProjectID: build.Project,
		BuildKey:  MakeBuildKey(c, hostname, build.Address()),
		BuilderID: NewBuilderID(build.Bucket, build.Builder).String(),
		BuildID:   "buildbucket/" + build.Address(),
		BuildSet:  buildset,
		ContextURI: []string{
			fmt.Sprintf("buildbucket://%s/build/%d", hostname, build.ID),
		},

		Created: build.CreationTime,
		Summary: model.Summary{
			Start:  build.StartTime,
			End:    build.CompletionTime,
			Status: parseStatus(build.Status),
		},

		Version: build.UpdateTime.UnixNano(),

		Experimental: build.Experimental,
	}

	if shost, sid := build.Tags.Get("swarming_hostname"), build.Tags.Get("swarming_task_id"); shost != "" && sid != "" {
		ret.ContextURI = append(ret.ContextURI, fmt.Sprintf("swarming://%s/task/%s", shost, sid))
	}

	// Informational compare of v1 and v2 API result.
	if diff := cmp.Diff(ret, bs, cmpopts.IgnoreUnexported(model.BuildSummary{})); diff != "" {
		logging.Errorf(c, "BuildSummary of v2 has Diff (-v1, +v2)\n%s", diff)
	} else {
		logging.Debugf(c, "BuildSummary between v1 and v2 are the same.")
	}
	// TODO(iannucci,nodir): get the bot context too

	// TODO(iannucci,nodir): support manifests/got_revision
	return ret, ret.AddManifestKeysFromBuildSets(c)
}

// pubSubHandlerImpl takes the http.Request, expects to find
// a common.PubSubSubscription JSON object in the Body, containing a bbPSEvent,
// and handles the contents with generateSummary.
func pubSubHandlerImpl(c context.Context, r *http.Request) error {
	// This is the default action. The code below will modify the values of some
	// or all of these parameters.
	isLUCI, bucket, status, action := false, "UNKNOWN", "UNKNOWN", "Rejected"

	defer func() {
		// closure for late binding
		buildCounter.Add(c, 1, bucket, isLUCI, status, action)
	}()

	msg := common.PubSubSubscription{}
	if err := json.NewDecoder(r.Body).Decode(&msg); err != nil {
		// This might be a transient error, e.g. when the json format changes
		// and Milo isn't updated yet.
		return errors.Annotate(err, "could not decode message").Tag(transient.Tag).Err()
	}
	if v, ok := msg.Message.Attributes["version"].(string); ok && v != "v1" {
		// TODO(nodir): switch to v2, crbug.com/826006
		logging.Debugf(c, "unsupported pubsub message version %q. Ignoring", v)
		return nil
	}
	bData, err := msg.GetData()
	if err != nil {
		return errors.Annotate(err, "could not parse pubsub message string").Err()
	}

	event := struct {
		Build    bbv1.ApiCommonBuildMessage `json:"build"`
		Hostname string                     `json:"hostname"`
	}{}
	if err := json.Unmarshal(bData, &event); err != nil {
		return errors.Annotate(err, "could not parse pubsub message data").Err()
	}

	build := buildbucket.Build{}
	if err := build.ParseMessage(&event.Build); err != nil {
		return errors.Annotate(err, "could not parse buildbucket.Build").Err()
	}

	bucket = build.Bucket
	status = build.Status.String()
	isLUCI = strings.HasPrefix(bucket, "luci.")

	logging.Debugf(c, "Received from %s: build %s/%s (%s)\n%v",
		event.Hostname, bucket, build.Builder, status, build)

	if !isLUCI || build.Builder == "" {
		logging.Infof(c, "This is not an ingestable build, ignoring")
		return nil
	}

	bs, err := generateSummary(c, event.Hostname, build)
	if err != nil {
		return err
	}

	return transient.Tag.Apply(datastore.RunInTransaction(c, func(c context.Context) error {
		curBS := &model.BuildSummary{BuildKey: bs.BuildKey}
		switch err := datastore.Get(c, curBS); err {
		case datastore.ErrNoSuchEntity:
			action = "Created"
		case nil:
			action = "Modified"
		default:
			return errors.Annotate(err, "reading current BuildSummary").Err()
		}

		if build.UpdateTime.UnixNano() <= curBS.Version {
			logging.Warningf(c, "current BuildSummary is newer: %d <= %d",
				build.UpdateTime.UnixNano(), curBS.Version)
			return nil
		}

		if err := datastore.Put(c, bs); err != nil {
			return err
		}

		return model.UpdateBuilderForBuild(c, bs)
	}, &datastore.TransactionOptions{XG: true}))
}

// MakeBuildKey returns a new datastore Key for a buildbucket.Build.
//
// There's currently no model associated with this key, but it's used as
// a parent for a model.BuildSummary.
func MakeBuildKey(c context.Context, host, buildAddress string) *datastore.Key {
	return datastore.MakeKey(c,
		"buildbucket.Build", fmt.Sprintf("%s:%s", host, buildAddress))
}
<commit_msg>[milo] Add infra.experimental to field mask for PubSub v2<commit_after>// Copyright 2017 The LUCI 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 buildbucket

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/golang/protobuf/ptypes"
	"github.com/golang/protobuf/ptypes/timestamp"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"go.chromium.org/gae/service/datastore"
	"go.chromium.org/luci/buildbucket"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	bbv1 "go.chromium.org/luci/common/api/buildbucket/buildbucket/v1"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/tsmon/field"
	"go.chromium.org/luci/common/tsmon/metric"
	"go.chromium.org/luci/milo/common"
	"go.chromium.org/luci/milo/common/model"
	"go.chromium.org/luci/server/auth"
	"go.chromium.org/luci/server/router"
	"google.golang.org/genproto/protobuf/field_mask"
)

var (
	buildCounter = metric.NewCounter(
		"luci/milo/buildbucket_pubsub/builds",
		"The number of buildbucket builds received by Milo from PubSub",
		nil,
		field.String("bucket"),
		// True for luci build, False for non-luci (ie buildbot) build.
		field.Bool("luci"),
		// Status can be "COMPLETED", "SCHEDULED", or "STARTED"
		field.String("status"),
		// Action can be one of 3 options.
		//   * "Created" - This is the first time Milo heard about this build
		//   * "Modified" - Milo updated some information about this build vs. what
		//     it knew before.
		//   * "Rejected" - Milo was unable to accept this build.
		field.String("action"))
)

// PubSubHandler is a webhook that stores the builds coming in from pubsub.
func PubSubHandler(ctx *router.Context) {
	err := pubSubHandlerImpl(ctx.Context, ctx.Request)
	if err != nil {
		logging.Errorf(ctx.Context, "error while handling pubsub event")
		errors.Log(ctx.Context, err)
	}
	if transient.Tag.In(err) {
		// Transient errors are 500 so that PubSub retries them.
		ctx.Writer.WriteHeader(http.StatusInternalServerError)
		return
	}
	// No errors or non-transient errors are 200s so that PubSub does not retry
	// them.
	ctx.Writer.WriteHeader(http.StatusOK)
}

func mustTimestamp(ts *timestamp.Timestamp) time.Time {
	if t, err := ptypes.Timestamp(ts); err == nil {
		return t
	}
	return time.Time{}
}

var summaryBuildMask = &field_mask.FieldMask{
	Paths: []string{
		"id",
		"builder",
		"number",
		"create_time",
		"start_time",
		"end_time",
		"update_time",
		"status",
		"summary_markdown",
		"tags",
		"infra.swarming",
		"input.experimental",
	},
}

// getSummary returns a model.BuildSummary representing a buildbucket build.
func getSummary(c context.Context, host string, project string, id int64) (*model.BuildSummary, error) {
	client, err := buildbucketClient(c, host, auth.AsProject, auth.WithProject(project))
	if err != nil {
		return nil, err
	}
	b, err := client.GetBuild(c, &buildbucketpb.GetBuildRequest{
		Id:     id,
		Fields: summaryBuildMask,
	})
	if err != nil {
		return nil, err
	}
	buildAddress := fmt.Sprintf("%d", b.Id)
	if b.Number != 0 {
		buildAddress = fmt.Sprintf("luci.%s.%s/%s/%d", b.Builder.Project, b.Builder.Bucket, b.Builder.Builder, b.Number)
	}

	// Note: The parent for buildbucket build summaries is currently a fake entity.
	// In the future, builds can be cached here, but we currently don't do that.
	buildKey := datastore.MakeKey(c, "buildbucket.Build", fmt.Sprintf("%s:%s", host, buildAddress))
	swarming := b.GetInfra().GetSwarming()

	return &model.BuildSummary{
		ProjectID: b.Builder.Project,
		BuildKey:  buildKey,
		BuilderID: BuilderID{*b.Builder}.String(),
		BuildID:   "buildbucket/" + buildAddress,
		BuildSet:  b.Buildsets(),
		ContextURI: []string{
			fmt.Sprintf("buildbucket://%s/build/%d", host, id),
			fmt.Sprintf("swarming://%s/task/%s", swarming.GetHostname(), swarming.GetTaskId()),
		},
		Created: mustTimestamp(b.CreateTime),
		Summary: model.Summary{
			Start:  mustTimestamp(b.StartTime),
			End:    mustTimestamp(b.EndTime),
			Status: statusMap[b.Status],
		},
		Version:      mustTimestamp(b.UpdateTime).UnixNano(),
		Experimental: b.GetInput().GetExperimental(),
	}, nil
}

// generateSummary takes a decoded buildbucket event and generates
// a model.BuildSummary from it.
//
// This is the portion of the summarization process which cannot fail (i.e. is
// pure-data).
func generateSummary(c context.Context, hostname string, build buildbucket.Build) (*model.BuildSummary, error) {
	bs, err := getSummary(c, hostname, build.Project, build.ID)
	if err != nil {
		logging.WithError(err).Errorf(c, "got error while getting summary")
	}

	buildset := build.Tags[bbv1.TagBuildSet]
	if buildset == nil {
		buildset = []string{}
	}
	ret := &model.BuildSummary{
		ProjectID: build.Project,
		BuildKey:  MakeBuildKey(c, hostname, build.Address()),
		BuilderID: NewBuilderID(build.Bucket, build.Builder).String(),
		BuildID:   "buildbucket/" + build.Address(),
		BuildSet:  buildset,
		ContextURI: []string{
			fmt.Sprintf("buildbucket://%s/build/%d", hostname, build.ID),
		},

		Created: build.CreationTime,
		Summary: model.Summary{
			Start:  build.StartTime,
			End:    build.CompletionTime,
			Status: parseStatus(build.Status),
		},

		Version: build.UpdateTime.UnixNano(),

		Experimental: build.Experimental,
	}

	if shost, sid := build.Tags.Get("swarming_hostname"), build.Tags.Get("swarming_task_id"); shost != "" && sid != "" {
		ret.ContextURI = append(ret.ContextURI, fmt.Sprintf("swarming://%s/task/%s", shost, sid))
	}

	// Informational compare of v1 and v2 API result.
	if diff := cmp.Diff(ret, bs, cmpopts.IgnoreUnexported(model.BuildSummary{})); diff != "" {
		logging.Errorf(c, "BuildSummary of v2 has Diff (-v1, +v2)\n%s", diff)
	} else {
		logging.Debugf(c, "BuildSummary between v1 and v2 are the same.")
	}
	// TODO(iannucci,nodir): get the bot context too

	// TODO(iannucci,nodir): support manifests/got_revision
	return ret, ret.AddManifestKeysFromBuildSets(c)
}

// pubSubHandlerImpl takes the http.Request, expects to find
// a common.PubSubSubscription JSON object in the Body, containing a bbPSEvent,
// and handles the contents with generateSummary.
func pubSubHandlerImpl(c context.Context, r *http.Request) error {
	// This is the default action. The code below will modify the values of some
	// or all of these parameters.
	isLUCI, bucket, status, action := false, "UNKNOWN", "UNKNOWN", "Rejected"

	defer func() {
		// closure for late binding
		buildCounter.Add(c, 1, bucket, isLUCI, status, action)
	}()

	msg := common.PubSubSubscription{}
	if err := json.NewDecoder(r.Body).Decode(&msg); err != nil {
		// This might be a transient error, e.g. when the json format changes
		// and Milo isn't updated yet.
		return errors.Annotate(err, "could not decode message").Tag(transient.Tag).Err()
	}
	if v, ok := msg.Message.Attributes["version"].(string); ok && v != "v1" {
		// TODO(nodir): switch to v2, crbug.com/826006
		logging.Debugf(c, "unsupported pubsub message version %q. Ignoring", v)
		return nil
	}
	bData, err := msg.GetData()
	if err != nil {
		return errors.Annotate(err, "could not parse pubsub message string").Err()
	}

	event := struct {
		Build    bbv1.ApiCommonBuildMessage `json:"build"`
		Hostname string                     `json:"hostname"`
	}{}
	if err := json.Unmarshal(bData, &event); err != nil {
		return errors.Annotate(err, "could not parse pubsub message data").Err()
	}

	build := buildbucket.Build{}
	if err := build.ParseMessage(&event.Build); err != nil {
		return errors.Annotate(err, "could not parse buildbucket.Build").Err()
	}

	bucket = build.Bucket
	status = build.Status.String()
	isLUCI = strings.HasPrefix(bucket, "luci.")

	logging.Debugf(c, "Received from %s: build %s/%s (%s)\n%v",
		event.Hostname, bucket, build.Builder, status, build)

	if !isLUCI || build.Builder == "" {
		logging.Infof(c, "This is not an ingestable build, ignoring")
		return nil
	}

	bs, err := generateSummary(c, event.Hostname, build)
	if err != nil {
		return err
	}

	return transient.Tag.Apply(datastore.RunInTransaction(c, func(c context.Context) error {
		curBS := &model.BuildSummary{BuildKey: bs.BuildKey}
		switch err := datastore.Get(c, curBS); err {
		case datastore.ErrNoSuchEntity:
			action = "Created"
		case nil:
			action = "Modified"
		default:
			return errors.Annotate(err, "reading current BuildSummary").Err()
		}

		if build.UpdateTime.UnixNano() <= curBS.Version {
			logging.Warningf(c, "current BuildSummary is newer: %d <= %d",
				build.UpdateTime.UnixNano(), curBS.Version)
			return nil
		}

		if err := datastore.Put(c, bs); err != nil {
			return err
		}

		return model.UpdateBuilderForBuild(c, bs)
	}, &datastore.TransactionOptions{XG: true}))
}

// MakeBuildKey returns a new datastore Key for a buildbucket.Build.
//
// There's currently no model associated with this key, but it's used as
// a parent for a model.BuildSummary.
func MakeBuildKey(c context.Context, host, buildAddress string) *datastore.Key {
	return datastore.MakeKey(c,
		"buildbucket.Build", fmt.Sprintf("%s:%s", host, buildAddress))
}
<|endoftext|>