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

package main

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os/exec"
	"strings"

	"github.com/appc/spec/schema"
	"github.com/appc/spec/schema/types"
	"github.com/coreos/rocket/Godeps/_workspace/src/github.com/gorilla/mux"
)

type metadata struct {
	manifest schema.ContainerRuntimeManifest
	apps     map[string]*schema.ImageManifest
}

var (
	metadataByIP  = make(map[string]*metadata)
	metadataByUID = make(map[types.UUID]*metadata)
	hmacKey       [sha256.Size]byte
)

const (
	myPort   = "4444"
	metaIP   = "169.254.169.255"
	metaPort = "80"
)

func setupIPTables() error {
	return exec.Command(
		"iptables",
		"-t", "nat",
		"-A", "PREROUTING",
		"-p", "tcp",
		"-d", metaIP,
		"--dport", metaPort,
		"-j", "REDIRECT",
		"--to-port", myPort,
	).Run()
}

func antiSpoof(brPort, ipAddr string) error {
	return exec.Command(
		"ebtables",
		"-t", "filter",
		"-I", "INPUT",
		"-i", brPort,
		"-p", "IPV4",
		"!", "--ip-source", ipAddr,
		"-j", "DROP",
	).Run()
}

func queryValue(u *url.URL, key string) string {
	vals, ok := u.Query()[key]
	if !ok || len(vals) != 1 {
		return ""
	}
	return vals[0]
}

func handleRegisterContainer(w http.ResponseWriter, r *http.Request) {
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	if _, ok := metadataByIP[remoteIP]; ok {
		// not allowed from container IP
		w.WriteHeader(http.StatusForbidden)
		return
	}

	containerIP := queryValue(r.URL, "container_ip")
	if containerIP == "" {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Print(w, "container_ip missing")
		return
	}
	containerBrPort := queryValue(r.URL, "container_brport")
	if containerBrPort == "" {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Print(w, "container_brport missing")
		return
	}

	m := &metadata{
		apps: make(map[string]*schema.ImageManifest),
	}

	if err := json.NewDecoder(r.Body).Decode(&m.manifest); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "JSON-decoding failed: %v", err)
		return
	}

	if err := antiSpoof(containerBrPort, containerIP); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "failed to set anti-spoofing: %v", err)
		return
	}

	metadataByIP[containerIP] = m
	metadataByUID[m.manifest.UUID] = m

	w.WriteHeader(http.StatusOK)
}

func handleRegisterApp(w http.ResponseWriter, r *http.Request) {
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	if _, ok := metadataByIP[remoteIP]; ok {
		// not allowed from container IP
		w.WriteHeader(http.StatusForbidden)
		return
	}

	uid, err := types.NewUUID(mux.Vars(r)["uid"])
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "UUID is missing or mulformed: %v", err)
		return
	}

	m, ok := metadataByUID[*uid]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprint(w, "Container with given UUID not found")
		return
	}

	an := mux.Vars(r)["app"]

	app := &schema.ImageManifest{}
	if err := json.NewDecoder(r.Body).Decode(&app); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "JSON-decoding failed: %v", err)
		return
	}

	m.apps[an] = app

	w.WriteHeader(http.StatusOK)
}

func containerGet(h func(w http.ResponseWriter, r *http.Request, m *metadata)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		remoteIP := strings.Split(r.RemoteAddr, ":")[0]
		m, ok := metadataByIP[remoteIP]
		if !ok {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "metadata by remoteIP (%v) not found", remoteIP)
			return
		}

		h(w, r, m)
	}
}

func appGet(h func(w http.ResponseWriter, r *http.Request, m *metadata, _ *schema.ImageManifest)) func(http.ResponseWriter, *http.Request) {
	return containerGet(func(w http.ResponseWriter, r *http.Request, m *metadata) {
		appname := mux.Vars(r)["app"]

		if im, ok := m.apps[appname]; ok {
			h(w, r, m, im)
		} else {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "App (%v) not found", appname)
		}
	})
}

func handleContainerAnnotations(w http.ResponseWriter, r *http.Request, m *metadata) {
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)

	for k, _ := range m.manifest.Annotations {
		fmt.Fprintln(w, k)
	}
}

func handleContainerAnnotation(w http.ResponseWriter, r *http.Request, m *metadata) {
	k, err := types.NewACName(mux.Vars(r)["name"])
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Container annotation is not a valid AC Name")
		return
	}

	v, ok := m.manifest.Annotations.Get(k.String())
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Container annotation (%v) not found", k)
		return
	}

	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(v))
}

func handleContainerManifest(w http.ResponseWriter, r *http.Request, m *metadata) {
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)

	if err := json.NewEncoder(w).Encode(m.manifest); err != nil {
		fmt.Println(err)
	}
}

func handleContainerUID(w http.ResponseWriter, r *http.Request, m *metadata) {
	uid := m.manifest.UUID.String()

	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(uid))
}

func mergeAppAnnotations(im *schema.ImageManifest, cm *schema.ContainerRuntimeManifest) types.Annotations {
	merged := types.Annotations{}

	for _, annot := range im.Annotations {
		merged.Set(annot.Name, annot.Value)
	}

	if app := cm.Apps.Get(im.Name); app != nil {
		for _, annot := range app.Annotations {
			merged.Set(annot.Name, annot.Value)
		}
	}

	return merged
}

func handleAppAnnotations(w http.ResponseWriter, r *http.Request, m *metadata, im *schema.ImageManifest) {
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)

	for _, annot := range mergeAppAnnotations(im, &m.manifest) {
		fmt.Fprintln(w, string(annot.Name))
	}
}

func handleAppAnnotation(w http.ResponseWriter, r *http.Request, m *metadata, im *schema.ImageManifest) {
	k, err := types.NewACName(mux.Vars(r)["name"])
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "App annotation is not a valid AC Name")
		return
	}

	merged := mergeAppAnnotations(im, &m.manifest)

	v, ok := merged.Get(k.String())
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "App annotation (%v) not found", k)
		return
	}

	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(v))
}

func handleImageManifest(w http.ResponseWriter, r *http.Request, m *metadata, im *schema.ImageManifest) {
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)

	if err := json.NewEncoder(w).Encode(*im); err != nil {
		fmt.Println(err)
	}
}

func handleAppID(w http.ResponseWriter, r *http.Request, m *metadata, im *schema.ImageManifest) {
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	a := m.manifest.Apps.Get(im.Name)
	if a == nil {
		panic("could not find app in manifest!")
	}
	w.Write([]byte(a.ImageID.String()))
}

func initCrypto() error {
	if n, err := rand.Reader.Read(hmacKey[:]); err != nil || n != len(hmacKey) {
		return fmt.Errorf("failed to generate HMAC Key")
	}
	return nil
}

func digest(r io.Reader) ([]byte, error) {
	digest := sha256.New()
	if _, err := io.Copy(digest, r); err != nil {
		return nil, err
	}
	return digest.Sum(nil), nil
}

func handleContainerSign(w http.ResponseWriter, r *http.Request) {
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	m, ok := metadataByIP[remoteIP]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Metadata by remoteIP (%v) not found", remoteIP)
		return
	}

	// compute message digest
	d, err := digest(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "Digest computation failed: %v", err)
		return
	}

	// HMAC(UID:digest)
	h := hmac.New(sha256.New, hmacKey[:])
	h.Write(m.manifest.UUID[:])
	h.Write(d)

	// Send back digest:HMAC as the signature
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	enc := base64.NewEncoder(base64.StdEncoding, w)
	enc.Write(d)
	enc.Write(h.Sum(nil))
	enc.Close()
}

func handleContainerVerify(w http.ResponseWriter, r *http.Request) {
	uid, err := types.NewUUID(r.FormValue("uid"))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "uid field missing or malformed: %v", err)
		return
	}

	sig, err := base64.StdEncoding.DecodeString(r.FormValue("signature"))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "signature field missing or corrupt: %v", err)
		return
	}

	digest := sig[:sha256.Size]
	sum := sig[sha256.Size:]

	h := hmac.New(sha256.New, hmacKey[:])
	h.Write(uid[:])
	h.Write(digest)

	if hmac.Equal(sum, h.Sum(nil)) {
		w.WriteHeader(http.StatusOK)
	} else {
		w.WriteHeader(http.StatusForbidden)
	}
}

type httpResp struct {
	writer http.ResponseWriter
	status int
}

func (r *httpResp) Header() http.Header {
	return r.writer.Header()
}

func (r *httpResp) Write(d []byte) (int, error) {
	return r.writer.Write(d)
}

func (r *httpResp) WriteHeader(status int) {
	r.status = status
	r.writer.WriteHeader(status)
}

func logReq(h func(w http.ResponseWriter, r *http.Request)) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		resp := &httpResp{w, 0}
		h(resp, r)
		fmt.Printf("%v %v - %v\n", r.Method, r.RequestURI, resp.status)
	}
}

func main() {
	if err := setupIPTables(); err != nil {
		log.Fatal(err)
	}

	if err := initCrypto(); err != nil {
		log.Fatal(err)
	}

	r := mux.NewRouter()
	r.HandleFunc("/containers/", logReq(handleRegisterContainer)).Methods("POST")
	r.HandleFunc("/containers/{uid}/{app:.*}", logReq(handleRegisterApp)).Methods("PUT")

	acRtr := r.Headers("Metadata-Flavor", "AppContainer header").
		PathPrefix("/acMetadata/v1").Subrouter()

	mr := acRtr.Methods("GET").Subrouter()

	mr.HandleFunc("/container/annotations/", logReq(containerGet(handleContainerAnnotations)))
	mr.HandleFunc("/container/annotations/{name}", logReq(containerGet(handleContainerAnnotation)))
	mr.HandleFunc("/container/manifest", logReq(containerGet(handleContainerManifest)))
	mr.HandleFunc("/container/uid", logReq(containerGet(handleContainerUID)))

	mr.HandleFunc("/apps/{app:.*}/annotations/", logReq(appGet(handleAppAnnotations)))
	mr.HandleFunc("/apps/{app:.*}/annotations/{name}", logReq(appGet(handleAppAnnotation)))
	mr.HandleFunc("/apps/{app:.*}/image/manifest", logReq(appGet(handleImageManifest)))
	mr.HandleFunc("/apps/{app:.*}/image/id", logReq(appGet(handleAppID)))

	acRtr.HandleFunc("/container/hmac/sign", logReq(handleContainerSign)).Methods("POST")
	acRtr.HandleFunc("/container/hmac/verify", logReq(handleContainerVerify)).Methods("POST")

	log.Fatal(http.ListenAndServe(":4444", r))
}
<commit_msg>metadatasvc: leverage http.HandlerFunc type in func returns<commit_after>// Copyright 2014 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os/exec"
	"strings"

	"github.com/appc/spec/schema"
	"github.com/appc/spec/schema/types"
	"github.com/coreos/rocket/Godeps/_workspace/src/github.com/gorilla/mux"
)

type metadata struct {
	manifest schema.ContainerRuntimeManifest
	apps     map[string]*schema.ImageManifest
}

var (
	metadataByIP  = make(map[string]*metadata)
	metadataByUID = make(map[types.UUID]*metadata)
	hmacKey       [sha256.Size]byte
)

const (
	myPort   = "4444"
	metaIP   = "169.254.169.255"
	metaPort = "80"
)

func setupIPTables() error {
	return exec.Command(
		"iptables",
		"-t", "nat",
		"-A", "PREROUTING",
		"-p", "tcp",
		"-d", metaIP,
		"--dport", metaPort,
		"-j", "REDIRECT",
		"--to-port", myPort,
	).Run()
}

func antiSpoof(brPort, ipAddr string) error {
	return exec.Command(
		"ebtables",
		"-t", "filter",
		"-I", "INPUT",
		"-i", brPort,
		"-p", "IPV4",
		"!", "--ip-source", ipAddr,
		"-j", "DROP",
	).Run()
}

func queryValue(u *url.URL, key string) string {
	vals, ok := u.Query()[key]
	if !ok || len(vals) != 1 {
		return ""
	}
	return vals[0]
}

func handleRegisterContainer(w http.ResponseWriter, r *http.Request) {
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	if _, ok := metadataByIP[remoteIP]; ok {
		// not allowed from container IP
		w.WriteHeader(http.StatusForbidden)
		return
	}

	containerIP := queryValue(r.URL, "container_ip")
	if containerIP == "" {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Print(w, "container_ip missing")
		return
	}
	containerBrPort := queryValue(r.URL, "container_brport")
	if containerBrPort == "" {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Print(w, "container_brport missing")
		return
	}

	m := &metadata{
		apps: make(map[string]*schema.ImageManifest),
	}

	if err := json.NewDecoder(r.Body).Decode(&m.manifest); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "JSON-decoding failed: %v", err)
		return
	}

	if err := antiSpoof(containerBrPort, containerIP); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "failed to set anti-spoofing: %v", err)
		return
	}

	metadataByIP[containerIP] = m
	metadataByUID[m.manifest.UUID] = m

	w.WriteHeader(http.StatusOK)
}

func handleRegisterApp(w http.ResponseWriter, r *http.Request) {
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	if _, ok := metadataByIP[remoteIP]; ok {
		// not allowed from container IP
		w.WriteHeader(http.StatusForbidden)
		return
	}

	uid, err := types.NewUUID(mux.Vars(r)["uid"])
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "UUID is missing or mulformed: %v", err)
		return
	}

	m, ok := metadataByUID[*uid]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprint(w, "Container with given UUID not found")
		return
	}

	an := mux.Vars(r)["app"]

	app := &schema.ImageManifest{}
	if err := json.NewDecoder(r.Body).Decode(&app); err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "JSON-decoding failed: %v", err)
		return
	}

	m.apps[an] = app

	w.WriteHeader(http.StatusOK)
}

func containerGet(h func(w http.ResponseWriter, r *http.Request, m *metadata)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		remoteIP := strings.Split(r.RemoteAddr, ":")[0]
		m, ok := metadataByIP[remoteIP]
		if !ok {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "metadata by remoteIP (%v) not found", remoteIP)
			return
		}

		h(w, r, m)
	}
}

func appGet(h func(w http.ResponseWriter, r *http.Request, m *metadata, _ *schema.ImageManifest)) http.HandlerFunc {
	return containerGet(func(w http.ResponseWriter, r *http.Request, m *metadata) {
		appname := mux.Vars(r)["app"]

		if im, ok := m.apps[appname]; ok {
			h(w, r, m, im)
		} else {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "App (%v) not found", appname)
		}
	})
}

func handleContainerAnnotations(w http.ResponseWriter, r *http.Request, m *metadata) {
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)

	for k, _ := range m.manifest.Annotations {
		fmt.Fprintln(w, k)
	}
}

func handleContainerAnnotation(w http.ResponseWriter, r *http.Request, m *metadata) {
	k, err := types.NewACName(mux.Vars(r)["name"])
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Container annotation is not a valid AC Name")
		return
	}

	v, ok := m.manifest.Annotations.Get(k.String())
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Container annotation (%v) not found", k)
		return
	}

	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(v))
}

func handleContainerManifest(w http.ResponseWriter, r *http.Request, m *metadata) {
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)

	if err := json.NewEncoder(w).Encode(m.manifest); err != nil {
		fmt.Println(err)
	}
}

func handleContainerUID(w http.ResponseWriter, r *http.Request, m *metadata) {
	uid := m.manifest.UUID.String()

	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(uid))
}

func mergeAppAnnotations(im *schema.ImageManifest, cm *schema.ContainerRuntimeManifest) types.Annotations {
	merged := types.Annotations{}

	for _, annot := range im.Annotations {
		merged.Set(annot.Name, annot.Value)
	}

	if app := cm.Apps.Get(im.Name); app != nil {
		for _, annot := range app.Annotations {
			merged.Set(annot.Name, annot.Value)
		}
	}

	return merged
}

func handleAppAnnotations(w http.ResponseWriter, r *http.Request, m *metadata, im *schema.ImageManifest) {
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)

	for _, annot := range mergeAppAnnotations(im, &m.manifest) {
		fmt.Fprintln(w, string(annot.Name))
	}
}

func handleAppAnnotation(w http.ResponseWriter, r *http.Request, m *metadata, im *schema.ImageManifest) {
	k, err := types.NewACName(mux.Vars(r)["name"])
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "App annotation is not a valid AC Name")
		return
	}

	merged := mergeAppAnnotations(im, &m.manifest)

	v, ok := merged.Get(k.String())
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "App annotation (%v) not found", k)
		return
	}

	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(v))
}

func handleImageManifest(w http.ResponseWriter, r *http.Request, m *metadata, im *schema.ImageManifest) {
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)

	if err := json.NewEncoder(w).Encode(*im); err != nil {
		fmt.Println(err)
	}
}

func handleAppID(w http.ResponseWriter, r *http.Request, m *metadata, im *schema.ImageManifest) {
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	a := m.manifest.Apps.Get(im.Name)
	if a == nil {
		panic("could not find app in manifest!")
	}
	w.Write([]byte(a.ImageID.String()))
}

func initCrypto() error {
	if n, err := rand.Reader.Read(hmacKey[:]); err != nil || n != len(hmacKey) {
		return fmt.Errorf("failed to generate HMAC Key")
	}
	return nil
}

func digest(r io.Reader) ([]byte, error) {
	digest := sha256.New()
	if _, err := io.Copy(digest, r); err != nil {
		return nil, err
	}
	return digest.Sum(nil), nil
}

func handleContainerSign(w http.ResponseWriter, r *http.Request) {
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	m, ok := metadataByIP[remoteIP]
	if !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Metadata by remoteIP (%v) not found", remoteIP)
		return
	}

	// compute message digest
	d, err := digest(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "Digest computation failed: %v", err)
		return
	}

	// HMAC(UID:digest)
	h := hmac.New(sha256.New, hmacKey[:])
	h.Write(m.manifest.UUID[:])
	h.Write(d)

	// Send back digest:HMAC as the signature
	w.Header().Add("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	enc := base64.NewEncoder(base64.StdEncoding, w)
	enc.Write(d)
	enc.Write(h.Sum(nil))
	enc.Close()
}

func handleContainerVerify(w http.ResponseWriter, r *http.Request) {
	uid, err := types.NewUUID(r.FormValue("uid"))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "uid field missing or malformed: %v", err)
		return
	}

	sig, err := base64.StdEncoding.DecodeString(r.FormValue("signature"))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "signature field missing or corrupt: %v", err)
		return
	}

	digest := sig[:sha256.Size]
	sum := sig[sha256.Size:]

	h := hmac.New(sha256.New, hmacKey[:])
	h.Write(uid[:])
	h.Write(digest)

	if hmac.Equal(sum, h.Sum(nil)) {
		w.WriteHeader(http.StatusOK)
	} else {
		w.WriteHeader(http.StatusForbidden)
	}
}

type httpResp struct {
	writer http.ResponseWriter
	status int
}

func (r *httpResp) Header() http.Header {
	return r.writer.Header()
}

func (r *httpResp) Write(d []byte) (int, error) {
	return r.writer.Write(d)
}

func (r *httpResp) WriteHeader(status int) {
	r.status = status
	r.writer.WriteHeader(status)
}

func logReq(h func(w http.ResponseWriter, r *http.Request)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		resp := &httpResp{w, 0}
		h(resp, r)
		fmt.Printf("%v %v - %v\n", r.Method, r.RequestURI, resp.status)
	}
}

func main() {
	if err := setupIPTables(); err != nil {
		log.Fatal(err)
	}

	if err := initCrypto(); err != nil {
		log.Fatal(err)
	}

	r := mux.NewRouter()
	r.HandleFunc("/containers/", logReq(handleRegisterContainer)).Methods("POST")
	r.HandleFunc("/containers/{uid}/{app:.*}", logReq(handleRegisterApp)).Methods("PUT")

	acRtr := r.Headers("Metadata-Flavor", "AppContainer header").
		PathPrefix("/acMetadata/v1").Subrouter()

	mr := acRtr.Methods("GET").Subrouter()

	mr.HandleFunc("/container/annotations/", logReq(containerGet(handleContainerAnnotations)))
	mr.HandleFunc("/container/annotations/{name}", logReq(containerGet(handleContainerAnnotation)))
	mr.HandleFunc("/container/manifest", logReq(containerGet(handleContainerManifest)))
	mr.HandleFunc("/container/uid", logReq(containerGet(handleContainerUID)))

	mr.HandleFunc("/apps/{app:.*}/annotations/", logReq(appGet(handleAppAnnotations)))
	mr.HandleFunc("/apps/{app:.*}/annotations/{name}", logReq(appGet(handleAppAnnotation)))
	mr.HandleFunc("/apps/{app:.*}/image/manifest", logReq(appGet(handleImageManifest)))
	mr.HandleFunc("/apps/{app:.*}/image/id", logReq(appGet(handleAppID)))

	acRtr.HandleFunc("/container/hmac/sign", logReq(handleContainerSign)).Methods("POST")
	acRtr.HandleFunc("/container/hmac/verify", logReq(handleContainerVerify)).Methods("POST")

	log.Fatal(http.ListenAndServe(":4444", r))
}
<|endoftext|>


<commit_before>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You 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 go_kafka_client

import (
	"encoding/json"
	"fmt"
	kafkaavro "github.com/elodina/go-kafka-avro"
	"github.com/stealthly/go-avro"
	avroline "github.com/stealthly/go_kafka_client/avro"
	"regexp"
	"strings"
)

type CodahaleKafkaReporter struct {
	topic    string
	producer Producer
}

func NewCodahaleKafkaReporter(topic string, schemaRegistryUrl string, producerConfig *ProducerConfig) *CodahaleKafkaReporter {
	producerConfig.ValueEncoder = kafkaavro.NewKafkaAvroEncoder(schemaRegistryUrl)

	return &CodahaleKafkaReporter{
		topic:    topic,
		producer: NewSaramaProducer(producerConfig),
	}
}

func (c *CodahaleKafkaReporter) Write(bytes []byte) (n int, err error) {
	metrics := make(map[string]interface{})
	if err := json.Unmarshal(bytes, &metrics); err != nil {
		return -1, err
	}

	schema, lookupNames := c.parseSchema(metrics)
	record := avro.NewGenericRecord(schema)
	c.fillRecord(record, schema.(*avro.RecordSchema), metrics, lookupNames)

	c.producer.Input() <- &ProducerMessage{Topic: c.topic, Value: record}

	return 0, nil
}

func (c *CodahaleKafkaReporter) parseSchema(metrics map[string]interface{}) (avro.Schema, map[string]string) {
	lookupNames := make(map[string]string)
	schema := &avro.RecordSchema{
		Name:      "Metrics",
		Namespace: "ly.stealth",
		Fields:    make([]*avro.SchemaField, 0),
	}
	for name, v := range metrics {
		lookupName := c.getLookupName(name)
		lookupNames[lookupName] = name

		field := &avro.SchemaField{
			Name:    lookupName,
			Default: nil,
			Type: &avro.UnionSchema{
				Types: []avro.Schema{&avro.NullSchema{}, c.parseRecordField(name, v.(map[string]interface{}), lookupNames)},
			},
		}
		schema.Fields = append(schema.Fields, field)
	}
	schema.Fields = append(schema.Fields, &avro.SchemaField{
		Name: "logLine",
		Type: avroline.NewLogLine().Schema(),
	})

	return schema, lookupNames
}

func (c *CodahaleKafkaReporter) getLookupName(name string) string {
	lookupName := name
	pattern := regexp.MustCompile("[0-9]+")
	if pattern.MatchString(name) {
		lookupName = fmt.Sprintf("metric%s", name)
	}
	lookupName = strings.Replace(lookupName, ".", "", -1)
	lookupName = strings.Replace(lookupName, "%", "", -1)
	lookupName = strings.Replace(lookupName, "-", "", -1)

	return lookupName
}

func (c *CodahaleKafkaReporter) parseRecordField(name string, v map[string]interface{}, lookupNames map[string]string) *avro.RecordSchema {
	lookupRecordName := c.getLookupName(name)
	lookupNames[lookupRecordName] = name
	record := &avro.RecordSchema{
		Name:      lookupRecordName,
		Namespace: "ly.stealth",
		Fields:    make([]*avro.SchemaField, 0),
	}

	for name, value := range v {
		lookupName := c.getLookupName(name)
		lookupNames[lookupName] = name

		var valueSchema avro.Schema
		switch value.(type) {
		case string:
			valueSchema = &avro.StringSchema{}
		case float64:
			valueSchema = &avro.DoubleSchema{}
		}

		field := &avro.SchemaField{
			Name:    lookupName,
			Default: nil,
			Type: &avro.UnionSchema{
				Types: []avro.Schema{&avro.NullSchema{}, valueSchema},
			},
		}
		record.Fields = append(record.Fields, field)
	}

	return record
}

func (c *CodahaleKafkaReporter) fillRecord(record *avro.GenericRecord, schema *avro.RecordSchema, metrics map[string]interface{}, lookupNames map[string]string) {
	for _, field := range schema.Fields {
		if field.Name == "logLine" {
			//TODO probably something more?
			logLine := avro.NewGenericRecord(field.Type)
			logLine.Set("source", "metrics")
			logLine.Set("logtypeid", MetricsLogTypeId)
			record.Set("logLine", logLine)
		} else {
			name := lookupNames[field.Name]
			switch fieldSchema := field.Type.(*avro.UnionSchema).Types[1].(type) {
			case *avro.RecordSchema:
				{
					fieldRecord := avro.NewGenericRecord(fieldSchema)
					c.fillRecord(fieldRecord, fieldSchema, metrics[name].(map[string]interface{}), lookupNames)
					record.Set(field.Name, fieldRecord)
				}
			default:
				{
					record.Set(field.Name, metrics[field.Name])
				}
			}
		}
	}
}

type KafkaMetricReporter struct {
	topic    string
	producer Producer
}

func NewKafkaMetricReporter(topic string, producerConfig *ProducerConfig) *KafkaMetricReporter {
	return &KafkaMetricReporter{
		topic:    topic,
		producer: NewSaramaProducer(producerConfig),
	}
}

func (k *KafkaMetricReporter) Write(bytes []byte) (n int, err error) {
	k.producer.Input() <- &ProducerMessage{Topic: k.topic, Value: bytes}

	return 0, nil
}
<commit_msg>Fix metric avro package<commit_after>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You 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 go_kafka_client

import (
	"encoding/json"
	"fmt"
	kafkaavro "github.com/elodina/go-kafka-avro"
	"github.com/elodina/go-avro"
	avroline "github.com/stealthly/go_kafka_client/avro"
	"regexp"
	"strings"
)

type CodahaleKafkaReporter struct {
	topic    string
	producer Producer
}

func NewCodahaleKafkaReporter(topic string, schemaRegistryUrl string, producerConfig *ProducerConfig) *CodahaleKafkaReporter {
	producerConfig.ValueEncoder = kafkaavro.NewKafkaAvroEncoder(schemaRegistryUrl)

	return &CodahaleKafkaReporter{
		topic:    topic,
		producer: NewSaramaProducer(producerConfig),
	}
}

func (c *CodahaleKafkaReporter) Write(bytes []byte) (n int, err error) {
	metrics := make(map[string]interface{})
	if err := json.Unmarshal(bytes, &metrics); err != nil {
		return -1, err
	}

	schema, lookupNames := c.parseSchema(metrics)
	record := avro.NewGenericRecord(schema)
	c.fillRecord(record, schema.(*avro.RecordSchema), metrics, lookupNames)

	c.producer.Input() <- &ProducerMessage{Topic: c.topic, Value: record}

	return 0, nil
}

func (c *CodahaleKafkaReporter) parseSchema(metrics map[string]interface{}) (avro.Schema, map[string]string) {
	lookupNames := make(map[string]string)
	schema := &avro.RecordSchema{
		Name:      "Metrics",
		Namespace: "ly.stealth",
		Fields:    make([]*avro.SchemaField, 0),
	}
	for name, v := range metrics {
		lookupName := c.getLookupName(name)
		lookupNames[lookupName] = name

		field := &avro.SchemaField{
			Name:    lookupName,
			Default: nil,
			Type: &avro.UnionSchema{
				Types: []avro.Schema{&avro.NullSchema{}, c.parseRecordField(name, v.(map[string]interface{}), lookupNames)},
			},
		}
		schema.Fields = append(schema.Fields, field)
	}
	schema.Fields = append(schema.Fields, &avro.SchemaField{
		Name: "logLine",
		Type: avroline.NewLogLine().Schema(),
	})

	return schema, lookupNames
}

func (c *CodahaleKafkaReporter) getLookupName(name string) string {
	lookupName := name
	pattern := regexp.MustCompile("[0-9]+")
	if pattern.MatchString(name) {
		lookupName = fmt.Sprintf("metric%s", name)
	}
	lookupName = strings.Replace(lookupName, ".", "", -1)
	lookupName = strings.Replace(lookupName, "%", "", -1)
	lookupName = strings.Replace(lookupName, "-", "", -1)

	return lookupName
}

func (c *CodahaleKafkaReporter) parseRecordField(name string, v map[string]interface{}, lookupNames map[string]string) *avro.RecordSchema {
	lookupRecordName := c.getLookupName(name)
	lookupNames[lookupRecordName] = name
	record := &avro.RecordSchema{
		Name:      lookupRecordName,
		Namespace: "ly.stealth",
		Fields:    make([]*avro.SchemaField, 0),
	}

	for name, value := range v {
		lookupName := c.getLookupName(name)
		lookupNames[lookupName] = name

		var valueSchema avro.Schema
		switch value.(type) {
		case string:
			valueSchema = &avro.StringSchema{}
		case float64:
			valueSchema = &avro.DoubleSchema{}
		}

		field := &avro.SchemaField{
			Name:    lookupName,
			Default: nil,
			Type: &avro.UnionSchema{
				Types: []avro.Schema{&avro.NullSchema{}, valueSchema},
			},
		}
		record.Fields = append(record.Fields, field)
	}

	return record
}

func (c *CodahaleKafkaReporter) fillRecord(record *avro.GenericRecord, schema *avro.RecordSchema, metrics map[string]interface{}, lookupNames map[string]string) {
	for _, field := range schema.Fields {
		if field.Name == "logLine" {
			//TODO probably something more?
			logLine := avro.NewGenericRecord(field.Type)
			logLine.Set("source", "metrics")
			logLine.Set("logtypeid", MetricsLogTypeId)
			record.Set("logLine", logLine)
		} else {
			name := lookupNames[field.Name]
			switch fieldSchema := field.Type.(*avro.UnionSchema).Types[1].(type) {
			case *avro.RecordSchema:
				{
					fieldRecord := avro.NewGenericRecord(fieldSchema)
					c.fillRecord(fieldRecord, fieldSchema, metrics[name].(map[string]interface{}), lookupNames)
					record.Set(field.Name, fieldRecord)
				}
			default:
				{
					record.Set(field.Name, metrics[field.Name])
				}
			}
		}
	}
}

type KafkaMetricReporter struct {
	topic    string
	producer Producer
}

func NewKafkaMetricReporter(topic string, producerConfig *ProducerConfig) *KafkaMetricReporter {
	return &KafkaMetricReporter{
		topic:    topic,
		producer: NewSaramaProducer(producerConfig),
	}
}

func (k *KafkaMetricReporter) Write(bytes []byte) (n int, err error) {
	k.producer.Input() <- &ProducerMessage{Topic: k.topic, Value: bytes}

	return 0, nil
}
<|endoftext|>


<commit_before>package sched

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"html/template"
	"io"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/StackExchange/bosun/_third_party/github.com/StackExchange/scollector/metadata"
	"github.com/StackExchange/bosun/_third_party/github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/bosun/conf"
	"github.com/StackExchange/bosun/expr"
	"github.com/StackExchange/bosun/expr/parse"
)

type Context struct {
	*State
	Alert *conf.Alert

	schedule    *Schedule
	Attachments []*conf.Attachment
}

func (s *Schedule) Data(st *State, a *conf.Alert, isEmail bool) *Context {
	c := Context{
		State:    st,
		Alert:    a,
		schedule: s,
	}
	if isEmail {
		c.Attachments = make([]*conf.Attachment, 0)
	}
	return &c
}

type unknownContext struct {
	Time  time.Time
	Name  string
	Group expr.AlertKeys

	schedule *Schedule
}

func (s *Schedule) unknownData(t time.Time, name string, group expr.AlertKeys) *unknownContext {
	return &unknownContext{
		Time:     t,
		Group:    group,
		Name:     name,
		schedule: s,
	}
}

// URL returns a prepopulated URL for external access, with path and query empty.
func (s *Schedule) URL() *url.URL {
	u := url.URL{
		Scheme: "http",
		Host:   s.Conf.HttpListen,
	}
	if strings.HasPrefix(s.Conf.HttpListen, ":") {
		h, err := os.Hostname()
		if err != nil {
			u.Host = "localhost" + u.Host
		} else {
			u.Host = h + u.Host
		}
	}
	return &u
}

// Ack returns the URL to acknowledge an alert.
func (c *Context) Ack() string {
	u := c.schedule.URL()
	u.Path = "/action"
	u.RawQuery = url.Values{
		"type": []string{"ack"},
		"key":  []string{c.Alert.Name + c.State.Group.String()},
	}.Encode()
	return u.String()
}

// HostView returns the URL to the host view page.
func (c *Context) HostView(host string) string {
	u := c.schedule.URL()
	u.Path = "/host"
	u.RawQuery = fmt.Sprintf("time=1d-ago&host=%s", host)
	return u.String()
}

func (c *Context) Expr(v string) (string, error) {
	q := "expr=" + base64.URLEncoding.EncodeToString([]byte(opentsdb.ReplaceTags(v, c.Group)))
	u := url.URL{
		Scheme:   "http",
		Host:     c.schedule.Conf.HttpListen,
		Path:     "/expr",
		RawQuery: q,
	}
	if strings.HasPrefix(c.schedule.Conf.HttpListen, ":") {
		h, err := os.Hostname()
		if err != nil {
			return "", err
		}
		u.Host = h + u.Host
	}
	return u.String(), nil
}

func (s *Schedule) ExecuteBody(w io.Writer, a *conf.Alert, st *State, isEmail bool) ([]*conf.Attachment, error) {
	t := a.Template
	if t == nil || t.Body == nil {
		return nil, nil
	}
	c := s.Data(st, a, isEmail)
	return c.Attachments, t.Body.Execute(w, c)
}

func (s *Schedule) ExecuteSubject(w io.Writer, a *conf.Alert, st *State) error {
	t := a.Template
	if t == nil || t.Subject == nil {
		return nil
	}
	return t.Subject.Execute(w, s.Data(st, a, false))
}

func (c *Context) eval(v interface{}, filter bool, series bool, autods int) ([]*expr.Result, error) {
	var e *expr.Expr
	var err error
	switch v := v.(type) {
	case string:
		e, err = expr.New(v)
	case *expr.Expr:
		e = v
	default:
		return nil, fmt.Errorf("expected string or expression, got %T (%v)", v, v)
	}
	if err != nil {
		return nil, fmt.Errorf("%v: %v", v, err)
	}
	if series && e.Root.Return() != parse.TYPE_SERIES {
		return nil, fmt.Errorf("egraph: requires an expression that returns a series")
	}
	res, _, err := e.Execute(c.schedule.cache, nil, c.schedule.CheckStart, autods, c.Alert.UnjoinedOK, c.schedule.Search, c.schedule.Lookups, c.schedule.Conf.AlertSquelched(c.Alert))
	if err != nil {
		return nil, fmt.Errorf("%s: %v", v, err)
	}
	if !filter {
		return res.Results, nil
	}
	for _, r := range res.Results {
		if r.Group.Equal(c.State.Group) {
			return []*expr.Result{r}, nil
		}
	}
	for _, r := range res.Results {
		if c.State.Group.Subset(r.Group) {
			return []*expr.Result{r}, nil
		}
		}
	}
	return nil, nil
}

// Lookup returns the value for a key in the lookup table for the context's tagset.
func (c *Context) Lookup(table, key string) (string, error) {
	l, ok := c.schedule.Lookups[table]
	if !ok {
		return "", fmt.Errorf("unknown lookup table %v", table)
	}
	if v, ok := l.Get(key, c.Group); ok {
		return v, nil
	} else {
		return "", fmt.Errorf("no entry for key %v in table %v for tagset %v", key, table, c.Group)
	}
}

// Eval executes the given expression and returns a value with corresponding
// tags to the context's tags. If no such result is found, the first result with
// nil tags is returned. If no such result is found, nil is returned.
func (c *Context) Eval(v interface{}) (interface{}, error) {
	res, err := c.eval(v, true, false, 0)
	if err != nil {
		return nil, err
	}
	if len(res) != 1 {
		return nil, fmt.Errorf("expected 1 result, got %v", len(res))
	}
	return res[0].Value, nil
}

// EvalAll returns the executed expression.
func (c *Context) EvalAll(v interface{}) (interface{}, error) {
	return c.eval(v, false, false, 0)
}

func (c *Context) IsEmail() bool {
	return c.Attachments != nil
}

func (c *Context) graph(v interface{}, filter bool) (interface{}, error) {
	res, err := c.eval(v, filter, true, 1000)
	if err != nil {
		return nil, err
	}
	var buf bytes.Buffer
	if err := c.schedule.ExprGraph(nil, &buf, res, fmt.Sprint(v), time.Now().UTC()); err != nil {
		return nil, err
	}
	if c.IsEmail() {
		name := fmt.Sprintf("%d.svg", len(c.Attachments)+1)
		c.Attachments = append(c.Attachments, &conf.Attachment{
			Data:        buf.Bytes(),
			Filename:    name,
			ContentType: "image/svg+xml",
		})
		return template.HTML(fmt.Sprintf(`<img alt="%s" src="cid:%s" />`,
			template.HTMLEscapeString(fmt.Sprint(v)),
			name,
		)), nil
	}
	return template.HTML(buf.String()), nil
}

func (c *Context) Graph(v interface{}) (interface{}, error) {
	return c.graph(v, true)
}

func (c *Context) GraphAll(v interface{}) (interface{}, error) {
	return c.graph(v, false)
}

func (c *Context) GetMeta(metric, name string, v interface{}) (interface{}, error) {
	var t opentsdb.TagSet
	switch v := v.(type) {
	case string:
		var err error
		t, err = opentsdb.ParseTags(v)
		if err != nil {
			return t, err
		}
	case opentsdb.TagSet:
		t = v
	}
	meta := c.schedule.GetMetadata(metric, t)
	if name == "" {
		return meta, nil
	}
	fm := make([]metadata.Metasend, 0)
	for _, m := range meta {
		if m.Name == name {
			fm = append(fm, m)
		}
	}
	return fm, nil
}

func (c *Context) LeftJoin(q ...interface{}) (interface{}, error) {
	if len(q) < 2 {
		return nil, fmt.Errorf("need at least two expressions, got %v", len(q))
	}
	matrix := make([][]*expr.Result, 0)
	results := make([][]*expr.Result, len(q))
	for col, v := range q {
		res, err := c.eval(v, false, false, 0)
		if err != nil {
			return nil, err
		}
		results[col] = res
	}
	for row, first := range results[0] {
		matrix = append(matrix, make([]*expr.Result, len(q)))
		matrix[row][0] = first
		for col, res := range results[1:] {
			for _, r := range res {
				if first.Group.Subset(r.Group) {
					matrix[row][col+1] = r
				}
			}
		}
	}
	return matrix, nil
}
<commit_msg>Also check for nil tags<commit_after>package sched

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"html/template"
	"io"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/StackExchange/bosun/_third_party/github.com/StackExchange/scollector/metadata"
	"github.com/StackExchange/bosun/_third_party/github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/bosun/conf"
	"github.com/StackExchange/bosun/expr"
	"github.com/StackExchange/bosun/expr/parse"
)

type Context struct {
	*State
	Alert *conf.Alert

	schedule    *Schedule
	Attachments []*conf.Attachment
}

func (s *Schedule) Data(st *State, a *conf.Alert, isEmail bool) *Context {
	c := Context{
		State:    st,
		Alert:    a,
		schedule: s,
	}
	if isEmail {
		c.Attachments = make([]*conf.Attachment, 0)
	}
	return &c
}

type unknownContext struct {
	Time  time.Time
	Name  string
	Group expr.AlertKeys

	schedule *Schedule
}

func (s *Schedule) unknownData(t time.Time, name string, group expr.AlertKeys) *unknownContext {
	return &unknownContext{
		Time:     t,
		Group:    group,
		Name:     name,
		schedule: s,
	}
}

// URL returns a prepopulated URL for external access, with path and query empty.
func (s *Schedule) URL() *url.URL {
	u := url.URL{
		Scheme: "http",
		Host:   s.Conf.HttpListen,
	}
	if strings.HasPrefix(s.Conf.HttpListen, ":") {
		h, err := os.Hostname()
		if err != nil {
			u.Host = "localhost" + u.Host
		} else {
			u.Host = h + u.Host
		}
	}
	return &u
}

// Ack returns the URL to acknowledge an alert.
func (c *Context) Ack() string {
	u := c.schedule.URL()
	u.Path = "/action"
	u.RawQuery = url.Values{
		"type": []string{"ack"},
		"key":  []string{c.Alert.Name + c.State.Group.String()},
	}.Encode()
	return u.String()
}

// HostView returns the URL to the host view page.
func (c *Context) HostView(host string) string {
	u := c.schedule.URL()
	u.Path = "/host"
	u.RawQuery = fmt.Sprintf("time=1d-ago&host=%s", host)
	return u.String()
}

func (c *Context) Expr(v string) (string, error) {
	q := "expr=" + base64.URLEncoding.EncodeToString([]byte(opentsdb.ReplaceTags(v, c.Group)))
	u := url.URL{
		Scheme:   "http",
		Host:     c.schedule.Conf.HttpListen,
		Path:     "/expr",
		RawQuery: q,
	}
	if strings.HasPrefix(c.schedule.Conf.HttpListen, ":") {
		h, err := os.Hostname()
		if err != nil {
			return "", err
		}
		u.Host = h + u.Host
	}
	return u.String(), nil
}

func (s *Schedule) ExecuteBody(w io.Writer, a *conf.Alert, st *State, isEmail bool) ([]*conf.Attachment, error) {
	t := a.Template
	if t == nil || t.Body == nil {
		return nil, nil
	}
	c := s.Data(st, a, isEmail)
	return c.Attachments, t.Body.Execute(w, c)
}

func (s *Schedule) ExecuteSubject(w io.Writer, a *conf.Alert, st *State) error {
	t := a.Template
	if t == nil || t.Subject == nil {
		return nil
	}
	return t.Subject.Execute(w, s.Data(st, a, false))
}

func (c *Context) eval(v interface{}, filter bool, series bool, autods int) ([]*expr.Result, error) {
	var e *expr.Expr
	var err error
	switch v := v.(type) {
	case string:
		e, err = expr.New(v)
	case *expr.Expr:
		e = v
	default:
		return nil, fmt.Errorf("expected string or expression, got %T (%v)", v, v)
	}
	if err != nil {
		return nil, fmt.Errorf("%v: %v", v, err)
	}
	if series && e.Root.Return() != parse.TYPE_SERIES {
		return nil, fmt.Errorf("egraph: requires an expression that returns a series")
	}
	res, _, err := e.Execute(c.schedule.cache, nil, c.schedule.CheckStart, autods, c.Alert.UnjoinedOK, c.schedule.Search, c.schedule.Lookups, c.schedule.Conf.AlertSquelched(c.Alert))
	if err != nil {
		return nil, fmt.Errorf("%s: %v", v, err)
	}
	if !filter {
		return res.Results, nil
	}
	for _, r := range res.Results {
		if r.Group.Equal(c.State.Group) {
			return []*expr.Result{r}, nil
		}
	}
	for _, r := range res.Results {
		if c.State.Group.Subset(r.Group) {
			return []*expr.Result{r}, nil
		}
	}
	for _, r := range res.Results {
		if r.Group == nil {
			return []*expr.Result{r}, nil
		}
	}
	return nil, nil
}

// Lookup returns the value for a key in the lookup table for the context's tagset.
func (c *Context) Lookup(table, key string) (string, error) {
	l, ok := c.schedule.Lookups[table]
	if !ok {
		return "", fmt.Errorf("unknown lookup table %v", table)
	}
	if v, ok := l.Get(key, c.Group); ok {
		return v, nil
	} else {
		return "", fmt.Errorf("no entry for key %v in table %v for tagset %v", key, table, c.Group)
	}
}

// Eval executes the given expression and returns a value with corresponding
// tags to the context's tags. If no such result is found, the first result with
// nil tags is returned. If no such result is found, nil is returned.
func (c *Context) Eval(v interface{}) (interface{}, error) {
	res, err := c.eval(v, true, false, 0)
	if err != nil {
		return nil, err
	}
	if len(res) != 1 {
		return nil, fmt.Errorf("expected 1 result, got %v", len(res))
	}
	return res[0].Value, nil
}

// EvalAll returns the executed expression.
func (c *Context) EvalAll(v interface{}) (interface{}, error) {
	return c.eval(v, false, false, 0)
}

func (c *Context) IsEmail() bool {
	return c.Attachments != nil
}

func (c *Context) graph(v interface{}, filter bool) (interface{}, error) {
	res, err := c.eval(v, filter, true, 1000)
	if err != nil {
		return nil, err
	}
	var buf bytes.Buffer
	if err := c.schedule.ExprGraph(nil, &buf, res, fmt.Sprint(v), time.Now().UTC()); err != nil {
		return nil, err
	}
	if c.IsEmail() {
		name := fmt.Sprintf("%d.svg", len(c.Attachments)+1)
		c.Attachments = append(c.Attachments, &conf.Attachment{
			Data:        buf.Bytes(),
			Filename:    name,
			ContentType: "image/svg+xml",
		})
		return template.HTML(fmt.Sprintf(`<img alt="%s" src="cid:%s" />`,
			template.HTMLEscapeString(fmt.Sprint(v)),
			name,
		)), nil
	}
	return template.HTML(buf.String()), nil
}

func (c *Context) Graph(v interface{}) (interface{}, error) {
	return c.graph(v, true)
}

func (c *Context) GraphAll(v interface{}) (interface{}, error) {
	return c.graph(v, false)
}

func (c *Context) GetMeta(metric, name string, v interface{}) (interface{}, error) {
	var t opentsdb.TagSet
	switch v := v.(type) {
	case string:
		var err error
		t, err = opentsdb.ParseTags(v)
		if err != nil {
			return t, err
		}
	case opentsdb.TagSet:
		t = v
	}
	meta := c.schedule.GetMetadata(metric, t)
	if name == "" {
		return meta, nil
	}
	fm := make([]metadata.Metasend, 0)
	for _, m := range meta {
		if m.Name == name {
			fm = append(fm, m)
		}
	}
	return fm, nil
}

func (c *Context) LeftJoin(q ...interface{}) (interface{}, error) {
	if len(q) < 2 {
		return nil, fmt.Errorf("need at least two expressions, got %v", len(q))
	}
	matrix := make([][]*expr.Result, 0)
	results := make([][]*expr.Result, len(q))
	for col, v := range q {
		res, err := c.eval(v, false, false, 0)
		if err != nil {
			return nil, err
		}
		results[col] = res
	}
	for row, first := range results[0] {
		matrix = append(matrix, make([]*expr.Result, len(q)))
		matrix[row][0] = first
		for col, res := range results[1:] {
			for _, r := range res {
				if first.Group.Subset(r.Group) {
					matrix[row][col+1] = r
				}
			}
		}
	}
	return matrix, nil
}
<|endoftext|>


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

package scrape

import (
	"encoding"
	"fmt"
	"hash/fnv"
	"net"
	"os"
	"reflect"
	"sync"
	"time"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/discovery/targetgroup"
	"github.com/prometheus/prometheus/pkg/labels"
	"github.com/prometheus/prometheus/storage"
)

var targetMetadataCache = newMetadataMetricsCollector()

// MetadataMetricsCollector is a Custom Collector for the metadata cache metrics.
type MetadataMetricsCollector struct {
	CacheEntries *prometheus.Desc
	CacheBytes   *prometheus.Desc

	scrapeManager *Manager
}

func newMetadataMetricsCollector() *MetadataMetricsCollector {
	return &MetadataMetricsCollector{
		CacheEntries: prometheus.NewDesc(
			"prometheus_target_metadata_cache_entries",
			"Total number of metric metadata entries in the cache",
			[]string{"scrape_job"},
			nil,
		),
		CacheBytes: prometheus.NewDesc(
			"prometheus_target_metadata_cache_bytes",
			"The number of bytes that are currently used for storing metric metadata in the cache",
			[]string{"scrape_job"},
			nil,
		),
	}
}

func (mc *MetadataMetricsCollector) registerManager(m *Manager) {
	mc.scrapeManager = m
}

// Describe sends the metrics descriptions to the channel.
func (mc *MetadataMetricsCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- mc.CacheEntries
	ch <- mc.CacheBytes
}

// Collect creates and sends the metrics for the metadata cache.
func (mc *MetadataMetricsCollector) Collect(ch chan<- prometheus.Metric) {
	if mc.scrapeManager == nil {
		return
	}

	for tset, targets := range mc.scrapeManager.TargetsActive() {
		var size, length int
		for _, t := range targets {
			size += t.MetadataSize()
			length += t.MetadataLength()
		}

		ch <- prometheus.MustNewConstMetric(
			mc.CacheEntries,
			prometheus.GaugeValue,
			float64(length),
			tset,
		)

		ch <- prometheus.MustNewConstMetric(
			mc.CacheBytes,
			prometheus.GaugeValue,
			float64(size),
			tset,
		)
	}
}

// NewManager is the Manager constructor
func NewManager(logger log.Logger, app storage.Appendable) *Manager {
	if logger == nil {
		logger = log.NewNopLogger()
	}
	m := &Manager{
		append:        app,
		logger:        logger,
		scrapeConfigs: make(map[string]*config.ScrapeConfig),
		scrapePools:   make(map[string]*scrapePool),
		graceShut:     make(chan struct{}),
		triggerReload: make(chan struct{}, 1),
	}
	targetMetadataCache.registerManager(m)

	return m
}

// Manager maintains a set of scrape pools and manages start/stop cycles
// when receiving new target groups form the discovery manager.
type Manager struct {
	logger    log.Logger
	append    storage.Appendable
	graceShut chan struct{}

	jitterSeed    uint64     // Global jitterSeed seed is used to spread scrape workload across HA setup.
	mtxScrape     sync.Mutex // Guards the fields below.
	scrapeConfigs map[string]*config.ScrapeConfig
	scrapePools   map[string]*scrapePool
	targetSets    map[string][]*targetgroup.Group

	triggerReload chan struct{}
}

// Run receives and saves target set updates and triggers the scraping loops reloading.
// Reloading happens in the background so that it doesn't block receiving targets updates.
func (m *Manager) Run(tsets <-chan map[string][]*targetgroup.Group) error {
	go m.reloader()
	for {
		select {
		case ts := <-tsets:
			m.updateTsets(ts)

			select {
			case m.triggerReload <- struct{}{}:
			default:
			}

		case <-m.graceShut:
			return nil
		}
	}
}

func (m *Manager) reloader() {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-m.graceShut:
			return
		case <-ticker.C:
			select {
			case <-m.triggerReload:
				m.reload()
			case <-m.graceShut:
				return
			}
		}
	}
}

func (m *Manager) reload() {
	m.mtxScrape.Lock()
	var wg sync.WaitGroup
	for setName, groups := range m.targetSets {
		if _, ok := m.scrapePools[setName]; !ok {
			scrapeConfig, ok := m.scrapeConfigs[setName]
			if !ok {
				level.Error(m.logger).Log("msg", "error reloading target set", "err", "invalid config id:"+setName)
				continue
			}
			sp, err := newScrapePool(scrapeConfig, m.append, m.jitterSeed, log.With(m.logger, "scrape_pool", setName))
			if err != nil {
				level.Error(m.logger).Log("msg", "error creating new scrape pool", "err", err, "scrape_pool", setName)
				continue
			}
			m.scrapePools[setName] = sp
		}

		wg.Add(1)
		// Run the sync in parallel as these take a while and at high load can't catch up.
		go func(sp *scrapePool, groups []*targetgroup.Group) {
			sp.Sync(groups)
			wg.Done()
		}(m.scrapePools[setName], groups)

	}
	m.mtxScrape.Unlock()
	wg.Wait()
}

// setJitterSeed calculates a global jitterSeed per server relying on extra label set.
func (m *Manager) setJitterSeed(labels labels.Labels) error {
	h := fnv.New64a()
	hostname, err := getFqdn()
	if err != nil {
		return err
	}
	if _, err := fmt.Fprintf(h, "%s%s", hostname, labels.String()); err != nil {
		return err
	}
	m.jitterSeed = h.Sum64()
	return nil
}

// Stop cancels all running scrape pools and blocks until all have exited.
func (m *Manager) Stop() {
	m.mtxScrape.Lock()
	defer m.mtxScrape.Unlock()

	for _, sp := range m.scrapePools {
		sp.stop()
	}
	close(m.graceShut)
}

func (m *Manager) updateTsets(tsets map[string][]*targetgroup.Group) {
	m.mtxScrape.Lock()
	m.targetSets = tsets
	m.mtxScrape.Unlock()
}

// ApplyConfig resets the manager's target providers and job configurations as defined by the new cfg.
func (m *Manager) ApplyConfig(cfg *config.Config) error {
	m.mtxScrape.Lock()
	defer m.mtxScrape.Unlock()

	c := make(map[string]*config.ScrapeConfig)
	for _, scfg := range cfg.ScrapeConfigs {
		c[scfg.JobName] = scfg
	}
	m.scrapeConfigs = c

	if err := m.setJitterSeed(cfg.GlobalConfig.ExternalLabels); err != nil {
		return err
	}

	// Cleanup and reload pool if the configuration has changed.
	var failed bool
	for name, sp := range m.scrapePools {
		if cfg, ok := m.scrapeConfigs[name]; !ok {
			sp.stop()
			delete(m.scrapePools, name)
		} else if !reflect.DeepEqual(sp.config, cfg) {
			err := sp.reload(cfg)
			if err != nil {
				level.Error(m.logger).Log("msg", "error reloading scrape pool", "err", err, "scrape_pool", name)
				failed = true
			}
		}
	}

	if failed {
		return errors.New("failed to apply the new configuration")
	}
	return nil
}

// TargetsAll returns active and dropped targets grouped by job_name.
func (m *Manager) TargetsAll() map[string][]*Target {
	m.mtxScrape.Lock()
	defer m.mtxScrape.Unlock()

	targets := make(map[string][]*Target, len(m.scrapePools))
	for tset, sp := range m.scrapePools {
		targets[tset] = append(sp.ActiveTargets(), sp.DroppedTargets()...)
	}
	return targets
}

// TargetsActive returns the active targets currently being scraped.
func (m *Manager) TargetsActive() map[string][]*Target {
	m.mtxScrape.Lock()
	defer m.mtxScrape.Unlock()

	var (
		wg  sync.WaitGroup
		mtx sync.Mutex
	)

	targets := make(map[string][]*Target, len(m.scrapePools))
	wg.Add(len(m.scrapePools))
	for tset, sp := range m.scrapePools {
		// Running in parallel limits the blocking time of scrapePool to scrape
		// interval when there's an update from SD.
		go func(tset string, sp *scrapePool) {
			mtx.Lock()
			targets[tset] = sp.ActiveTargets()
			mtx.Unlock()
			wg.Done()
		}(tset, sp)
	}
	wg.Wait()
	return targets
}

// TargetsDropped returns the dropped targets during relabelling.
func (m *Manager) TargetsDropped() map[string][]*Target {
	m.mtxScrape.Lock()
	defer m.mtxScrape.Unlock()

	targets := make(map[string][]*Target, len(m.scrapePools))
	for tset, sp := range m.scrapePools {
		targets[tset] = sp.DroppedTargets()
	}
	return targets
}

// getFqdn returns a FQDN if it's possible, otherwise falls back to hostname.
func getFqdn() (string, error) {
	hostname, err := os.Hostname()
	if err != nil {
		return "", err
	}

	ips, err := net.LookupIP(hostname)
	if err != nil {
		// Return the system hostname if we can't look up the IP address.
		return hostname, nil
	}

	lookup := func(ipStr encoding.TextMarshaler) (string, error) {
		ip, err := ipStr.MarshalText()
		if err != nil {
			return "", err
		}
		hosts, err := net.LookupAddr(string(ip))
		if err != nil || len(hosts) == 0 {
			return "", err
		}
		return hosts[0], nil
	}

	for _, addr := range ips {
		if ip := addr.To4(); ip != nil {
			if fqdn, err := lookup(ip); err == nil {
				return fqdn, nil
			}

		}

		if ip := addr.To16(); ip != nil {
			if fqdn, err := lookup(ip); err == nil {
				return fqdn, nil
			}

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

package scrape

import (
	"encoding"
	"fmt"
	"hash/fnv"
	"net"
	"os"
	"reflect"
	"sync"
	"time"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/discovery/targetgroup"
	"github.com/prometheus/prometheus/pkg/labels"
	"github.com/prometheus/prometheus/storage"
)

var targetMetadataCache = newMetadataMetricsCollector()

// MetadataMetricsCollector is a Custom Collector for the metadata cache metrics.
type MetadataMetricsCollector struct {
	CacheEntries *prometheus.Desc
	CacheBytes   *prometheus.Desc

	scrapeManager *Manager
}

func newMetadataMetricsCollector() *MetadataMetricsCollector {
	return &MetadataMetricsCollector{
		CacheEntries: prometheus.NewDesc(
			"prometheus_target_metadata_cache_entries",
			"Total number of metric metadata entries in the cache",
			[]string{"scrape_job"},
			nil,
		),
		CacheBytes: prometheus.NewDesc(
			"prometheus_target_metadata_cache_bytes",
			"The number of bytes that are currently used for storing metric metadata in the cache",
			[]string{"scrape_job"},
			nil,
		),
	}
}

func (mc *MetadataMetricsCollector) registerManager(m *Manager) {
	mc.scrapeManager = m
}

// Describe sends the metrics descriptions to the channel.
func (mc *MetadataMetricsCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- mc.CacheEntries
	ch <- mc.CacheBytes
}

// Collect creates and sends the metrics for the metadata cache.
func (mc *MetadataMetricsCollector) Collect(ch chan<- prometheus.Metric) {
	if mc.scrapeManager == nil {
		return
	}

	for tset, targets := range mc.scrapeManager.TargetsActive() {
		var size, length int
		for _, t := range targets {
			size += t.MetadataSize()
			length += t.MetadataLength()
		}

		ch <- prometheus.MustNewConstMetric(
			mc.CacheEntries,
			prometheus.GaugeValue,
			float64(length),
			tset,
		)

		ch <- prometheus.MustNewConstMetric(
			mc.CacheBytes,
			prometheus.GaugeValue,
			float64(size),
			tset,
		)
	}
}

// NewManager is the Manager constructor
func NewManager(logger log.Logger, app storage.Appendable) *Manager {
	if logger == nil {
		logger = log.NewNopLogger()
	}
	m := &Manager{
		append:        app,
		logger:        logger,
		scrapeConfigs: make(map[string]*config.ScrapeConfig),
		scrapePools:   make(map[string]*scrapePool),
		graceShut:     make(chan struct{}),
		triggerReload: make(chan struct{}, 1),
	}
	targetMetadataCache.registerManager(m)

	return m
}

// Manager maintains a set of scrape pools and manages start/stop cycles
// when receiving new target groups from the discovery manager.
type Manager struct {
	logger    log.Logger
	append    storage.Appendable
	graceShut chan struct{}

	jitterSeed    uint64     // Global jitterSeed seed is used to spread scrape workload across HA setup.
	mtxScrape     sync.Mutex // Guards the fields below.
	scrapeConfigs map[string]*config.ScrapeConfig
	scrapePools   map[string]*scrapePool
	targetSets    map[string][]*targetgroup.Group

	triggerReload chan struct{}
}

// Run receives and saves target set updates and triggers the scraping loops reloading.
// Reloading happens in the background so that it doesn't block receiving targets updates.
func (m *Manager) Run(tsets <-chan map[string][]*targetgroup.Group) error {
	go m.reloader()
	for {
		select {
		case ts := <-tsets:
			m.updateTsets(ts)

			select {
			case m.triggerReload <- struct{}{}:
			default:
			}

		case <-m.graceShut:
			return nil
		}
	}
}

func (m *Manager) reloader() {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-m.graceShut:
			return
		case <-ticker.C:
			select {
			case <-m.triggerReload:
				m.reload()
			case <-m.graceShut:
				return
			}
		}
	}
}

func (m *Manager) reload() {
	m.mtxScrape.Lock()
	var wg sync.WaitGroup
	for setName, groups := range m.targetSets {
		if _, ok := m.scrapePools[setName]; !ok {
			scrapeConfig, ok := m.scrapeConfigs[setName]
			if !ok {
				level.Error(m.logger).Log("msg", "error reloading target set", "err", "invalid config id:"+setName)
				continue
			}
			sp, err := newScrapePool(scrapeConfig, m.append, m.jitterSeed, log.With(m.logger, "scrape_pool", setName))
			if err != nil {
				level.Error(m.logger).Log("msg", "error creating new scrape pool", "err", err, "scrape_pool", setName)
				continue
			}
			m.scrapePools[setName] = sp
		}

		wg.Add(1)
		// Run the sync in parallel as these take a while and at high load can't catch up.
		go func(sp *scrapePool, groups []*targetgroup.Group) {
			sp.Sync(groups)
			wg.Done()
		}(m.scrapePools[setName], groups)

	}
	m.mtxScrape.Unlock()
	wg.Wait()
}

// setJitterSeed calculates a global jitterSeed per server relying on extra label set.
func (m *Manager) setJitterSeed(labels labels.Labels) error {
	h := fnv.New64a()
	hostname, err := getFqdn()
	if err != nil {
		return err
	}
	if _, err := fmt.Fprintf(h, "%s%s", hostname, labels.String()); err != nil {
		return err
	}
	m.jitterSeed = h.Sum64()
	return nil
}

// Stop cancels all running scrape pools and blocks until all have exited.
func (m *Manager) Stop() {
	m.mtxScrape.Lock()
	defer m.mtxScrape.Unlock()

	for _, sp := range m.scrapePools {
		sp.stop()
	}
	close(m.graceShut)
}

func (m *Manager) updateTsets(tsets map[string][]*targetgroup.Group) {
	m.mtxScrape.Lock()
	m.targetSets = tsets
	m.mtxScrape.Unlock()
}

// ApplyConfig resets the manager's target providers and job configurations as defined by the new cfg.
func (m *Manager) ApplyConfig(cfg *config.Config) error {
	m.mtxScrape.Lock()
	defer m.mtxScrape.Unlock()

	c := make(map[string]*config.ScrapeConfig)
	for _, scfg := range cfg.ScrapeConfigs {
		c[scfg.JobName] = scfg
	}
	m.scrapeConfigs = c

	if err := m.setJitterSeed(cfg.GlobalConfig.ExternalLabels); err != nil {
		return err
	}

	// Cleanup and reload pool if the configuration has changed.
	var failed bool
	for name, sp := range m.scrapePools {
		if cfg, ok := m.scrapeConfigs[name]; !ok {
			sp.stop()
			delete(m.scrapePools, name)
		} else if !reflect.DeepEqual(sp.config, cfg) {
			err := sp.reload(cfg)
			if err != nil {
				level.Error(m.logger).Log("msg", "error reloading scrape pool", "err", err, "scrape_pool", name)
				failed = true
			}
		}
	}

	if failed {
		return errors.New("failed to apply the new configuration")
	}
	return nil
}

// TargetsAll returns active and dropped targets grouped by job_name.
func (m *Manager) TargetsAll() map[string][]*Target {
	m.mtxScrape.Lock()
	defer m.mtxScrape.Unlock()

	targets := make(map[string][]*Target, len(m.scrapePools))
	for tset, sp := range m.scrapePools {
		targets[tset] = append(sp.ActiveTargets(), sp.DroppedTargets()...)
	}
	return targets
}

// TargetsActive returns the active targets currently being scraped.
func (m *Manager) TargetsActive() map[string][]*Target {
	m.mtxScrape.Lock()
	defer m.mtxScrape.Unlock()

	var (
		wg  sync.WaitGroup
		mtx sync.Mutex
	)

	targets := make(map[string][]*Target, len(m.scrapePools))
	wg.Add(len(m.scrapePools))
	for tset, sp := range m.scrapePools {
		// Running in parallel limits the blocking time of scrapePool to scrape
		// interval when there's an update from SD.
		go func(tset string, sp *scrapePool) {
			mtx.Lock()
			targets[tset] = sp.ActiveTargets()
			mtx.Unlock()
			wg.Done()
		}(tset, sp)
	}
	wg.Wait()
	return targets
}

// TargetsDropped returns the dropped targets during relabelling.
func (m *Manager) TargetsDropped() map[string][]*Target {
	m.mtxScrape.Lock()
	defer m.mtxScrape.Unlock()

	targets := make(map[string][]*Target, len(m.scrapePools))
	for tset, sp := range m.scrapePools {
		targets[tset] = sp.DroppedTargets()
	}
	return targets
}

// getFqdn returns a FQDN if it's possible, otherwise falls back to hostname.
func getFqdn() (string, error) {
	hostname, err := os.Hostname()
	if err != nil {
		return "", err
	}

	ips, err := net.LookupIP(hostname)
	if err != nil {
		// Return the system hostname if we can't look up the IP address.
		return hostname, nil
	}

	lookup := func(ipStr encoding.TextMarshaler) (string, error) {
		ip, err := ipStr.MarshalText()
		if err != nil {
			return "", err
		}
		hosts, err := net.LookupAddr(string(ip))
		if err != nil || len(hosts) == 0 {
			return "", err
		}
		return hosts[0], nil
	}

	for _, addr := range ips {
		if ip := addr.To4(); ip != nil {
			if fqdn, err := lookup(ip); err == nil {
				return fqdn, nil
			}

		}

		if ip := addr.To16(); ip != nil {
			if fqdn, err := lookup(ip); err == nil {
				return fqdn, nil
			}

		}
	}
	return hostname, nil
}
<|endoftext|>


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

package v3rpc

import (
	"io"

	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
	"github.com/coreos/etcd/storage"
)

type watchServer struct {
	watchable storage.Watchable
}

func NewWatchServer(w storage.Watchable) pb.WatchServer {
	return &watchServer{w}
}

func (ws *watchServer) Watch(stream pb.Watch_WatchServer) error {
	closec := make(chan struct{})
	defer close(closec)

	watcher := ws.watchable.NewWatcher()
	defer watcher.Close()

	go sendLoop(stream, watcher, closec)

	for {
		req, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}

		var prefix bool
		toWatch := req.Key
		if len(req.Key) == 0 {
			toWatch = req.Prefix
			prefix = true
		}
		// TODO: support cancellation
		watcher.Watch(toWatch, prefix, req.StartRevision)
	}
}

func sendLoop(stream pb.Watch_WatchServer, watcher storage.Watcher, closec chan struct{}) {
	for {
		select {
		case e, ok := <-watcher.Chan():
			if !ok {
				return
			}
			err := stream.Send(&pb.WatchResponse{Event: &e})
			storage.ReportEventReceived()
			if err != nil {
				return
			}
		case <-closec:
			// drain the chan to clean up pending events
			for {
				_, ok := <-watcher.Chan()
				if !ok {
					return
				}
				storage.ReportEventReceived()
			}
		}
	}
}
<commit_msg>etcdserver/api/v3rpc: watch.go with events slice<commit_after>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package v3rpc

import (
	"io"

	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
	"github.com/coreos/etcd/storage"
	"github.com/coreos/etcd/storage/storagepb"
)

type watchServer struct {
	watchable storage.Watchable
}

func NewWatchServer(w storage.Watchable) pb.WatchServer {
	return &watchServer{w}
}

func (ws *watchServer) Watch(stream pb.Watch_WatchServer) error {
	closec := make(chan struct{})
	defer close(closec)

	watcher := ws.watchable.NewWatcher()
	defer watcher.Close()

	go sendLoop(stream, watcher, closec)

	for {
		req, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}

		var prefix bool
		toWatch := req.Key
		if len(req.Key) == 0 {
			toWatch = req.Prefix
			prefix = true
		}
		// TODO: support cancellation
		watcher.Watch(toWatch, prefix, req.StartRevision)
	}
}

func sendLoop(stream pb.Watch_WatchServer, watcher storage.Watcher, closec chan struct{}) {
	for {
		select {
		case evs, ok := <-watcher.Chan():
			if !ok {
				return
			}

			// TODO: evs is []storagepb.Event type
			// either return []*storagepb.Event from storage package
			// or define protocol buffer with []storagepb.Event.
			events := make([]*storagepb.Event, len(evs))
			for i := range evs {
				events[i] = &evs[i]
			}

			err := stream.Send(&pb.WatchResponse{Events: events})
			storage.ReportEventReceived()
			if err != nil {
				return
			}

		case <-closec:
			// drain the chan to clean up pending events
			for {
				_, ok := <-watcher.Chan()
				if !ok {
					return
				}
				storage.ReportEventReceived()
			}
		}
	}
}
<|endoftext|>


<commit_before>package xweb

import (
	"crypto/md5"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/howeyc/fsnotify"
)

var DefaultStaticVerMgr *StaticVerMgr = new(StaticVerMgr)

type StaticVerMgr struct {
	Caches        map[string]string
	Combined      map[string][]string
	Combines      map[string]bool
	mutex         *sync.Mutex
	Path          string
	Ignores       map[string]bool
	app           *App
	TimerCallback func() bool
}

func (self *StaticVerMgr) Moniter(staticPath string) error {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return err
	}
	defer watcher.Close()
	done := make(chan bool)
	go func() {
		for {
			select {
			case ev := <-watcher.Event:
				if ev == nil {
					break
				}
				if _, ok := self.Ignores[filepath.Base(ev.Name)]; ok {
					break
				}
				d, err := os.Stat(ev.Name)
				if err != nil {
					break
				}

				if ev.IsCreate() {
					if d.IsDir() {
						watcher.Watch(ev.Name)
					} else {
						url := ev.Name[len(self.Path)+1:]
						self.CacheItem(url)
					}
				} else if ev.IsDelete() {
					if d.IsDir() {
						watcher.RemoveWatch(ev.Name)
					} else {
						url := ev.Name[len(self.Path)+1:]
						self.CacheDelete(url)
						self.DeleteCombined(url)
					}
				} else if ev.IsModify() {
					if d.IsDir() {
					} else {
						url := ev.Name[len(staticPath)+1:]
						self.CacheItem(url)
						self.DeleteCombined(url)
					}
				} else if ev.IsRename() {
					if d.IsDir() {
						watcher.RemoveWatch(ev.Name)
					} else {
						url := ev.Name[len(staticPath)+1:]
						self.CacheDelete(url)
						self.DeleteCombined(url)
					}
				}
			case err := <-watcher.Error:
				self.app.Errorf("error: %v", err)
			case <-time.After(time.Second * 2):
				if self.TimerCallback != nil {
					if self.TimerCallback() == false {
						close(done)
						return
					}
				}
			}
		}
	}()

	err = filepath.Walk(staticPath, func(f string, info os.FileInfo, err error) error {
		if info.IsDir() {
			return watcher.Watch(f)
		}
		return nil
	})

	if err != nil {
		fmt.Println(err)
		return err
	}

	<-done
	return nil
}

func (self *StaticVerMgr) Init(app *App, staticPath string) error {
	self.Path = staticPath
	self.Caches = make(map[string]string)
	self.Combined = make(map[string][]string)
	self.Combines = make(map[string]bool)
	self.mutex = &sync.Mutex{}
	self.Ignores = map[string]bool{".DS_Store": true}
	self.app = app

	if dirExists(staticPath) {
		//self.CacheAll(staticPath)
		if self.TimerCallback == nil {
			self.TimerCallback = func() bool {
				//更改模板主题后，关闭当前监控，重新监控新目录
				if self.app.AppConfig.StaticDir == self.Path {
					return true
				}
				self.Caches = make(map[string]string)
				self.Combined = make(map[string][]string)
				self.Combines = make(map[string]bool)
				self.Path = self.app.AppConfig.StaticDir
				go self.Moniter(self.Path)
				return false
			}
		}
		go self.Moniter(staticPath)
	}

	return nil
}

func (self *StaticVerMgr) getFileVer(url string) string {
	//content, err := ioutil.ReadFile(path.Join(self.Path, url))
	fPath := filepath.Join(self.Path, url)
	self.app.Debug("loaded static ", fPath)
	f, err := os.Open(fPath)
	if err != nil {
		return ""
	}
	defer f.Close()

	fInfo, err := f.Stat()
	if err != nil {
		return ""
	}

	content := make([]byte, int(fInfo.Size()))
	_, err = f.Read(content)
	if err == nil {
		h := md5.New()
		io.WriteString(h, string(content))
		return fmt.Sprintf("%x", h.Sum(nil))[0:4]
	}
	return ""
}

func (self *StaticVerMgr) CacheAll(staticPath string) error {
	self.mutex.Lock()
	defer self.mutex.Unlock()
	fmt.Print("Getting static file version number, please wait... ")
	err := filepath.Walk(staticPath, func(f string, info os.FileInfo, err error) error {
		if info.IsDir() {
			return nil
		}
		rp := f[len(staticPath)+1:]
		rp = FixDirSeparator(rp)
		if _, ok := self.Ignores[filepath.Base(rp)]; !ok {
			self.Caches[rp] = self.getFileVer(rp)
		}
		return nil
	})
	fmt.Println("Complete.")
	return err
}

func (self *StaticVerMgr) GetVersion(url string) string {
	self.mutex.Lock()
	defer self.mutex.Unlock()
	if ver, ok := self.Caches[url]; ok {
		return ver
	}

	ver := self.getFileVer(url)
	if ver != "" {
		self.Caches[url] = ver
	}
	return ver
}

func (self *StaticVerMgr) CacheDelete(url string) {
	url = FixDirSeparator(url)
	self.mutex.Lock()
	defer self.mutex.Unlock()
	if _, ok := self.Caches[url]; ok {
		delete(self.Caches, url)
		self.app.Infof("static file %s is deleted.\n", url)
	}
}

func (self *StaticVerMgr) CacheItem(url string) {
	url = FixDirSeparator(url)
	ver := self.getFileVer(url)
	if ver != "" {
		self.mutex.Lock()
		defer self.mutex.Unlock()
		self.Caches[url] = ver
		self.app.Infof("static file %s is created.", url)
	}
}

func (self *StaticVerMgr) DeleteCombined(url string) {
	url = "/" + FixDirSeparator(url)
	self.mutex.Lock()
	defer self.mutex.Unlock()
	if val, ok := self.Combined[url]; ok {
		for _, v := range val {
			if _, has := self.Combines[v]; !has {
				continue
			}
			err := os.Remove(self.app.AppConfig.StaticDir + v)
			delete(self.Combines, v)
			if err != nil {
				fmt.Println(err)
			}
		}
	}
}

func (self *StaticVerMgr) RecordCombined(fromUrl string, combineUrl string) {
	if self.Combined == nil {
		return
	}
	self.mutex.Lock()
	defer self.mutex.Unlock()
	if _, ok := self.Combined[fromUrl]; !ok {
		self.Combined[fromUrl] = make([]string, 0)
	}
	self.Combined[fromUrl] = append(self.Combined[fromUrl], combineUrl)
}

func (self *StaticVerMgr) RecordCombines(combineUrl string) {
	if self.Combines == nil {
		return
	}
	self.mutex.Lock()
	defer self.mutex.Unlock()
	self.Combines[combineUrl] = true
}

func (self *StaticVerMgr) IsCombined(combineUrl string) (ok bool) {
	if self.Combines == nil {
		return
	}
	self.mutex.Lock()
	defer self.mutex.Unlock()
	_, ok = self.Combines[combineUrl]
	return
}
<commit_msg>update<commit_after>package xweb

import (
	"crypto/md5"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/howeyc/fsnotify"
)

var DefaultStaticVerMgr *StaticVerMgr = new(StaticVerMgr)

type StaticVerMgr struct {
	Caches        map[string]string
	Combined      map[string][]string
	Combines      map[string]bool
	mutex         *sync.Mutex
	Path          string
	Ignores       map[string]bool
	app           *App
	TimerCallback func() bool
}

func (self *StaticVerMgr) Moniter(staticPath string) error {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return err
	}
	defer watcher.Close()
	done := make(chan bool)
	go func() {
		for {
			select {
			case ev := <-watcher.Event:
				if ev == nil {
					break
				}
				if _, ok := self.Ignores[filepath.Base(ev.Name)]; ok {
					break
				}
				d, err := os.Stat(ev.Name)
				if err != nil {
					break
				}

				if ev.IsCreate() {
					if d.IsDir() {
						watcher.Watch(ev.Name)
					} else {
						url := ev.Name[len(self.Path)+1:]
						self.CacheItem(url)
					}
				} else if ev.IsDelete() {
					if d.IsDir() {
						watcher.RemoveWatch(ev.Name)
					} else {
						url := ev.Name[len(self.Path)+1:]
						self.CacheDelete(url)
						self.DeleteCombined(url)
					}
				} else if ev.IsModify() {
					if d.IsDir() {
					} else {
						url := ev.Name[len(staticPath)+1:]
						self.CacheItem(url)
						self.DeleteCombined(url)
					}
				} else if ev.IsRename() {
					if d.IsDir() {
						watcher.RemoveWatch(ev.Name)
					} else {
						url := ev.Name[len(staticPath)+1:]
						self.CacheDelete(url)
						self.DeleteCombined(url)
					}
				}
			case err := <-watcher.Error:
				self.app.Errorf("error: %v", err)
			case <-time.After(time.Second * 2):
				if self.TimerCallback != nil {
					if self.TimerCallback() == false {
						close(done)
						return
					}
				}
			}
		}
	}()

	err = filepath.Walk(staticPath, func(f string, info os.FileInfo, err error) error {
		if info.IsDir() {
			return watcher.Watch(f)
		}
		return nil
	})

	if err != nil {
		fmt.Println(err)
		return err
	}

	<-done
	return nil
}

func (self *StaticVerMgr) Init(app *App, staticPath string) error {
	self.Path = staticPath
	self.Caches = make(map[string]string)
	self.Combined = make(map[string][]string)
	self.Combines = make(map[string]bool)
	self.mutex = &sync.Mutex{}
	self.Ignores = map[string]bool{".DS_Store": true}
	self.app = app

	if dirExists(staticPath) {
		//self.CacheAll(staticPath)
		if self.TimerCallback == nil {
			self.TimerCallback = func() bool {
				//更改模板主题后，关闭当前监控，重新监控新目录
				if self.app.AppConfig.StaticDir == self.Path {
					return true
				}
				for f, _ := range self.Combines {
					os.Remove(self.Path + f)
				}
				self.Caches = make(map[string]string)
				self.Combined = make(map[string][]string)
				self.Combines = make(map[string]bool)
				self.Path = self.app.AppConfig.StaticDir
				go self.Moniter(self.Path)
				return false
			}
		}
		go self.Moniter(staticPath)
	}

	return nil
}

func (self *StaticVerMgr) getFileVer(url string) string {
	//content, err := ioutil.ReadFile(path.Join(self.Path, url))
	fPath := filepath.Join(self.Path, url)
	self.app.Debug("loaded static ", fPath)
	f, err := os.Open(fPath)
	if err != nil {
		return ""
	}
	defer f.Close()

	fInfo, err := f.Stat()
	if err != nil {
		return ""
	}

	content := make([]byte, int(fInfo.Size()))
	_, err = f.Read(content)
	if err == nil {
		h := md5.New()
		io.WriteString(h, string(content))
		return fmt.Sprintf("%x", h.Sum(nil))[0:4]
	}
	return ""
}

func (self *StaticVerMgr) CacheAll(staticPath string) error {
	self.mutex.Lock()
	defer self.mutex.Unlock()
	fmt.Print("Getting static file version number, please wait... ")
	err := filepath.Walk(staticPath, func(f string, info os.FileInfo, err error) error {
		if info.IsDir() {
			return nil
		}
		rp := f[len(staticPath)+1:]
		rp = FixDirSeparator(rp)
		if _, ok := self.Ignores[filepath.Base(rp)]; !ok {
			self.Caches[rp] = self.getFileVer(rp)
		}
		return nil
	})
	fmt.Println("Complete.")
	return err
}

func (self *StaticVerMgr) GetVersion(url string) string {
	self.mutex.Lock()
	defer self.mutex.Unlock()
	if ver, ok := self.Caches[url]; ok {
		return ver
	}

	ver := self.getFileVer(url)
	if ver != "" {
		self.Caches[url] = ver
	}
	return ver
}

func (self *StaticVerMgr) CacheDelete(url string) {
	url = FixDirSeparator(url)
	self.mutex.Lock()
	defer self.mutex.Unlock()
	if _, ok := self.Caches[url]; ok {
		delete(self.Caches, url)
		self.app.Infof("static file %s is deleted.\n", url)
	}
}

func (self *StaticVerMgr) CacheItem(url string) {
	url = FixDirSeparator(url)
	ver := self.getFileVer(url)
	if ver != "" {
		self.mutex.Lock()
		defer self.mutex.Unlock()
		self.Caches[url] = ver
		self.app.Infof("static file %s is created.", url)
	}
}

func (self *StaticVerMgr) DeleteCombined(url string) {
	url = "/" + FixDirSeparator(url)
	self.mutex.Lock()
	defer self.mutex.Unlock()
	if val, ok := self.Combined[url]; ok {
		for _, v := range val {
			if _, has := self.Combines[v]; !has {
				continue
			}
			err := os.Remove(self.Path + v)
			delete(self.Combines, v)
			if err != nil {
				fmt.Println(err)
			}
		}
	}
}

func (self *StaticVerMgr) RecordCombined(fromUrl string, combineUrl string) {
	if self.Combined == nil {
		return
	}
	self.mutex.Lock()
	defer self.mutex.Unlock()
	if _, ok := self.Combined[fromUrl]; !ok {
		self.Combined[fromUrl] = make([]string, 0)
	}
	self.Combined[fromUrl] = append(self.Combined[fromUrl], combineUrl)
}

func (self *StaticVerMgr) RecordCombines(combineUrl string) {
	if self.Combines == nil {
		return
	}
	self.mutex.Lock()
	defer self.mutex.Unlock()
	self.Combines[combineUrl] = true
}

func (self *StaticVerMgr) IsCombined(combineUrl string) (ok bool) {
	if self.Combines == nil {
		return
	}
	self.mutex.Lock()
	defer self.mutex.Unlock()
	_, ok = self.Combines[combineUrl]
	return
}
<|endoftext|>


<commit_before>package merkle

import (
	"hash"
	"log"
	"runtime"
)

// NewHash provides a hash.Hash to generate a merkle.Tree checksum, given a
// HashMaker for the checksums of the blocks written and the blockSize of each
// block per node in the tree.
func NewHash(hm HashMaker, merkleBlockLength int) HashTreeer {
	return newMerkleHash(hm, merkleBlockLength)
}

func newMerkleHash(hm HashMaker, merkleBlockLength int) *merkleHash {
	mh := new(merkleHash)
	mh.blockSize = merkleBlockLength
	mh.hm = hm
	mh.tree = &Tree{Nodes: []*Node{}, BlockLength: merkleBlockLength}
	mh.lastBlock = make([]byte, merkleBlockLength)
	return mh
}

// Treeer (Tree-er) provides access to the Merkle tree internals
type Treeer interface {
	Nodes() []*Node
	Root() *Node
}

// HashTreeer can be used as a hash.Hash but also provide access to the Merkle
// tree internals
type HashTreeer interface {
	hash.Hash
	Treeer
}

// TODO make a similar hash.Hash, that accepts an argument of a merkle.Tree,
// that will validate nodes as the new bytes are written. If a new written
// block fails checksum, then return an error on the io.Writer

type merkleHash struct {
	blockSize       int
	tree            *Tree
	hm              HashMaker
	lastBlock       []byte // as needed, for Sum()
	lastBlockLen    int
	partialLastNode bool // true when Sum() has appended a Node for a partial block
}

func (mh *merkleHash) Reset() {
	mh1 := newMerkleHash(mh.hm, mh.blockSize)
	*mh = *mh1
}

func (mh merkleHash) Nodes() []*Node {
	return mh.tree.Nodes
}

func (mh merkleHash) Root() *Node {
	return mh.tree.Root()
}

// XXX this will be tricky, as the last block can be less than the BlockSize.
// if they get the sum, it will be mh.tree.Root().Checksum() at that point.
//
// But if they continue writing, it would mean a continuation of the bytes in
// the last block. So popping the last node, and having a buffer for the bytes
// in that last partial block.
//
// if that last block was complete, then no worries. start the next node.
func (mh *merkleHash) Sum(b []byte) []byte {
	var (
		curBlock = []byte{}
		offset   int
	)
	if mh.partialLastNode {
		// if this is true, then we need to pop the last node
		mh.tree.Nodes = mh.tree.Nodes[:len(mh.tree.Nodes)-1]
		mh.partialLastNode = false
	}

	if mh.lastBlockLen > 0 {
		curBlock = append(curBlock[:], mh.lastBlock[:mh.lastBlockLen]...)
		mh.lastBlockLen = 0
	}

	if b != nil && len(b) > 0 {
		curBlock = append(curBlock, b...)
	}

	// incase we're at a new or reset state
	if len(mh.tree.Nodes) == 0 && len(curBlock) == 0 {
		return nil
	}

	for i := 0; i < len(curBlock)/mh.blockSize; i++ {
		n, err := NewNodeHashBlock(mh.hm, curBlock[offset:(offset+mh.blockSize)])
		if err != nil {
			// XXX i hate to swallow an error here, but the `Sum() []byte` signature
			// :-\
			sBuf := make([]byte, 1024)
			runtime.Stack(sBuf, false)
			log.Printf("[ERROR]: %s %q", err, string(sBuf))
			return nil
		}
		mh.tree.Nodes = append(mh.tree.Nodes, n)
		offset = offset + mh.blockSize
	}

	// If there is remainder, we'll need to make a partial node and stash it
	if m := (len(curBlock) % mh.blockSize); m != 0 {
		mh.lastBlockLen = copy(mh.lastBlock, curBlock[offset:])

		n, err := NewNodeHashBlock(mh.hm, curBlock[offset:])
		if err != nil {
			sBuf := make([]byte, 1024)
			runtime.Stack(sBuf, false)
			log.Printf("[ERROR]: %s %q", err, string(sBuf))
			return nil
		}
		mh.tree.Nodes = append(mh.tree.Nodes, n)
		mh.partialLastNode = true
	}

	sum, err := mh.tree.Root().Checksum()
	if err != nil {
		// XXX i hate to swallow an error here, but the `Sum() []byte` signature
		// :-\
		sBuf := make([]byte, 1024)
		runtime.Stack(sBuf, false)
		log.Printf("[ERROR]: %s %q", err, string(sBuf))
		return nil
	}
	return sum
}

func (mh *merkleHash) Write(b []byte) (int, error) {
	// basically we need to:
	// * include prior partial lastBlock, if any
	// * chunk these writes into blockSize
	// * create Node of the sum
	// * add the Node to the tree
	// * stash remainder in the mh.lastBlock

	var (
		curBlock   = make([]byte, mh.blockSize)
		numBytes   int
		numWritten int
		offset     int
	)
	if mh.lastBlock != nil && mh.lastBlockLen > 0 {
		if (mh.lastBlockLen + len(b)) < mh.blockSize {
			mh.lastBlockLen += copy(mh.lastBlock[mh.lastBlockLen:], b[:])
			return len(b), nil
		}

		//                                         XXX off by one?
		numBytes = copy(curBlock[:], mh.lastBlock[:mh.lastBlockLen])
		// not adding to numWritten, since these blocks were accounted for in a
		// prior Write()

		// then we'll chunk the front of the incoming bytes
		end := mh.blockSize - numBytes
		if end > len(b) {
			end = len(b)
		}
		offset = copy(curBlock[numBytes:], b[:end])
		n, err := NewNodeHashBlock(mh.hm, curBlock)
		if err != nil {
			// XXX might need to stash again the prior lastBlock and first little chunk
			return numWritten, err
		}
		mh.tree.Nodes = append(mh.tree.Nodes, n)
		numWritten += offset
	}

	numBytes = (len(b) - offset)
	for i := 0; i < numBytes/mh.blockSize; i++ {
		//fmt.Printf("%s", b[offset:offset+mh.blockSize])
		numWritten += copy(curBlock, b[offset:offset+mh.blockSize])
		n, err := NewNodeHashBlock(mh.hm, curBlock)
		if err != nil {
			// XXX might need to stash again the prior lastBlock and first little chunk
			return numWritten, err
		}
		mh.tree.Nodes = append(mh.tree.Nodes, n)
		offset = offset + mh.blockSize
	}

	mh.lastBlockLen = numBytes % mh.blockSize
	//                                       XXX off by one?
	numWritten += copy(mh.lastBlock[:], b[(len(b)-mh.lastBlockLen):])

	return numWritten, nil
}

// likely not the best to pass this through and not use our own node block
// size, but let's revisit this.
func (mh *merkleHash) BlockSize() int { return mh.hm().BlockSize() }
func (mh *merkleHash) Size() int      { return mh.hm().Size() }
<commit_msg>stream: reduce allocs in Reset()<commit_after>package merkle

import (
	"hash"
	"log"
	"runtime"
)

// NewHash provides a hash.Hash to generate a merkle.Tree checksum, given a
// HashMaker for the checksums of the blocks written and the blockSize of each
// block per node in the tree.
func NewHash(hm HashMaker, merkleBlockLength int) HashTreeer {
	return newMerkleHash(hm, merkleBlockLength)
}

func newMerkleHash(hm HashMaker, merkleBlockLength int) *merkleHash {
	mh := new(merkleHash)
	mh.blockSize = merkleBlockLength
	mh.hm = hm
	mh.tree = &Tree{Nodes: []*Node{}, BlockLength: merkleBlockLength}
	mh.lastBlock = make([]byte, merkleBlockLength)
	return mh
}

// Treeer (Tree-er) provides access to the Merkle tree internals
type Treeer interface {
	Nodes() []*Node
	Root() *Node
}

// HashTreeer can be used as a hash.Hash but also provide access to the Merkle
// tree internals
type HashTreeer interface {
	hash.Hash
	Treeer
}

// TODO make a similar hash.Hash, that accepts an argument of a merkle.Tree,
// that will validate nodes as the new bytes are written. If a new written
// block fails checksum, then return an error on the io.Writer

type merkleHash struct {
	blockSize       int
	tree            *Tree
	hm              HashMaker
	lastBlock       []byte // as needed, for Sum()
	lastBlockLen    int
	partialLastNode bool // true when Sum() has appended a Node for a partial block
}

func (mh *merkleHash) Reset() {
	mh.tree = &Tree{Nodes: []*Node{}, BlockLength: mh.blockSize}
	mh.lastBlockLen = 0
	mh.partialLastNode = false
}

func (mh merkleHash) Nodes() []*Node {
	return mh.tree.Nodes
}

func (mh merkleHash) Root() *Node {
	return mh.tree.Root()
}

// XXX this will be tricky, as the last block can be less than the BlockSize.
// if they get the sum, it will be mh.tree.Root().Checksum() at that point.
//
// But if they continue writing, it would mean a continuation of the bytes in
// the last block. So popping the last node, and having a buffer for the bytes
// in that last partial block.
//
// if that last block was complete, then no worries. start the next node.
func (mh *merkleHash) Sum(b []byte) []byte {
	var (
		curBlock = []byte{}
		offset   int
	)
	if mh.partialLastNode {
		// if this is true, then we need to pop the last node
		mh.tree.Nodes = mh.tree.Nodes[:len(mh.tree.Nodes)-1]
		mh.partialLastNode = false
	}

	if mh.lastBlockLen > 0 {
		curBlock = append(curBlock[:], mh.lastBlock[:mh.lastBlockLen]...)
		mh.lastBlockLen = 0
	}

	if b != nil && len(b) > 0 {
		curBlock = append(curBlock, b...)
	}

	// incase we're at a new or reset state
	if len(mh.tree.Nodes) == 0 && len(curBlock) == 0 {
		return nil
	}

	for i := 0; i < len(curBlock)/mh.blockSize; i++ {
		n, err := NewNodeHashBlock(mh.hm, curBlock[offset:(offset+mh.blockSize)])
		if err != nil {
			// XXX i hate to swallow an error here, but the `Sum() []byte` signature
			// :-\
			sBuf := make([]byte, 1024)
			runtime.Stack(sBuf, false)
			log.Printf("[ERROR]: %s %q", err, string(sBuf))
			return nil
		}
		mh.tree.Nodes = append(mh.tree.Nodes, n)
		offset = offset + mh.blockSize
	}

	// If there is remainder, we'll need to make a partial node and stash it
	if m := (len(curBlock) % mh.blockSize); m != 0 {
		mh.lastBlockLen = copy(mh.lastBlock, curBlock[offset:])

		n, err := NewNodeHashBlock(mh.hm, curBlock[offset:])
		if err != nil {
			sBuf := make([]byte, 1024)
			runtime.Stack(sBuf, false)
			log.Printf("[ERROR]: %s %q", err, string(sBuf))
			return nil
		}
		mh.tree.Nodes = append(mh.tree.Nodes, n)
		mh.partialLastNode = true
	}

	sum, err := mh.tree.Root().Checksum()
	if err != nil {
		// XXX i hate to swallow an error here, but the `Sum() []byte` signature
		// :-\
		sBuf := make([]byte, 1024)
		runtime.Stack(sBuf, false)
		log.Printf("[ERROR]: %s %q", err, string(sBuf))
		return nil
	}
	return sum
}

func (mh *merkleHash) Write(b []byte) (int, error) {
	// basically we need to:
	// * include prior partial lastBlock, if any
	// * chunk these writes into blockSize
	// * create Node of the sum
	// * add the Node to the tree
	// * stash remainder in the mh.lastBlock

	var (
		curBlock   = make([]byte, mh.blockSize)
		numBytes   int
		numWritten int
		offset     int
	)
	if mh.lastBlock != nil && mh.lastBlockLen > 0 {
		if (mh.lastBlockLen + len(b)) < mh.blockSize {
			mh.lastBlockLen += copy(mh.lastBlock[mh.lastBlockLen:], b[:])
			return len(b), nil
		}

		//                                         XXX off by one?
		numBytes = copy(curBlock[:], mh.lastBlock[:mh.lastBlockLen])
		// not adding to numWritten, since these blocks were accounted for in a
		// prior Write()

		// then we'll chunk the front of the incoming bytes
		end := mh.blockSize - numBytes
		if end > len(b) {
			end = len(b)
		}
		offset = copy(curBlock[numBytes:], b[:end])
		n, err := NewNodeHashBlock(mh.hm, curBlock)
		if err != nil {
			// XXX might need to stash again the prior lastBlock and first little chunk
			return numWritten, err
		}
		mh.tree.Nodes = append(mh.tree.Nodes, n)
		numWritten += offset
	}

	numBytes = (len(b) - offset)
	for i := 0; i < numBytes/mh.blockSize; i++ {
		//fmt.Printf("%s", b[offset:offset+mh.blockSize])
		numWritten += copy(curBlock, b[offset:offset+mh.blockSize])
		n, err := NewNodeHashBlock(mh.hm, curBlock)
		if err != nil {
			// XXX might need to stash again the prior lastBlock and first little chunk
			return numWritten, err
		}
		mh.tree.Nodes = append(mh.tree.Nodes, n)
		offset = offset + mh.blockSize
	}

	mh.lastBlockLen = numBytes % mh.blockSize
	//                                       XXX off by one?
	numWritten += copy(mh.lastBlock[:], b[(len(b)-mh.lastBlockLen):])

	return numWritten, nil
}

// likely not the best to pass this through and not use our own node block
// size, but let's revisit this.
func (mh *merkleHash) BlockSize() int { return mh.hm().BlockSize() }
func (mh *merkleHash) Size() int      { return mh.hm().Size() }
<|endoftext|>