<commit_before>package main

import (
	"encoding/json"
	"github.com/dustin/go-humanize"
	"github.com/flosch/pongo2"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"time"
)

type Server struct {
	Url  string
	Path string
}

type Result struct {
	Servers   []Server
	Filename  string
	Size      uint64
	HumanSize string
}

type SearchConfig struct {
	Frontend *Frontend
}

type Search struct {
	cfg  SearchConfig
	tmpl *pongo2.Template
}

func CreateSearch(cfg SearchConfig) (search *Search, err error) {
	search = &Search{cfg: cfg}

	// Load the results template
	search.tmpl, err = search.cfg.Frontend.templates.FromFile("results.tmpl")
	if err != nil {
		return
	}

	return
}

func (search *Search) Handler(w http.ResponseWriter, r *http.Request) {
	start := time.Now()

	defer func() {
		if err, ok := recover().(error); ok {
			log.Println(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}
	}()

	query := r.FormValue("q")
	page, err := strconv.Atoi(r.FormValue("p"))
	if err != nil {
		page = 0
	}

	resp, err := search.cfg.Frontend.elasticSearch.Search(query, search.cfg.Frontend.cfg.PerPage, page)
	if err != nil {
		search.cfg.Frontend.Log.Panic(err)
	}

	// API like a bauss
	if r.FormValue("f") == "json" {
		output, err := json.Marshal(resp.Hits)
		if err != nil {
			search.cfg.Frontend.Log.Panic(err)
		}

		w.Header().Set("Content-Type", "application/json")
		_, err = w.Write(output)
		if err != nil {
			search.cfg.Frontend.Log.Panic(err)
		}

		return
	}

	var results []Result
	for _, qr := range resp.Hits.Hits {
		raw, err := qr.Source.MarshalJSON()
		if err != nil {
			search.cfg.Frontend.Log.Panic(err)
		}

		var result Result
		err = json.Unmarshal(raw, &result)
		if err != nil {
			search.cfg.Frontend.Log.Panic(err)
		}

		result.HumanSize = humanize.Bytes(result.Size)
		results = append(results, result)
	}

	search.tmpl.ExecuteWriter(pongo2.Context{
		"query": query,

		"page":     page,
		"frompage": search.cfg.Frontend.cfg.PerPage * page,
		"maxpages": resp.Hits.Total / search.cfg.Frontend.cfg.PerPage,
		"prevpage": search.getPageLink(page-1, r.URL),
		"nextpage": search.getPageLink(page+1, r.URL),

		"elapsed":  time.Since(start) / time.Millisecond,
		"response": resp,
		"results":  results,
	}, w)
}

func (search *Search) getPageLink(page int, inURL *url.URL) (outURL string) {
	if page < 1 {
		page = 0
	}

	qry := inURL.Query()
	qry.Set("p", strconv.Itoa(page))
	inURL.RawQuery = qry.Encode()

	outURL = inURL.String()

	return
}
<commit_msg>Frontend: Add some extra documentation to the code (related to #13)<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/dustin/go-humanize"
	"github.com/flosch/pongo2"
	"net/http"
	"net/url"
	"strconv"
	"time"
)

type Server struct {
	Url  string
	Path string
}

type Result struct {
	Servers   []Server
	Filename  string
	Size      uint64
	HumanSize string
}

type SearchConfig struct {
	Frontend *Frontend
}

type Search struct {
	cfg  SearchConfig
	tmpl *pongo2.Template
}

func CreateSearch(cfg SearchConfig) (search *Search, err error) {
	search = &Search{cfg: cfg}

	// Load the results template
	search.tmpl, err = search.cfg.Frontend.templates.FromFile("results.tmpl")
	if err != nil {
		return
	}

	return
}

func (search *Search) Handler(w http.ResponseWriter, r *http.Request) {
	start := time.Now()

	// Catch errors in the following code, log them and return a HTTP 500
	defer func() {
		if err := recover(); err != nil {
			search.cfg.Frontend.Log.Println(err)
			http.Error(w, fmt.Sprintf("%s", err), http.StatusInternalServerError)
		}
	}()

	/* Parse GET parameters
	 * q: Search query
	 * f: Format. Default is HTML, currently supported options: "json"
	 * p: Current page. Is zero if not a number
	 */
	query := r.FormValue("q")
	format := r.FormValue("f")
	page, err := strconv.Atoi(r.FormValue("p"))
	if err != nil {
		page = 0
	}

	// Do the actual search
	resp, err := search.cfg.Frontend.elasticSearch.Search(query, search.cfg.Frontend.cfg.PerPage, page)
	if err != nil {
		search.cfg.Frontend.Log.Panic(err)
	}

	// Format: JSON
	if format == "json" {
		output, err := json.Marshal(resp.Hits)
		if err != nil {
			search.cfg.Frontend.Log.Panic(err)
		}

		w.Header().Set("Content-Type", "application/json")
		_, err = w.Write(output)
		if err != nil {
			search.cfg.Frontend.Log.Panic(err)
		}

		return
	}

	// Format: HTML (default)
	var results []Result
	for _, qr := range resp.Hits.Hits {
		// Parse the search result into a Result struct
		var result Result
		err := unmarshalRawJson(qr.Source, &result)
		if err != nil {
			search.cfg.Frontend.Log.Panic(err)
		}

		// Humanize the file size
		result.HumanSize = humanize.Bytes(result.Size)
		results = append(results, result)
	}

	search.tmpl.ExecuteWriter(pongo2.Context{
		"query": query,

		"page":     page,
		"frompage": search.cfg.Frontend.cfg.PerPage * page,
		"maxpages": resp.Hits.Total / search.cfg.Frontend.cfg.PerPage,
		"prevpage": getPageLink(page-1, r.URL),
		"nextpage": getPageLink(page+1, r.URL),

		"elapsed":  time.Since(start) / time.Millisecond,
		"response": resp,
		"results":  results,
	}, w)
}

func getPageLink(page int, inURL *url.URL) (outURL string) {
	if page < 1 {
		page = 0
	}

	qry := inURL.Query()
	qry.Set("p", strconv.Itoa(page))
	inURL.RawQuery = qry.Encode()

	outURL = inURL.String()

	return
}

func unmarshalRawJson(input *json.RawMessage, output interface{}) (err error) {
	raw, err := input.MarshalJSON()
	if err != nil {
		return
	}

	err = json.Unmarshal(raw, &output)
	if err != nil {
		return
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/richo/roving/types"
)

type Server struct {
	hostport string
}

type Fuzzer struct {
	cmd     *exec.Cmd
	Id      string
	started bool
}

var preExisting bool = false

func newFuzzer() Fuzzer {
	name, err := os.Hostname()
	if err != nil {
		log.Fatal("Couldn't get hostname", err)
	}

	number := types.RandInt() & 0xffff

	return Fuzzer{
		Id:      fmt.Sprintf("%s-%x", name, number),
		started: false,
	}
}

type WatchDog struct {
	Interval time.Duration
	Fuzzer   *Fuzzer
	Server   *Server
}

func (f *Fuzzer) run() {
	f.cmd = exec.Command(f.path(),
		"-o", "output",
		"-i", "input",
		"-S", f.Id,
		"./target",
	)
	stdout, err := f.cmd.StdoutPipe()
	if err != nil {
		log.Fatalf("Couldn't get stdout handle", err)
	}

	stderr, err := f.cmd.StderrPipe()
	if err != nil {
		log.Fatalf("Couldn't get stderr handle", err)
	}

	if err := f.cmd.Start(); err != nil {
		log.Fatalf("Couldn't start fuzzer", err)
	}

	go func() {
		io.Copy(os.Stdout, stdout)
	}()

	go func() {
		io.Copy(os.Stderr, stderr)
	}()

	log.Printf("Waiting for fuzzer to exit")
	f.started = true

	err = f.cmd.Wait()
	if err != nil {
		log.Fatal(err)
	}
}

func (f *Fuzzer) stop() {
	log.Printf("Stopping the fuzzer")
	f.cmd.Process.Signal(syscall.SIGSTOP)
}

func (f *Fuzzer) start() {
	log.Printf("Starting the fuzzer")
	f.cmd.Process.Signal(syscall.SIGCONT)
}

func (f *Fuzzer) State() types.State {
	state := types.State{
		Id:      f.Id,
		Queue:   types.ReadQueue(fmt.Sprintf("output/%s/queue", f.Id)),
		Crashes: types.ReadDir(fmt.Sprintf("output/%s/crashes", f.Id)),
		Hangs:   types.ReadDir(fmt.Sprintf("output/%s/hangs", f.Id)),
	}

	return state
}

func (f *Fuzzer) path() string {
	root := os.Getenv("AFL")
	if root == "" { // Not found, hopefully it's in PATH
		return "afl-fuzz"
	}
	return fmt.Sprintf("%s/afl-fuzz", root)
}

func (w *WatchDog) run() {
	log.Printf("Starting watchdog goroutine")
	ticker := time.NewTicker(w.Interval)

	for {
		select {
		case <-ticker.C:
			w.syncState()
		}
	}
}

func (w *WatchDog) syncState() {
	if w.Fuzzer.started {
		w.Fuzzer.stop()
		defer w.Fuzzer.start()
	}

	log.Printf("Uploading our corpus")
	state := w.Fuzzer.State()
	w.Server.UploadState(state)

	log.Printf("Downloading their corpus")
	other := w.Server.FetchState(w.Fuzzer.Id)

	log.Printf("Unpacking their state")
	w.Fuzzer.UnpackStates(other)
}

func (s *Server) fetchToFile(resource, file string) error {
	target := s.getPath(resource)
	defer target.Body.Close()

	f, err := os.OpenFile(file, os.O_WRONLY|os.O_CREATE, 0755)

	if err != nil {
		return err
	}

	io.Copy(f, target.Body)

	return nil
}

func (s *Server) getPath(path string) *http.Response {
	resource := fmt.Sprintf("http://%s/%s", s.hostport, path)
	log.Printf("Fetching %s", resource)

	resp, err := http.Get(resource)
	if err != nil {
		log.Panicf("Couldn't fetch target", err)
	}

	return resp
}

func (s *Server) FetchTarget() {
	if err := s.fetchToFile("target", "target"); err != nil {
		if preExisting {
			log.Printf("Couldn't write target, ignoring since this tree is preexisting")
		} else {
			log.Fatalf("Couldn't fetch target: %s", err)
		}
	}
}

func base64ToPath(content, path string) {
	f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0755)

	if err != nil {
		log.Panicf("Couldn't open %s for writing", path, err)
	}

	body, err := base64.StdEncoding.DecodeString(content)
	if err != nil {
		log.Fatalf("Couldn't decode body for %s", path, err)
	}
	f.Write([]byte(body))
}

func (s *Server) FetchInputs() {
	os.Mkdir("input", 0755)

	inputs := s.getPath("inputs")
	defer inputs.Body.Close()

	inps := &types.InputCorpus{}

	encoder := json.NewDecoder(inputs.Body)
	encoder.Decode(&inps)

	for _, inp := range inps.Inputs {
		path := fmt.Sprintf("input/%s", inp.Name)
		base64ToPath(inp.Body, path)
	}
}

func (s *Fuzzer) UnpackStates(other []types.State) {
	for _, state := range other {
		cmd := exec.Command("tar", "-kxjf", "-")

		stdin, err := cmd.StdinPipe()
		if err != nil {
			log.Fatalf("Couldn't get stdin handle", err)
		}

		body, err := base64.StdEncoding.DecodeString(state.Queue)
		if err != nil {
			log.Fatal("Couldn't decode queue from %s", state.Id, err)
		}

		go func() {
			_, err := stdin.Write(body)
			if err != nil {
				log.Fatal("Error writing data into tar", err)
			}
			stdin.Close()
		}()

		_ = cmd.Run()
		// YOLO
		// if err != nil {
		// 	log.Fatalf("Couldn't untar state from %s", state.Id, err)
		// }
	}
}

func (s *Server) FetchState(id string) []types.State {
	var state []types.State
	resp := s.getPath(fmt.Sprintf("state/%s", id))

	encoder := json.NewDecoder(resp.Body)
	encoder.Decode(&state)

	return state
}

func (s *Server) UploadState(state types.State) {
	data, err := json.Marshal(state)
	if err != nil {
		log.Fatal("Couldn't marshall state", err)
	}

	buffer := bytes.NewReader(data)

	resource := fmt.Sprintf("http://%s/%s", s.hostport, "state")
	_, err = http.Post(resource, "application/json", buffer)

	if err != nil {
		log.Panicf("Couldn't upload state", err)
	}
}

func setupWorkDir() {
	var err error
	// TODO(richo) Ephemeral workdirs for concurrency
	if err = os.Mkdir("work", 0755); err != nil {
		log.Println("Workdir already exists, assuming we're joining an existing run")
		preExisting = true
	}
	if err = os.Chdir("work"); err != nil {
		log.Panicf("Couldn't change to workdir", err)
	}
}

func main() {
	args := os.Args
	if len(args) != 2 {
		log.Printf("Usage: ./client server:port")
		return
	}

	setupWorkDir()

	server := Server{args[1]}
	server.FetchTarget()
	server.FetchInputs()

	fuzzer := newFuzzer()

	log.Printf("Brought up a fuzzer with id %s", fuzzer.Id)

	var interval time.Duration
	if os.Getenv("SHORT_INTERVAL") != "" {
		interval = 3 * time.Second
	} else {
		interval = 10 * time.Minute
	}

	watchdog := WatchDog{
		Interval: interval,
		Fuzzer:   &fuzzer,
		Server:   &server,
	}

	log.Printf("Priming with upstream state")
	watchdog.syncState()

	go watchdog.run()

	fuzzer.run()

	// TODO(richo) Fetch the input corpus
}
<commit_msg>It's also not possible to upload a queue that doesn't exist<commit_after>package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/richo/roving/types"
)

type Server struct {
	hostport string
}

type Fuzzer struct {
	cmd     *exec.Cmd
	Id      string
	started bool
}

var preExisting bool = false

func newFuzzer() Fuzzer {
	name, err := os.Hostname()
	if err != nil {
		log.Fatal("Couldn't get hostname", err)
	}

	number := types.RandInt() & 0xffff

	return Fuzzer{
		Id:      fmt.Sprintf("%s-%x", name, number),
		started: false,
	}
}

type WatchDog struct {
	Interval time.Duration
	Fuzzer   *Fuzzer
	Server   *Server
}

func (f *Fuzzer) run() {
	f.cmd = exec.Command(f.path(),
		"-o", "output",
		"-i", "input",
		"-S", f.Id,
		"./target",
	)
	stdout, err := f.cmd.StdoutPipe()
	if err != nil {
		log.Fatalf("Couldn't get stdout handle", err)
	}

	stderr, err := f.cmd.StderrPipe()
	if err != nil {
		log.Fatalf("Couldn't get stderr handle", err)
	}

	if err := f.cmd.Start(); err != nil {
		log.Fatalf("Couldn't start fuzzer", err)
	}

	go func() {
		io.Copy(os.Stdout, stdout)
	}()

	go func() {
		io.Copy(os.Stderr, stderr)
	}()

	log.Printf("Waiting for fuzzer to exit")
	f.started = true

	err = f.cmd.Wait()
	if err != nil {
		log.Fatal(err)
	}
}

func (f *Fuzzer) stop() {
	log.Printf("Stopping the fuzzer")
	f.cmd.Process.Signal(syscall.SIGSTOP)
}

func (f *Fuzzer) start() {
	log.Printf("Starting the fuzzer")
	f.cmd.Process.Signal(syscall.SIGCONT)
}

func (f *Fuzzer) State() types.State {
	state := types.State{
		Id:      f.Id,
		Queue:   types.ReadQueue(fmt.Sprintf("output/%s/queue", f.Id)),
		Crashes: types.ReadDir(fmt.Sprintf("output/%s/crashes", f.Id)),
		Hangs:   types.ReadDir(fmt.Sprintf("output/%s/hangs", f.Id)),
	}

	return state
}

func (f *Fuzzer) path() string {
	root := os.Getenv("AFL")
	if root == "" { // Not found, hopefully it's in PATH
		return "afl-fuzz"
	}
	return fmt.Sprintf("%s/afl-fuzz", root)
}

func (w *WatchDog) run() {
	log.Printf("Starting watchdog goroutine")
	ticker := time.NewTicker(w.Interval)

	for {
		select {
		case <-ticker.C:
			w.syncState()
		}
	}
}

func (w *WatchDog) syncState() {
	if w.Fuzzer.started {
		w.Fuzzer.stop()
		defer w.Fuzzer.start()

		log.Printf("Uploading our corpus")
		state := w.Fuzzer.State()
		w.Server.UploadState(state)
	}

	log.Printf("Downloading their corpus")
	other := w.Server.FetchState(w.Fuzzer.Id)

	log.Printf("Unpacking their state")
	w.Fuzzer.UnpackStates(other)
}

func (s *Server) fetchToFile(resource, file string) error {
	target := s.getPath(resource)
	defer target.Body.Close()

	f, err := os.OpenFile(file, os.O_WRONLY|os.O_CREATE, 0755)

	if err != nil {
		return err
	}

	io.Copy(f, target.Body)

	return nil
}

func (s *Server) getPath(path string) *http.Response {
	resource := fmt.Sprintf("http://%s/%s", s.hostport, path)
	log.Printf("Fetching %s", resource)

	resp, err := http.Get(resource)
	if err != nil {
		log.Panicf("Couldn't fetch target", err)
	}

	return resp
}

func (s *Server) FetchTarget() {
	if err := s.fetchToFile("target", "target"); err != nil {
		if preExisting {
			log.Printf("Couldn't write target, ignoring since this tree is preexisting")
		} else {
			log.Fatalf("Couldn't fetch target: %s", err)
		}
	}
}

func base64ToPath(content, path string) {
	f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0755)

	if err != nil {
		log.Panicf("Couldn't open %s for writing", path, err)
	}

	body, err := base64.StdEncoding.DecodeString(content)
	if err != nil {
		log.Fatalf("Couldn't decode body for %s", path, err)
	}
	f.Write([]byte(body))
}

func (s *Server) FetchInputs() {
	os.Mkdir("input", 0755)

	inputs := s.getPath("inputs")
	defer inputs.Body.Close()

	inps := &types.InputCorpus{}

	encoder := json.NewDecoder(inputs.Body)
	encoder.Decode(&inps)

	for _, inp := range inps.Inputs {
		path := fmt.Sprintf("input/%s", inp.Name)
		base64ToPath(inp.Body, path)
	}
}

func (s *Fuzzer) UnpackStates(other []types.State) {
	for _, state := range other {
		cmd := exec.Command("tar", "-kxjf", "-")

		stdin, err := cmd.StdinPipe()
		if err != nil {
			log.Fatalf("Couldn't get stdin handle", err)
		}

		body, err := base64.StdEncoding.DecodeString(state.Queue)
		if err != nil {
			log.Fatal("Couldn't decode queue from %s", state.Id, err)
		}

		go func() {
			_, err := stdin.Write(body)
			if err != nil {
				log.Fatal("Error writing data into tar", err)
			}
			stdin.Close()
		}()

		_ = cmd.Run()
		// YOLO
		// if err != nil {
		// 	log.Fatalf("Couldn't untar state from %s", state.Id, err)
		// }
	}
}

func (s *Server) FetchState(id string) []types.State {
	var state []types.State
	resp := s.getPath(fmt.Sprintf("state/%s", id))

	encoder := json.NewDecoder(resp.Body)
	encoder.Decode(&state)

	return state
}

func (s *Server) UploadState(state types.State) {
	data, err := json.Marshal(state)
	if err != nil {
		log.Fatal("Couldn't marshall state", err)
	}

	buffer := bytes.NewReader(data)

	resource := fmt.Sprintf("http://%s/%s", s.hostport, "state")
	_, err = http.Post(resource, "application/json", buffer)

	if err != nil {
		log.Panicf("Couldn't upload state", err)
	}
}

func setupWorkDir() {
	var err error
	// TODO(richo) Ephemeral workdirs for concurrency
	if err = os.Mkdir("work", 0755); err != nil {
		log.Println("Workdir already exists, assuming we're joining an existing run")
		preExisting = true
	}
	if err = os.Chdir("work"); err != nil {
		log.Panicf("Couldn't change to workdir", err)
	}
}

func main() {
	args := os.Args
	if len(args) != 2 {
		log.Printf("Usage: ./client server:port")
		return
	}

	setupWorkDir()

	server := Server{args[1]}
	server.FetchTarget()
	server.FetchInputs()

	fuzzer := newFuzzer()

	log.Printf("Brought up a fuzzer with id %s", fuzzer.Id)

	var interval time.Duration
	if os.Getenv("SHORT_INTERVAL") != "" {
		interval = 3 * time.Second
	} else {
		interval = 10 * time.Minute
	}

	watchdog := WatchDog{
		Interval: interval,
		Fuzzer:   &fuzzer,
		Server:   &server,
	}

	log.Printf("Priming with upstream state")
	watchdog.syncState()

	go watchdog.run()

	fuzzer.run()

	// TODO(richo) Fetch the input corpus
}
<|endoftext|>


<commit_before>package trace

import (
	"context"
	"fmt"
	"github.com/shawnfeng/sutil/sconf/center"
	"github.com/shawnfeng/sutil/slog/slog"
	"net/http"
	"strings"
	"sync"
)

const FilterUrls = "span_filter_urls"

const ListConfigSep = ","

var once sync.Once

var apolloCenter center.ConfigCenter

var apolloSpanFilterConfig *spanFilterConfig

func InitTraceSpanFilter() error {
	fun := "TraceSpanFilter.init --> "
	ctx := context.Background()

	initApolloCenter(ctx)
	if apolloCenter == nil {
		return fmt.Errorf("trace apollo center is nil")
	}

	urls, ok := apolloCenter.GetStringWithNamespace(ctx, center.DefaultApolloTraceNamespace, FilterUrls)
	if !ok {
		return fmt.Errorf("not get %s from apollo namespace %s", FilterUrls, center.DefaultApolloTraceNamespace)
	}
	slog.Infof(ctx, "%s get %s from apollo: %s", fun, FilterUrls, urls)

	urlList := strings.Split(urls, ListConfigSep)

	apolloSpanFilterConfig = &spanFilterConfig{
		urls: urlList,
	}

	apolloCenter.RegisterObserver(ctx, apolloSpanFilterConfig)
	return nil
}

func initApolloCenter(ctx context.Context) {
	fun := "ApolloCenter.init --> "

	if apolloCenter != nil {
		return
	}

	once.Do(func() {
		var err error
		apolloCenter, err = center.NewConfigCenter(center.ApolloConfigCenter)
		if err != nil {
			slog.Errorf(ctx, "%s new config center error, center type: %d, err: %s", fun, center.ApolloConfigCenter, err.Error())
			return
		}

		err = apolloCenter.Init(ctx, center.DefaultApolloMiddlewareService, []string{center.DefaultApolloTraceNamespace})
		if err != nil {
			slog.Errorf(ctx, "%s init apollo config center error, service name: %s, namespaces: %s, err: %s",
				fun, center.DefaultApolloMiddlewareService, center.DefaultApolloTraceNamespace, err.Error())
			return
		}

		apolloCenter.StartWatchUpdate(ctx)
	})
}

type spanFilterConfig struct {
	mu sync.RWMutex

	urls []string
}

func (m *spanFilterConfig) updateUrls(urls []string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.urls = urls
}

func (m *spanFilterConfig) filterUrl(url string) bool {
	m.mu.RLock()
	defer m.mu.RUnlock()

	for _, u := range m.urls {
		if u == url {
			return false
		}
	}

	return true
}

func (m *spanFilterConfig) HandleChangeEvent(event *center.ChangeEvent) {
	fun := "spanFilterConfig.HandleChangeEvent --> "
	ctx := context.Background()

	if event.Namespace != center.DefaultApolloTraceNamespace {
		return
	}

	for key, change := range event.Changes {
		if key == FilterUrls {
			slog.Infof(ctx, "%s get key %s from apollo, old value: %s, new value: %s", fun, key, change.OldValue, change.NewValue)
			urlList := strings.Split(change.NewValue, ListConfigSep)
			m.updateUrls(urlList)
		}
	}
}

func UrlSpanFilter(r *http.Request) bool {
	if apolloSpanFilterConfig != nil {
		return apolloSpanFilterConfig.filterUrl(r.URL.Path)
	}

	return true
}
<commit_msg>Modify initApolloCenter to throw error<commit_after>package trace

import (
	"context"
	"fmt"
	"github.com/shawnfeng/sutil/sconf/center"
	"github.com/shawnfeng/sutil/slog/slog"
	"net/http"
	"strings"
	"sync"
)

const FilterUrls = "span_filter_urls"

const ListConfigSep = ","

var initApolloLock sync.Mutex

var apolloCenter center.ConfigCenter

var apolloSpanFilterConfig *spanFilterConfig

func InitTraceSpanFilter() error {
	fun := "TraceSpanFilter.init --> "
	ctx := context.Background()

	if err := initApolloCenter(ctx); err != nil {
		return err
	}

	urls, ok := apolloCenter.GetStringWithNamespace(ctx, center.DefaultApolloTraceNamespace, FilterUrls)
	if !ok {
		return fmt.Errorf("not get %s from apollo namespace %s", FilterUrls, center.DefaultApolloTraceNamespace)
	}
	slog.Infof(ctx, "%s get %s from apollo: %s", fun, FilterUrls, urls)

	urlList := strings.Split(urls, ListConfigSep)

	apolloSpanFilterConfig = &spanFilterConfig{
		urls: urlList,
	}

	apolloCenter.RegisterObserver(ctx, apolloSpanFilterConfig)
	return nil
}

func initApolloCenter(ctx context.Context) error {
	if apolloCenter != nil {
		return nil
	}

	initApolloLock.Lock()
	defer initApolloLock.Unlock()

	if apolloCenter != nil {
		return nil
	}

	var err error
	apolloCenter, err = center.NewConfigCenter(center.ApolloConfigCenter)
	if err != nil {
		return fmt.Errorf("new config center error, %s", err.Error())
	}

	namespaceList := []string{center.DefaultApolloTraceNamespace}
	err = apolloCenter.Init(ctx, center.DefaultApolloMiddlewareService, namespaceList)
	if err != nil {
		return fmt.Errorf("init apollo with service %s namespace %s error, %s",
			center.DefaultApolloMiddlewareService, strings.Join(namespaceList, " "), err.Error())
	}

	apolloCenter.StartWatchUpdate(ctx)
	return nil
}

type spanFilterConfig struct {
	mu sync.RWMutex

	urls []string
}

func (m *spanFilterConfig) updateUrls(urls []string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.urls = urls
}

func (m *spanFilterConfig) filterUrl(url string) bool {
	m.mu.RLock()
	defer m.mu.RUnlock()

	for _, u := range m.urls {
		if u == url {
			return false
		}
	}

	return true
}

func (m *spanFilterConfig) HandleChangeEvent(event *center.ChangeEvent) {
	fun := "spanFilterConfig.HandleChangeEvent --> "
	ctx := context.Background()

	if event.Namespace != center.DefaultApolloTraceNamespace {
		return
	}

	for key, change := range event.Changes {
		if key == FilterUrls {
			slog.Infof(ctx, "%s get key %s from apollo, old value: %s, new value: %s", fun, key, change.OldValue, change.NewValue)
			urlList := strings.Split(change.NewValue, ListConfigSep)
			m.updateUrls(urlList)
		}
	}
}

func UrlSpanFilter(r *http.Request) bool {
	if apolloSpanFilterConfig != nil {
		return apolloSpanFilterConfig.filterUrl(r.URL.Path)
	}

	return true
}
<|endoftext|>


<commit_before>// Copyright 2014 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.

package tracker

import (
	"net"

	"github.com/chihaya/chihaya/stats"
	"github.com/chihaya/chihaya/tracker/models"
)

// HandleAnnounce encapsulates all of the logic of handling a BitTorrent
// client's Announce without being coupled to any transport protocol.
func (tkr *Tracker) HandleAnnounce(ann *models.Announce, w Writer) error {
	conn, err := tkr.Pool.Get()
	if err != nil {
		return err
	}

	defer conn.Close()

	if tkr.cfg.ClientWhitelistEnabled {
		if err = conn.FindClient(ann.ClientID()); err != nil {
			return err
		}
	}

	var user *models.User
	if tkr.cfg.PrivateEnabled {
		if user, err = conn.FindUser(ann.Passkey); err != nil {
			return err
		}
	}

	torrent, err := conn.FindTorrent(ann.Infohash)

	if err == models.ErrTorrentDNE && !tkr.cfg.PrivateEnabled {
		torrent = &models.Torrent{
			Infohash: ann.Infohash,
			Seeders:  models.PeerMap{},
			Leechers: models.PeerMap{},
		}

		err = conn.PutTorrent(torrent)
		if err != nil {
			return err
		}
		stats.RecordEvent(stats.NewTorrent)
	} else if err != nil {
		return err
	}

	ann.BuildPeer(user, torrent)

	var uploaded, downloaded uint64
	if tkr.cfg.PrivateEnabled {
		uploaded, downloaded = delta(ann)
	}

	created, err := updateSwarm(conn, ann)
	if err != nil {
		return err
	}

	snatched, err := handleEvent(conn, ann)
	if err != nil {
		return err
	}

	if tkr.cfg.PrivateEnabled {
		err = tkr.backend.RecordAnnounce(&models.AnnounceDelta{
			Peer:       ann.Peer,
			Torrent:    ann.Torrent,
			User:       ann.User,
			Created:    created,
			Snatched:   snatched,
			Uploaded:   uploaded,
			Downloaded: downloaded,
		})
		if err != nil {
			return err
		}
	} else if tkr.cfg.PurgeInactiveTorrents && torrent.PeerCount() == 0 {
		// Rather than deleting the torrent explicitly, let the tracker driver
		// ensure there are no race conditions.
		conn.PurgeInactiveTorrent(torrent.Infohash)
		stats.RecordEvent(stats.DeletedTorrent)
	}

	return w.WriteAnnounce(newAnnounceResponse(ann))
}

func delta(a *models.Announce) (uploaded, downloaded uint64) {
	var (
		rawDeltaUp   = a.Peer.Uploaded - a.Uploaded
		rawDeltaDown uint64
	)

	if !a.Config.FreeleechEnabled {
		rawDeltaDown = a.Peer.Downloaded - a.Downloaded
	}

	// Restarting a torrent may cause a delta to be negative.
	if rawDeltaUp < 0 {
		rawDeltaUp = 0
	}

	if rawDeltaDown < 0 {
		rawDeltaDown = 0
	}

	uploaded = uint64(float64(rawDeltaUp) * a.User.UpMultiplier * a.Torrent.UpMultiplier)
	downloaded = uint64(float64(rawDeltaDown) * a.User.DownMultiplier * a.Torrent.DownMultiplier)

	return
}

// updateSwarm handles the changes to a torrent's swarm given an announce.
func updateSwarm(c Conn, ann *models.Announce) (created bool, err error) {
	var createdv4, createdv6 bool
	c.TouchTorrent(ann.Torrent.Infohash)

	if ann.HasIPv4() {
		createdv4, err = updatePeer(c, ann, ann.PeerV4)
		if err != nil {
			return
		}
	}
	if ann.HasIPv6() {
		createdv6, err = updatePeer(c, ann, ann.PeerV6)
		if err != nil {
			return
		}
	}

	return createdv4 || createdv6, nil
}

func updatePeer(c Conn, ann *models.Announce, peer *models.Peer) (created bool, err error) {
	p, t := ann.Peer, ann.Torrent

	switch {
	case t.InSeederPool(p):
		err = c.PutSeeder(t.Infohash, p)
		if err != nil {
			return
		}

	case t.InLeecherPool(p):
		err = c.PutLeecher(t.Infohash, p)
		if err != nil {
			return
		}

	default:
		if ann.Event != "" && ann.Event != "started" {
			err = models.ErrBadRequest
			return
		}

		if ann.Left == 0 {
			err = c.PutSeeder(t.Infohash, p)
			if err != nil {
				return
			}
			stats.RecordPeerEvent(stats.NewSeed, p.HasIPv6())

		} else {
			err = c.PutLeecher(t.Infohash, p)
			if err != nil {
				return
			}
			stats.RecordPeerEvent(stats.NewLeech, p.HasIPv6())
		}
		created = true
	}
	return
}

// handleEvent checks to see whether an announce has an event and if it does,
// properly handles that event.
func handleEvent(c Conn, ann *models.Announce) (snatched bool, err error) {
	var snatchedv4, snatchedv6 bool

	if ann.HasIPv4() {
		snatchedv4, err = handlePeerEvent(c, ann, ann.PeerV4)
		if err != nil {
			return
		}
	}
	if ann.HasIPv6() {
		snatchedv6, err = handlePeerEvent(c, ann, ann.PeerV6)
		if err != nil {
			return
		}
	}

	if snatchedv4 || snatchedv6 {
		err = c.IncrementTorrentSnatches(ann.Torrent.Infohash)
		if err != nil {
			return
		}
		ann.Torrent.Snatches++
		return true, nil
	}
	return false, nil
}

func handlePeerEvent(c Conn, ann *models.Announce, p *models.Peer) (snatched bool, err error) {
	p, t := ann.Peer, ann.Torrent

	switch {
	case ann.Event == "stopped" || ann.Event == "paused":
		// updateSwarm checks if the peer is active on the torrent,
		// so one of these branches must be followed.
		if t.InSeederPool(p) {
			err = c.DeleteSeeder(t.Infohash, p)
			if err != nil {
				return
			}
			stats.RecordPeerEvent(stats.DeletedSeed, p.HasIPv6())

		} else if t.InLeecherPool(p) {
			err = c.DeleteLeecher(t.Infohash, p)
			if err != nil {
				return
			}
			stats.RecordPeerEvent(stats.DeletedLeech, p.HasIPv6())
		}

	case ann.Event == "completed":
		_, v4seed := t.Seeders[models.NewPeerKey(p.ID, false)]
		_, v6seed := t.Seeders[models.NewPeerKey(p.ID, true)]

		if t.InLeecherPool(p) {
			err = leecherFinished(c, t, p)
		} else {
			err = models.ErrBadRequest
		}

		// If one of the dual-stacked peers is already a seeder, they have
		// already snatched.
		if !(v4seed || v6seed) {
			snatched = true
		}

	case t.InLeecherPool(p) && ann.Left == 0:
		// A leecher completed but the event was never received.
		err = leecherFinished(c, t, p)
	}

	return
}

// leecherFinished moves a peer from the leeching pool to the seeder pool.
func leecherFinished(c Conn, t *models.Torrent, p *models.Peer) error {
	if err := c.DeleteLeecher(t.Infohash, p); err != nil {
		return err
	}
	if err := c.PutSeeder(t.Infohash, p); err != nil {
		return err
	}
	stats.RecordPeerEvent(stats.Completed, p.HasIPv6())
	return nil
}

func newAnnounceResponse(ann *models.Announce) *models.AnnounceResponse {
	seedCount := len(ann.Torrent.Seeders)
	leechCount := len(ann.Torrent.Leechers)

	res := &models.AnnounceResponse{
		Complete:    seedCount,
		Incomplete:  leechCount,
		Interval:    ann.Config.Announce.Duration,
		MinInterval: ann.Config.MinAnnounce.Duration,
		Compact:     ann.Compact,
	}

	if ann.NumWant > 0 && ann.Event != "stopped" && ann.Event != "paused" {
		res.IPv4Peers, res.IPv6Peers = getPeers(ann)
	}

	return res
}

// getPeers returns lists IPv4 and IPv6 peers on a given torrent sized according
// to the wanted parameter.
func getPeers(ann *models.Announce) (ipv4s, ipv6s models.PeerList) {
	ipv4s, ipv6s = models.PeerList{}, models.PeerList{}

	if ann.Left == 0 {
		// If they're seeding, give them only leechers.
		return appendPeers(ipv4s, ipv6s, ann, ann.Torrent.Leechers, ann.NumWant)
	}

	// If they're leeching, prioritize giving them seeders.
	ipv4s, ipv6s = appendPeers(ipv4s, ipv6s, ann, ann.Torrent.Seeders, ann.NumWant)
	return appendPeers(ipv4s, ipv6s, ann, ann.Torrent.Leechers, ann.NumWant-len(ipv4s)-len(ipv6s))
}

// appendPeers implements the logic of adding peers to the IPv4 or IPv6 lists.
func appendPeers(ipv4s, ipv6s models.PeerList, ann *models.Announce, peers models.PeerMap, wanted int) (models.PeerList, models.PeerList) {
	if ann.Config.PreferredSubnet {
		return appendSubnetPeers(ipv4s, ipv6s, ann, peers, wanted)
	}

	count := 0

	for _, peer := range peers {
		if count >= wanted {
			break
		} else if peersEquivalent(&peer, ann.Peer) {
			continue
		}

		if ann.HasIPv6() && peer.HasIPv6() {
			ipv6s = append(ipv6s, peer)
			count++
		} else if peer.HasIPv4() {
			ipv4s = append(ipv4s, peer)
			count++
		}
	}

	return ipv4s, ipv6s
}

// appendSubnetPeers is an alternative version of appendPeers used when the
// config variable PreferredSubnet is enabled.
func appendSubnetPeers(ipv4s, ipv6s models.PeerList, ann *models.Announce, peers models.PeerMap, wanted int) (models.PeerList, models.PeerList) {
	var subnetIPv4 net.IPNet
	var subnetIPv6 net.IPNet

	if ann.HasIPv4() {
		subnetIPv4 = net.IPNet{ann.IPv4, net.CIDRMask(ann.Config.PreferredIPv4Subnet, 32)}
	}

	if ann.HasIPv6() {
		subnetIPv6 = net.IPNet{ann.IPv6, net.CIDRMask(ann.Config.PreferredIPv6Subnet, 128)}
	}

	// Iterate over the peers twice: first add only peers in the same subnet and
	// if we still need more peers grab ones that haven't already been added.
	count := 0
	for _, checkInSubnet := range [2]bool{true, false} {
		for _, peer := range peers {
			if count >= wanted {
				break
			}

			inSubnet4 := peer.HasIPv4() && subnetIPv4.Contains(peer.IP)
			inSubnet6 := peer.HasIPv6() && subnetIPv6.Contains(peer.IP)

			if peersEquivalent(&peer, ann.Peer) || checkInSubnet != (inSubnet4 || inSubnet6) {
				continue
			}

			if ann.HasIPv6() && peer.HasIPv6() {
				ipv6s = append(ipv6s, peer)
				count++
			} else if peer.HasIPv4() {
				ipv4s = append(ipv4s, peer)
				count++
			}
		}
	}

	return ipv4s, ipv6s
}

// peersEquivalent checks if two peers represent the same entity.
func peersEquivalent(a, b *models.Peer) bool {
	return a.ID == b.ID || a.UserID != 0 && a.UserID == b.UserID
}
<commit_msg>Calculate delta based off of old peer<commit_after>// Copyright 2014 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.

package tracker

import (
	"net"

	"github.com/chihaya/chihaya/stats"
	"github.com/chihaya/chihaya/tracker/models"
)

// HandleAnnounce encapsulates all of the logic of handling a BitTorrent
// client's Announce without being coupled to any transport protocol.
func (tkr *Tracker) HandleAnnounce(ann *models.Announce, w Writer) error {
	conn, err := tkr.Pool.Get()
	if err != nil {
		return err
	}

	defer conn.Close()

	if tkr.cfg.ClientWhitelistEnabled {
		if err = conn.FindClient(ann.ClientID()); err != nil {
			return err
		}
	}

	var user *models.User
	if tkr.cfg.PrivateEnabled {
		if user, err = conn.FindUser(ann.Passkey); err != nil {
			return err
		}
	}

	torrent, err := conn.FindTorrent(ann.Infohash)

	if err == models.ErrTorrentDNE && !tkr.cfg.PrivateEnabled {
		torrent = &models.Torrent{
			Infohash: ann.Infohash,
			Seeders:  models.PeerMap{},
			Leechers: models.PeerMap{},
		}

		err = conn.PutTorrent(torrent)
		if err != nil {
			return err
		}
		stats.RecordEvent(stats.NewTorrent)
	} else if err != nil {
		return err
	}

	ann.BuildPeer(user, torrent)

	var uploaded, downloaded uint64
	if tkr.cfg.PrivateEnabled {
		uploaded, downloaded = delta(ann, torrent)
	}

	created, err := updateSwarm(conn, ann)
	if err != nil {
		return err
	}

	snatched, err := handleEvent(conn, ann)
	if err != nil {
		return err
	}

	if tkr.cfg.PrivateEnabled {
		err = tkr.backend.RecordAnnounce(&models.AnnounceDelta{
			Peer:       ann.Peer,
			Torrent:    ann.Torrent,
			User:       ann.User,
			Created:    created,
			Snatched:   snatched,
			Uploaded:   uploaded,
			Downloaded: downloaded,
		})
		if err != nil {
			return err
		}
	} else if tkr.cfg.PurgeInactiveTorrents && torrent.PeerCount() == 0 {
		// Rather than deleting the torrent explicitly, let the tracker driver
		// ensure there are no race conditions.
		conn.PurgeInactiveTorrent(torrent.Infohash)
		stats.RecordEvent(stats.DeletedTorrent)
	}

	return w.WriteAnnounce(newAnnounceResponse(ann))
}

func delta(a *models.Announce, t *models.Torrent) (uploaded, downloaded uint64) {
	var oldUp, oldDown uint64
	switch {
	case t.InSeederPool(a.Peer):
		oldPeer := t.Seeders[a.Peer.Key()]
		oldUp = oldPeer.Uploaded
		oldDown = oldPeer.Downloaded
	case t.InLeecherPool(a.Peer):
		oldPeer := t.Leechers[a.Peer.Key()]
		oldUp = oldPeer.Uploaded
		oldDown = oldPeer.Downloaded
	}

	var (
		rawDeltaUp   = a.Peer.Uploaded - oldUp
		rawDeltaDown uint64
	)

	if !a.Config.FreeleechEnabled {
		rawDeltaDown = a.Peer.Downloaded - oldDown
	}

	// Restarting a torrent may cause a delta to be negative.
	if rawDeltaUp < 0 {
		rawDeltaUp = 0
	}

	if rawDeltaDown < 0 {
		rawDeltaDown = 0
	}

	uploaded = uint64(float64(rawDeltaUp) * a.User.UpMultiplier * a.Torrent.UpMultiplier)
	downloaded = uint64(float64(rawDeltaDown) * a.User.DownMultiplier * a.Torrent.DownMultiplier)

	return
}

// updateSwarm handles the changes to a torrent's swarm given an announce.
func updateSwarm(c Conn, ann *models.Announce) (created bool, err error) {
	var createdv4, createdv6 bool
	c.TouchTorrent(ann.Torrent.Infohash)

	if ann.HasIPv4() {
		createdv4, err = updatePeer(c, ann, ann.PeerV4)
		if err != nil {
			return
		}
	}
	if ann.HasIPv6() {
		createdv6, err = updatePeer(c, ann, ann.PeerV6)
		if err != nil {
			return
		}
	}

	return createdv4 || createdv6, nil
}

func updatePeer(c Conn, ann *models.Announce, peer *models.Peer) (created bool, err error) {
	p, t := ann.Peer, ann.Torrent

	switch {
	case t.InSeederPool(p):
		err = c.PutSeeder(t.Infohash, p)
		if err != nil {
			return
		}

	case t.InLeecherPool(p):
		err = c.PutLeecher(t.Infohash, p)
		if err != nil {
			return
		}

	default:
		if ann.Event != "" && ann.Event != "started" {
			err = models.ErrBadRequest
			return
		}

		if ann.Left == 0 {
			err = c.PutSeeder(t.Infohash, p)
			if err != nil {
				return
			}
			stats.RecordPeerEvent(stats.NewSeed, p.HasIPv6())

		} else {
			err = c.PutLeecher(t.Infohash, p)
			if err != nil {
				return
			}
			stats.RecordPeerEvent(stats.NewLeech, p.HasIPv6())
		}
		created = true
	}
	return
}

// handleEvent checks to see whether an announce has an event and if it does,
// properly handles that event.
func handleEvent(c Conn, ann *models.Announce) (snatched bool, err error) {
	var snatchedv4, snatchedv6 bool

	if ann.HasIPv4() {
		snatchedv4, err = handlePeerEvent(c, ann, ann.PeerV4)
		if err != nil {
			return
		}
	}
	if ann.HasIPv6() {
		snatchedv6, err = handlePeerEvent(c, ann, ann.PeerV6)
		if err != nil {
			return
		}
	}

	if snatchedv4 || snatchedv6 {
		err = c.IncrementTorrentSnatches(ann.Torrent.Infohash)
		if err != nil {
			return
		}
		ann.Torrent.Snatches++
		return true, nil
	}
	return false, nil
}

func handlePeerEvent(c Conn, ann *models.Announce, p *models.Peer) (snatched bool, err error) {
	p, t := ann.Peer, ann.Torrent

	switch {
	case ann.Event == "stopped" || ann.Event == "paused":
		// updateSwarm checks if the peer is active on the torrent,
		// so one of these branches must be followed.
		if t.InSeederPool(p) {
			err = c.DeleteSeeder(t.Infohash, p)
			if err != nil {
				return
			}
			stats.RecordPeerEvent(stats.DeletedSeed, p.HasIPv6())

		} else if t.InLeecherPool(p) {
			err = c.DeleteLeecher(t.Infohash, p)
			if err != nil {
				return
			}
			stats.RecordPeerEvent(stats.DeletedLeech, p.HasIPv6())
		}

	case ann.Event == "completed":
		_, v4seed := t.Seeders[models.NewPeerKey(p.ID, false)]
		_, v6seed := t.Seeders[models.NewPeerKey(p.ID, true)]

		if t.InLeecherPool(p) {
			err = leecherFinished(c, t, p)
		} else {
			err = models.ErrBadRequest
		}

		// If one of the dual-stacked peers is already a seeder, they have
		// already snatched.
		if !(v4seed || v6seed) {
			snatched = true
		}

	case t.InLeecherPool(p) && ann.Left == 0:
		// A leecher completed but the event was never received.
		err = leecherFinished(c, t, p)
	}

	return
}

// leecherFinished moves a peer from the leeching pool to the seeder pool.
func leecherFinished(c Conn, t *models.Torrent, p *models.Peer) error {
	if err := c.DeleteLeecher(t.Infohash, p); err != nil {
		return err
	}
	if err := c.PutSeeder(t.Infohash, p); err != nil {
		return err
	}
	stats.RecordPeerEvent(stats.Completed, p.HasIPv6())
	return nil
}

func newAnnounceResponse(ann *models.Announce) *models.AnnounceResponse {
	seedCount := len(ann.Torrent.Seeders)
	leechCount := len(ann.Torrent.Leechers)

	res := &models.AnnounceResponse{
		Complete:    seedCount,
		Incomplete:  leechCount,
		Interval:    ann.Config.Announce.Duration,
		MinInterval: ann.Config.MinAnnounce.Duration,
		Compact:     ann.Compact,
	}

	if ann.NumWant > 0 && ann.Event != "stopped" && ann.Event != "paused" {
		res.IPv4Peers, res.IPv6Peers = getPeers(ann)
	}

	return res
}

// getPeers returns lists IPv4 and IPv6 peers on a given torrent sized according
// to the wanted parameter.
func getPeers(ann *models.Announce) (ipv4s, ipv6s models.PeerList) {
	ipv4s, ipv6s = models.PeerList{}, models.PeerList{}

	if ann.Left == 0 {
		// If they're seeding, give them only leechers.
		return appendPeers(ipv4s, ipv6s, ann, ann.Torrent.Leechers, ann.NumWant)
	}

	// If they're leeching, prioritize giving them seeders.
	ipv4s, ipv6s = appendPeers(ipv4s, ipv6s, ann, ann.Torrent.Seeders, ann.NumWant)
	return appendPeers(ipv4s, ipv6s, ann, ann.Torrent.Leechers, ann.NumWant-len(ipv4s)-len(ipv6s))
}

// appendPeers implements the logic of adding peers to the IPv4 or IPv6 lists.
func appendPeers(ipv4s, ipv6s models.PeerList, ann *models.Announce, peers models.PeerMap, wanted int) (models.PeerList, models.PeerList) {
	if ann.Config.PreferredSubnet {
		return appendSubnetPeers(ipv4s, ipv6s, ann, peers, wanted)
	}

	count := 0

	for _, peer := range peers {
		if count >= wanted {
			break
		} else if peersEquivalent(&peer, ann.Peer) {
			continue
		}

		if ann.HasIPv6() && peer.HasIPv6() {
			ipv6s = append(ipv6s, peer)
			count++
		} else if peer.HasIPv4() {
			ipv4s = append(ipv4s, peer)
			count++
		}
	}

	return ipv4s, ipv6s
}

// appendSubnetPeers is an alternative version of appendPeers used when the
// config variable PreferredSubnet is enabled.
func appendSubnetPeers(ipv4s, ipv6s models.PeerList, ann *models.Announce, peers models.PeerMap, wanted int) (models.PeerList, models.PeerList) {
	var subnetIPv4 net.IPNet
	var subnetIPv6 net.IPNet

	if ann.HasIPv4() {
		subnetIPv4 = net.IPNet{ann.IPv4, net.CIDRMask(ann.Config.PreferredIPv4Subnet, 32)}
	}

	if ann.HasIPv6() {
		subnetIPv6 = net.IPNet{ann.IPv6, net.CIDRMask(ann.Config.PreferredIPv6Subnet, 128)}
	}

	// Iterate over the peers twice: first add only peers in the same subnet and
	// if we still need more peers grab ones that haven't already been added.
	count := 0
	for _, checkInSubnet := range [2]bool{true, false} {
		for _, peer := range peers {
			if count >= wanted {
				break
			}

			inSubnet4 := peer.HasIPv4() && subnetIPv4.Contains(peer.IP)
			inSubnet6 := peer.HasIPv6() && subnetIPv6.Contains(peer.IP)

			if peersEquivalent(&peer, ann.Peer) || checkInSubnet != (inSubnet4 || inSubnet6) {
				continue
			}

			if ann.HasIPv6() && peer.HasIPv6() {
				ipv6s = append(ipv6s, peer)
				count++
			} else if peer.HasIPv4() {
				ipv4s = append(ipv4s, peer)
				count++
			}
		}
	}

	return ipv4s, ipv6s
}

// peersEquivalent checks if two peers represent the same entity.
func peersEquivalent(a, b *models.Peer) bool {
	return a.ID == b.ID || a.UserID != 0 && a.UserID == b.UserID
}
<|endoftext|>


<commit_before>// TODO: IPv4 parsing?
// TODO: use dnsstub library
// TODO: match outbound to inbound queries
package main

import (
	"encoding/binary"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcapgo"
	"github.com/miekg/dns"
	"log"
	"net"
	"os"
	"time"
)

var (
	// We store the configuration of our local resolver in a global
	// variable for convenience.
	resolv_conf *dns.ClientConfig

	// Use a global debug flag.
	debug bool
)

// If we were passed name server addresses, parse them with this function.
func parse_root_server_addresses(addrs []string) map[string]bool {
	if debug {
		fmt.Fprintf(os.Stderr, "pcap2ymmv parse_root_server_addresses()\n")
		fmt.Fprintf(os.Stderr, "pcap2ymmv addrs:%s\n", addrs)
	}
	root_addresses := make(map[string]bool)
	for _, addr := range addrs {
		ip := net.ParseIP(addr)
		if ip == nil {
			log.Fatal("Error parsing address '%s'", addr)
		}
		root_addresses[ip.String()] = true
		if debug {
			fmt.Fprintf(os.Stderr, "pcap2ymmv checking for %s\n", ip.String())
		}
	}
	return root_addresses
}

// We have a goroutine to act as a stub resolver, and use this
// structure to send the question in and get the results out.
type stub_resolve_info struct {
	ownername string
	rtype     uint16
	answer    *dns.Msg
}

// A goroutine which performs stub lookups from a queue, writing
// the results to another queue.
func stub_resolve(questions <-chan stub_resolve_info,
	results chan<- stub_resolve_info) {
	// make a client for our lookups
	dnsClient := new(dns.Client)
	dnsClient.Net = "tcp"
	// read each question on our channel
	for question := range questions {
		// build our answer
		var result stub_resolve_info = question
		result.answer = nil
		// make a DNS query based on our question
		query := new(dns.Msg)
		query.RecursionDesired = true
		query.SetQuestion(question.ownername, question.rtype)
		// check each resolver in turn
		for _, server := range resolv_conf.Servers {
			resolver := server + ":53"
			r, _, err := dnsClient.Exchange(query, resolver)
			// if we got an answer, use that and stop trying
			if (err == nil) && (r != nil) && (r.Rcode == dns.RcodeSuccess) {
				result.answer = r
				break
			}
		}
		// send back our answer (might be nil)
		results <- result
	}
}

// Lookup all of the IP addresses associated with the root name servers.
// Return two maps based on the results found, which have the keys of
// the binary values of the IPv4 and IPv6 addresses. (It's a bit clumsy,
// but it allows us to do quick and easy lookups of the addresses in the
// pcap later.)
func lookup_root_server_addresses() map[string]bool {
	if debug {
		fmt.Fprintf(os.Stderr, "pcap2ymmv lookup_root_server_addresses()\n")
	}
	// look up the NS of the IANA root
	root_client := new(dns.Client)
	root_client.Net = "tcp"
	ns_query := new(dns.Msg)
	ns_query.SetQuestion(".", dns.TypeNS)
	// TODO: avoid hard-coding a particular root server here
	ns_response, _, err := root_client.Exchange(ns_query,
		"f.root-servers.net:53")
	if err != nil {
		log.Fatal("Error looking up root name servers")
	}
	var root_servers []string
	for _, root_server := range ns_response.Answer {
		switch root_server.(type) {
		case *dns.NS:
			ns := root_server.(*dns.NS).Ns
			root_servers = append(root_servers, ns)
		}
	}
	// look up the addresses of the root servers
	questions := make(chan stub_resolve_info, 16)
	results := make(chan stub_resolve_info, len(root_servers)*2)
	for i := 0; i < 8; i++ {
		go stub_resolve(questions, results)
	}
	for _, ns := range root_servers {
		info := new(stub_resolve_info)
		info.ownername = ns
		info.rtype = dns.TypeAAAA
		questions <- *info
		info = new(stub_resolve_info)
		info.ownername = ns
		info.rtype = dns.TypeA
		questions <- *info
	}
	root_addresses := make(map[string]bool)
	for i := 0; i < len(root_servers)*2; i++ {
		response := <-results
		if response.answer == nil {
			log.Fatalf("Error looking up root server %s",
				response.ownername)
		}
		for _, root_address := range response.answer.Answer {
			switch root_address.(type) {
			case *dns.AAAA:
				aaaa_s := root_address.(*dns.AAAA).AAAA.String()
				if debug {
					fmt.Fprintf(os.Stderr, "pcap2ymmv IANA server: %s\n", aaaa_s)
				}
				root_addresses[aaaa_s] = true
			case *dns.A:
				a_s := root_address.(*dns.A).A.String()
				if debug {
					fmt.Fprintf(os.Stderr, "pcap2ymmv IANA server: %s\n", a_s)
				}
				root_addresses[a_s] = true
			}
		}
	}
	close(questions)
	return root_addresses
}

func ymmv_write(ip_family int, addr net.IP, query dns.Msg,
	answer_time time.Time, answer dns.Msg) {
	// output magic value
	_, err := os.Stdout.Write([]byte("ymmv"))
	if err != nil {
		log.Fatal(err)
	}

	// output address family
	if ip_family == 4 {
		_, err = os.Stdout.Write([]byte("4"))
	} else if ip_family == 6 {
		_, err = os.Stdout.Write([]byte("6"))
	} else {
		log.Fatalf("Unknown ip_family %d\n", ip_family)
	}
	if err != nil {
		log.Fatal(err)
	}

	// output (U)DP or (T)CP
	_, err = os.Stdout.Write([]byte("u")) // only support UDP for now...
	if err != nil {
		log.Fatal(err)
	}

	// output actual address
	_, err = os.Stdout.Write(addr)
	if err != nil {
		log.Fatal(err)
	}

	// output when the query happened (we don't know, so use 0)
	_, err = os.Stdout.Write([]byte{0, 0, 0, 0, 0, 0, 0, 0})

	// write the byte count of our query
	query_bytes, err := query.Pack()
	if err != nil {
		log.Fatal(err)
	}
	query_len := uint16(len(query_bytes))
	err = binary.Write(os.Stdout, binary.BigEndian, query_len)
	if err != nil {
		log.Fatal(err)
	}

	// write our query
	_, err = os.Stdout.Write(query_bytes)
	if err != nil {
		log.Fatal(err)
	}

	// output when the answer arrived
	seconds := uint32(answer_time.Unix())
	err = binary.Write(os.Stdout, binary.BigEndian, seconds)
	if err != nil {
		log.Fatal(err)
	}
	nanoseconds := uint32(answer_time.Nanosecond())
	err = binary.Write(os.Stdout, binary.BigEndian, nanoseconds)
	if err != nil {
		log.Fatal(err)
	}

	// write the byte count of our answer
	answer_bytes, err := answer.Pack()
	if err != nil {
		log.Fatal(err)
	}
	answer_len := uint16(len(answer_bytes))
	err = binary.Write(os.Stdout, binary.BigEndian, answer_len)
	if err != nil {
		log.Fatal(err)
	}

	// write our answer
	_, err = os.Stdout.Write(answer_bytes)
	if err != nil {
		log.Fatal(err)
	}
	os.Stdout.Sync()

	if debug {
		fmt.Fprintf(os.Stderr, "pcap2ymmv wrote ymmv record of %d bytes\n", len(answer_bytes))
	}
}

func parse_query(raw_answer []byte) (*dns.Msg, *dns.Msg, error) {
	// parse the query
	answer := new(dns.Msg)
	err := answer.Unpack(raw_answer)
	if err != nil {
		return nil, nil, err
	}
	answer.Id = 0
	// infer the answer and build that
	query := answer.Copy()
	query.Response = false
	query.Authoritative = false
	query.Truncated = false
	query.AuthenticatedData = true
	query.CheckingDisabled = true
	query.Rcode = 0
	query.Answer = nil
	query.Ns = nil
	query.Extra = nil
	return query, answer, nil
}

// Look in the named file and find any packets that are from our root
// servers on port 53.
func pcap2ymmv(fname string, root_addresses map[string]bool) {
	// open our pcap file
	var file *os.File
	if fname == "-" {
		file = os.Stdin
	} else {
		named_file, err := os.Open(fname)
		if err != nil {
			log.Fatal(err)
		}
		file = named_file
	}
	pcap_file, err := pcapgo.NewReader(file)
	if err != nil {
		log.Fatal(err)
	}

	for {
		// reach each packet
		pkt_bytes, ci, err := pcap_file.ReadPacketData()
		if err != nil {
			fmt.Fprintf(os.Stderr, "pcap2ymmv error reading packet; %s\n", err)
			break
		}
		if debug {
			fmt.Fprintf(os.Stderr, "pcap2ymmv read packet (len:%d)\n", len(pkt_bytes))
		}

		// check for match against IP addresses that we care about
		ip_match := false
		var ip_family int
		var ip_addr net.IP
		var udp *layers.UDP
		ipv6packet := gopacket.NewPacket(pkt_bytes, layers.LayerTypeIPv6, gopacket.Default)
		var ipv6 *layers.IPv6
		if ipv6packet != nil {
			ipv6, _ = ipv6packet.Layer(layers.LayerTypeIPv6).(*layers.IPv6)
			//			ipv6, _ = ipv6Layer.(*layers.IPv6)
		}
		if ipv6 != nil {
			if debug {
				fmt.Fprintf(os.Stderr, "pcap2ymmv IPv6 %s\n", ipv6.SrcIP.String())
			}
			ip_family = 6
			ip_addr = ipv6.SrcIP
			_, ip_match = root_addresses[ipv6.SrcIP.String()]
			if ip_match {
				udp, _ = ipv6packet.Layer(layers.LayerTypeUDP).(*layers.UDP)
			}
		} else {
			ipv4packet := gopacket.NewPacket(pkt_bytes, layers.LayerTypeIPv4, gopacket.Default)
			var ipv4 *layers.IPv4
			if ipv4packet != nil {
				ipv4, _ = ipv4packet.Layer(layers.LayerTypeIPv4).(*layers.IPv4)
				//				ipv4, _ = ipv4Layer.(*layers.IPv4)
			}
			if ipv4 != nil {
				if debug {
					fmt.Fprintf(os.Stderr, "pcap2ymmv IPv4 %s\n", ipv4.SrcIP.String())
				}
				ip_family = 4
				ip_addr = ipv4.SrcIP
				_, ip_match = root_addresses[ipv4.SrcIP.String()]
				if ip_match {
					udp, _ = ipv4packet.Layer(layers.LayerTypeUDP).(*layers.UDP)
				}
			}
		}
		if debug {
			fmt.Fprintf(os.Stderr, "pcap2ymmv IP match %t\n", ip_match)
		}
		if !ip_match {
			continue
		}

		// we only want port 53
		if debug {
			fmt.Fprintf(os.Stderr, "pcap2ymmv UDP port: %d\n", udp.SrcPort)
		}
		if udp.SrcPort != 53 {
			continue
		}

		// if we got a valid IP and UDP packet, process it
		if debug {
			fmt.Fprintf(os.Stderr, "pcap2ymmv matched packet\n")
		}

		// parse the payload as the DNS message
		query, answer, err := parse_query(udp.Payload)
		if err != nil {
			fmt.Fprintf(os.Stderr, "pcap2ymmv error unpacking DNS message: %s\n", err)
		} else {
			ymmv_write(ip_family, ip_addr, *query, ci.Timestamp, *answer)
			os.Stdout.Sync()
		}
	}
	file.Close()
}

// Main function.
func main() {
	// turn on debugging if desired
	if (len(os.Args) > 1) && (os.Args[1] == "-d") {
		debug = true
		os.Args = append(os.Args[:1], os.Args[2:]...)
	}

	// initialize our stub resolver
	var (
		err error
	)
	resolv_conf, err = dns.ClientConfigFromFile("/etc/resolv.conf")
	if err != nil {
		log.Fatal(err)
	}

	// get root server addresses
	var root_addresses map[string]bool
	if len(os.Args) > 1 {
		root_addresses = parse_root_server_addresses(os.Args[1:])
	} else {
		root_addresses = lookup_root_server_addresses()
	}

	// process stdin as a pcap file
	pcap2ymmv("-", root_addresses)
}
<commit_msg>Handle case where UDP can't parse (perhaps for an IP fragment)<commit_after>// TODO: IPv4 parsing?
// TODO: use dnsstub library
// TODO: match outbound to inbound queries
package main

import (
	"encoding/binary"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcapgo"
	"github.com/miekg/dns"
	"log"
	"net"
	"os"
	"time"
)

var (
	// We store the configuration of our local resolver in a global
	// variable for convenience.
	resolv_conf *dns.ClientConfig

	// Use a global debug flag.
	debug bool
)

// If we were passed name server addresses, parse them with this function.
func parse_root_server_addresses(addrs []string) map[string]bool {
	if debug {
		fmt.Fprintf(os.Stderr, "pcap2ymmv parse_root_server_addresses()\n")
		fmt.Fprintf(os.Stderr, "pcap2ymmv addrs:%s\n", addrs)
	}
	root_addresses := make(map[string]bool)
	for _, addr := range addrs {
		ip := net.ParseIP(addr)
		if ip == nil {
			log.Fatal("Error parsing address '%s'", addr)
		}
		root_addresses[ip.String()] = true
		if debug {
			fmt.Fprintf(os.Stderr, "pcap2ymmv checking for %s\n", ip.String())
		}
	}
	return root_addresses
}

// We have a goroutine to act as a stub resolver, and use this
// structure to send the question in and get the results out.
type stub_resolve_info struct {
	ownername string
	rtype     uint16
	answer    *dns.Msg
}

// A goroutine which performs stub lookups from a queue, writing
// the results to another queue.
func stub_resolve(questions <-chan stub_resolve_info,
	results chan<- stub_resolve_info) {
	// make a client for our lookups
	dnsClient := new(dns.Client)
	dnsClient.Net = "tcp"
	// read each question on our channel
	for question := range questions {
		// build our answer
		var result stub_resolve_info = question
		result.answer = nil
		// make a DNS query based on our question
		query := new(dns.Msg)
		query.RecursionDesired = true
		query.SetQuestion(question.ownername, question.rtype)
		// check each resolver in turn
		for _, server := range resolv_conf.Servers {
			resolver := server + ":53"
			r, _, err := dnsClient.Exchange(query, resolver)
			// if we got an answer, use that and stop trying
			if (err == nil) && (r != nil) && (r.Rcode == dns.RcodeSuccess) {
				result.answer = r
				break
			}
		}
		// send back our answer (might be nil)
		results <- result
	}
}

// Lookup all of the IP addresses associated with the root name servers.
// Return two maps based on the results found, which have the keys of
// the binary values of the IPv4 and IPv6 addresses. (It's a bit clumsy,
// but it allows us to do quick and easy lookups of the addresses in the
// pcap later.)
func lookup_root_server_addresses() map[string]bool {
	if debug {
		fmt.Fprintf(os.Stderr, "pcap2ymmv lookup_root_server_addresses()\n")
	}
	// look up the NS of the IANA root
	root_client := new(dns.Client)
	root_client.Net = "tcp"
	ns_query := new(dns.Msg)
	ns_query.SetQuestion(".", dns.TypeNS)
	// TODO: avoid hard-coding a particular root server here
	ns_response, _, err := root_client.Exchange(ns_query,
		"f.root-servers.net:53")
	if err != nil {
		log.Fatal("Error looking up root name servers")
	}
	var root_servers []string
	for _, root_server := range ns_response.Answer {
		switch root_server.(type) {
		case *dns.NS:
			ns := root_server.(*dns.NS).Ns
			root_servers = append(root_servers, ns)
		}
	}
	// look up the addresses of the root servers
	questions := make(chan stub_resolve_info, 16)
	results := make(chan stub_resolve_info, len(root_servers)*2)
	for i := 0; i < 8; i++ {
		go stub_resolve(questions, results)
	}
	for _, ns := range root_servers {
		info := new(stub_resolve_info)
		info.ownername = ns
		info.rtype = dns.TypeAAAA
		questions <- *info
		info = new(stub_resolve_info)
		info.ownername = ns
		info.rtype = dns.TypeA
		questions <- *info
	}
	root_addresses := make(map[string]bool)
	for i := 0; i < len(root_servers)*2; i++ {
		response := <-results
		if response.answer == nil {
			log.Fatalf("Error looking up root server %s",
				response.ownername)
		}
		for _, root_address := range response.answer.Answer {
			switch root_address.(type) {
			case *dns.AAAA:
				aaaa_s := root_address.(*dns.AAAA).AAAA.String()
				if debug {
					fmt.Fprintf(os.Stderr, "pcap2ymmv IANA server: %s\n", aaaa_s)
				}
				root_addresses[aaaa_s] = true
			case *dns.A:
				a_s := root_address.(*dns.A).A.String()
				if debug {
					fmt.Fprintf(os.Stderr, "pcap2ymmv IANA server: %s\n", a_s)
				}
				root_addresses[a_s] = true
			}
		}
	}
	close(questions)
	return root_addresses
}

func ymmv_write(ip_family int, addr net.IP, query dns.Msg,
	answer_time time.Time, answer dns.Msg) {
	// output magic value
	_, err := os.Stdout.Write([]byte("ymmv"))
	if err != nil {
		log.Fatal(err)
	}

	// output address family
	if ip_family == 4 {
		_, err = os.Stdout.Write([]byte("4"))
	} else if ip_family == 6 {
		_, err = os.Stdout.Write([]byte("6"))
	} else {
		log.Fatalf("Unknown ip_family %d\n", ip_family)
	}
	if err != nil {
		log.Fatal(err)
	}

	// output (U)DP or (T)CP
	_, err = os.Stdout.Write([]byte("u")) // only support UDP for now...
	if err != nil {
		log.Fatal(err)
	}

	// output actual address
	_, err = os.Stdout.Write(addr)
	if err != nil {
		log.Fatal(err)
	}

	// output when the query happened (we don't know, so use 0)
	_, err = os.Stdout.Write([]byte{0, 0, 0, 0, 0, 0, 0, 0})

	// write the byte count of our query
	query_bytes, err := query.Pack()
	if err != nil {
		log.Fatal(err)
	}
	query_len := uint16(len(query_bytes))
	err = binary.Write(os.Stdout, binary.BigEndian, query_len)
	if err != nil {
		log.Fatal(err)
	}

	// write our query
	_, err = os.Stdout.Write(query_bytes)
	if err != nil {
		log.Fatal(err)
	}

	// output when the answer arrived
	seconds := uint32(answer_time.Unix())
	err = binary.Write(os.Stdout, binary.BigEndian, seconds)
	if err != nil {
		log.Fatal(err)
	}
	nanoseconds := uint32(answer_time.Nanosecond())
	err = binary.Write(os.Stdout, binary.BigEndian, nanoseconds)
	if err != nil {
		log.Fatal(err)
	}

	// write the byte count of our answer
	answer_bytes, err := answer.Pack()
	if err != nil {
		log.Fatal(err)
	}
	answer_len := uint16(len(answer_bytes))
	err = binary.Write(os.Stdout, binary.BigEndian, answer_len)
	if err != nil {
		log.Fatal(err)
	}

	// write our answer
	_, err = os.Stdout.Write(answer_bytes)
	if err != nil {
		log.Fatal(err)
	}
	os.Stdout.Sync()

	if debug {
		fmt.Fprintf(os.Stderr, "pcap2ymmv wrote ymmv record of %d bytes\n", len(answer_bytes))
	}
}

func parse_query(raw_answer []byte) (*dns.Msg, *dns.Msg, error) {
	// parse the query
	answer := new(dns.Msg)
	err := answer.Unpack(raw_answer)
	if err != nil {
		return nil, nil, err
	}
	answer.Id = 0
	// infer the answer and build that
	query := answer.Copy()
	query.Response = false
	query.Authoritative = false
	query.Truncated = false
	query.AuthenticatedData = true
	query.CheckingDisabled = true
	query.Rcode = 0
	query.Answer = nil
	query.Ns = nil
	query.Extra = nil
	return query, answer, nil
}

// Look in the named file and find any packets that are from our root
// servers on port 53.
func pcap2ymmv(fname string, root_addresses map[string]bool) {
	// open our pcap file
	var file *os.File
	if fname == "-" {
		file = os.Stdin
	} else {
		named_file, err := os.Open(fname)
		if err != nil {
			log.Fatal(err)
		}
		file = named_file
	}
	pcap_file, err := pcapgo.NewReader(file)
	if err != nil {
		log.Fatal(err)
	}

	for {
		// reach each packet
		pkt_bytes, ci, err := pcap_file.ReadPacketData()
		if err != nil {
			fmt.Fprintf(os.Stderr, "pcap2ymmv error reading packet; %s\n", err)
			break
		}
		if debug {
			fmt.Fprintf(os.Stderr, "pcap2ymmv read packet (len:%d)\n", len(pkt_bytes))
		}

		// check for match against IP addresses that we care about
		ip_match := false
		var ip_family int
		var ip_addr net.IP
		var udp *layers.UDP
		ipv6packet := gopacket.NewPacket(pkt_bytes, layers.LayerTypeIPv6, gopacket.Default)
		var ipv6 *layers.IPv6
		if ipv6packet != nil {
			ipv6, _ = ipv6packet.Layer(layers.LayerTypeIPv6).(*layers.IPv6)
			//			ipv6, _ = ipv6Layer.(*layers.IPv6)
		}
		if ipv6 != nil {
			if debug {
				fmt.Fprintf(os.Stderr, "pcap2ymmv IPv6 %s\n", ipv6.SrcIP.String())
			}
			ip_family = 6
			ip_addr = ipv6.SrcIP
			_, ip_match = root_addresses[ipv6.SrcIP.String()]
			if ip_match {
				udp, _ = ipv6packet.Layer(layers.LayerTypeUDP).(*layers.UDP)
			}
		} else {
			ipv4packet := gopacket.NewPacket(pkt_bytes, layers.LayerTypeIPv4, gopacket.Default)
			var ipv4 *layers.IPv4
			if ipv4packet != nil {
				ipv4, _ = ipv4packet.Layer(layers.LayerTypeIPv4).(*layers.IPv4)
				//				ipv4, _ = ipv4Layer.(*layers.IPv4)
			}
			if ipv4 != nil {
				if debug {
					fmt.Fprintf(os.Stderr, "pcap2ymmv IPv4 %s\n", ipv4.SrcIP.String())
				}
				ip_family = 4
				ip_addr = ipv4.SrcIP
				_, ip_match = root_addresses[ipv4.SrcIP.String()]
				if ip_match {
					udp, _ = ipv4packet.Layer(layers.LayerTypeUDP).(*layers.UDP)
				}
			}
		}
		if debug {
			fmt.Fprintf(os.Stderr, "pcap2ymmv IP match %t\n", ip_match)
		}
		if !ip_match {
			continue
		}

		// if we were unable to parse the UDP for some reason, report it
		if udp == nil {
			fmt.Fprintf(os.Stderr, "pcap2ymmv unable to parse UDP in packet from %s\n", ip_addr.String())
			continue
		}

		// we only want port 53
		if debug {
			fmt.Fprintf(os.Stderr, "pcap2ymmv UDP port: %d\n", udp.SrcPort)
		}
		if udp.SrcPort != 53 {
			continue
		}

		// if we got a valid IP and UDP packet, process it
		if debug {
			fmt.Fprintf(os.Stderr, "pcap2ymmv matched packet\n")
		}

		// parse the payload as the DNS message
		query, answer, err := parse_query(udp.Payload)
		if err != nil {
			fmt.Fprintf(os.Stderr, "pcap2ymmv error unpacking DNS message: %s\n", err)
		} else {
			ymmv_write(ip_family, ip_addr, *query, ci.Timestamp, *answer)
			os.Stdout.Sync()
		}
	}
	file.Close()
}

// Main function.
func main() {
	// turn on debugging if desired
	if (len(os.Args) > 1) && (os.Args[1] == "-d") {
		debug = true
		os.Args = append(os.Args[:1], os.Args[2:]...)
	}

	// initialize our stub resolver
	var (
		err error
	)
	resolv_conf, err = dns.ClientConfigFromFile("/etc/resolv.conf")
	if err != nil {
		log.Fatal(err)
	}

	// get root server addresses
	var root_addresses map[string]bool
	if len(os.Args) > 1 {
		root_addresses = parse_root_server_addresses(os.Args[1:])
	} else {
		root_addresses = lookup_root_server_addresses()
	}

	// process stdin as a pcap file
	pcap2ymmv("-", root_addresses)
}
<|endoftext|>


<commit_before>// +build integration

package riak

import (
	"fmt"
	"testing"
	"time"
)

func init() {
	integrationTestsBuildCluster()
}

// FetchIndex
// StoreIndex

func TestStoreFetchAndDeleteAYokozunaIndex(t *testing.T) {
	var err error
	var cmd Command
	indexName := "indexName"
	sbuilder := NewStoreIndexCommandBuilder()
	cmd, err = sbuilder.WithIndexName(indexName).WithTimeout(time.Second * 10).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}
	if scmd, ok := cmd.(*StoreIndexCommand); ok {
		if expected, actual := true, scmd.Response; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.FailNow()
	}

	fbuilder := NewFetchIndexCommandBuilder()
	cmd, err = fbuilder.WithIndexName(indexName).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}
	if fcmd, ok := cmd.(*FetchIndexCommand); ok {
		if fcmd.Response == nil {
			t.Errorf("expected non-nil Response")
		}
		idx := fcmd.Response[0]
		if expected, actual := indexName, idx.Name; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := "_yz_default", idx.Schema; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := uint32(3), idx.NVal; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.FailNow()
	}

	dbuilder := NewDeleteIndexCommandBuilder()
	cmd, err = dbuilder.WithIndexName(indexName).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}
	if dcmd, ok := cmd.(*DeleteIndexCommand); ok {
		if expected, actual := true, dcmd.Response; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.FailNow()
	}
}

// FetchSchema
// StoreSchema

func TestStoreFetchAndDeleteAYokozunaSchema(t *testing.T) {
	var err error
	var cmd Command
	defaultSchemaName := "_yz_default"
	schemaName := "schemaName"
	schemaXml := "dummy"

	fbuilder := NewFetchSchemaCommandBuilder()
	cmd, err = fbuilder.WithSchemaName(defaultSchemaName).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}
	if fcmd, ok := cmd.(*FetchSchemaCommand); ok {
		if fcmd.Response == nil {
			t.Errorf("expected non-nil Response")
		}
		sch := fcmd.Response
		if expected, actual := defaultSchemaName, sch.Name; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		schemaXml = sch.Content
	} else {
		t.FailNow()
	}

	sbuilder := NewStoreSchemaCommandBuilder()
	cmd, err = sbuilder.WithSchemaName(schemaName).WithSchema(schemaXml).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}
	if scmd, ok := cmd.(*StoreSchemaCommand); ok {
		if expected, actual := true, scmd.Response; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.FailNow()
	}
}

// Search

func TestSearchViaYokozunaIndex(t *testing.T) {
	var err error
	var cmd Command
	indexName := "myIndex"

	b1 := NewStoreIndexCommandBuilder()
	cmd, err = b1.WithIndexName(indexName).WithTimeout(time.Second * 20).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}
	if scmd, ok := cmd.(*StoreIndexCommand); ok {
		if expected, actual := true, scmd.Response; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.FailNow()
	}

	searchBucket := fmt.Sprintf("%s_search", testBucketName)
	b2 := NewStoreBucketPropsCommandBuilder()
	cmd, err = b2.WithBucket(searchBucket).WithSearchIndex(indexName).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}

	stuffToStore := [...]string{
		"{ \"content_s\":\"Alice was beginning to get very tired of sitting by her sister on the bank, and of having nothing to do: once or twice she had peeped into the book her sister was reading, but it had no pictures or conversations in it, 'and what is the use of a book,' thought Alice 'without pictures or conversation?'\"}",
		"{ \"content_s\":\"So she was considering in her own mind (as well as she could, for the hot day made her feel very sleepy and stupid), whether the pleasure of making a daisy-chain would be worth the trouble of getting up and picking the daisies, when suddenly a White Rabbit with pink eyes ran close by her.\", \"multi_ss\":[\"this\",\"that\"]}",
		"{ \"content_s\":\"The rabbit-hole went straight on like a tunnel for some way, and then dipped suddenly down, so suddenly that Alice had not a moment to think about stopping herself before she found herself falling down a very deep well.\"}",
	}

	for _, s := range stuffToStore {
		b3 := NewStoreValueCommandBuilder()
		obj := &Object{
			ContentType: "application/json", // NB: *very* important for extractor
			Value:       []byte(s),
		}
		cmd, err = b3.WithBucket(searchBucket).WithContent(obj).Build()
		if err != nil {
			t.Fatal(err.Error())
		}
		if err = cluster.Execute(cmd); err != nil {
			t.Fatal(err.Error())
		}
	}

	for count := 0; count < 10; count++ {
		time.Sleep(time.Millisecond * 500)
		b4 := NewSearchCommandBuilder()
		cmd, err = b4.WithIndexName(indexName).WithQuery("multi_ss:t*").Build()
		if err != nil {
			t.Fatal(err.Error())
		}
		if err = cluster.Execute(cmd); err != nil {
			t.Fatal(err.Error())
		}
		if scmd, ok := cmd.(*SearchCommand); ok {
			resp := scmd.Response
			if expected, actual := 1, len(resp.Docs); expected != actual {
				if count < 10 {
					t.Logf("expected %v, got %v - RETRYING", expected, actual)
				} else {
					t.Errorf("expected %v, got %v - DONE", expected, actual)
				}
			}
		} else {
			t.FailNow()
		}
	}
}
<commit_msg>Fix integration test<commit_after>// +build integration

package riak

import (
	"fmt"
	"testing"
	"time"
)

func init() {
	integrationTestsBuildCluster()
}

// FetchIndex
// StoreIndex

func TestStoreFetchAndDeleteAYokozunaIndex(t *testing.T) {
	var err error
	var cmd Command
	indexName := "indexName"
	sbuilder := NewStoreIndexCommandBuilder()
	cmd, err = sbuilder.WithIndexName(indexName).WithTimeout(time.Second * 30).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}
	if scmd, ok := cmd.(*StoreIndexCommand); ok {
		if expected, actual := true, scmd.Response; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.FailNow()
	}

	fbuilder := NewFetchIndexCommandBuilder()
	cmd, err = fbuilder.WithIndexName(indexName).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}
	if fcmd, ok := cmd.(*FetchIndexCommand); ok {
		if fcmd.Response == nil {
			t.Errorf("expected non-nil Response")
		}
		idx := fcmd.Response[0]
		if expected, actual := indexName, idx.Name; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := "_yz_default", idx.Schema; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
		if expected, actual := uint32(3), idx.NVal; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.FailNow()
	}

	dbuilder := NewDeleteIndexCommandBuilder()
	cmd, err = dbuilder.WithIndexName(indexName).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}
	if dcmd, ok := cmd.(*DeleteIndexCommand); ok {
		if expected, actual := true, dcmd.Response; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.FailNow()
	}
}

// FetchSchema
// StoreSchema

func TestStoreFetchAndDeleteAYokozunaSchema(t *testing.T) {
	var err error
	var cmd Command
	defaultSchemaName := "_yz_default"
	schemaName := "schemaName"
	schemaXml := "dummy"

	fbuilder := NewFetchSchemaCommandBuilder()
	cmd, err = fbuilder.WithSchemaName(defaultSchemaName).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err)
	}

	fcmd := cmd.(*FetchSchemaCommand)
	if fcmd.Response == nil {
		t.Fatal("expected non-nil Response")
	}

	sch := fcmd.Response
	if expected, actual := defaultSchemaName, sch.Name; expected != actual {
		t.Errorf("expected %v, got %v", expected, actual)
	}

	schemaXml = sch.Content

	sbuilder := NewStoreSchemaCommandBuilder()
	cmd, err = sbuilder.WithSchemaName(schemaName).WithSchema(schemaXml).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}
	if scmd, ok := cmd.(*StoreSchemaCommand); ok {
		if expected, actual := true, scmd.Response; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.FailNow()
	}
}

// Search

func TestSearchViaYokozunaIndex(t *testing.T) {
	var err error
	var cmd Command
	indexName := "myIndex"

	b1 := NewStoreIndexCommandBuilder()
	cmd, err = b1.WithIndexName(indexName).WithTimeout(time.Second * 20).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}
	if scmd, ok := cmd.(*StoreIndexCommand); ok {
		if expected, actual := true, scmd.Response; expected != actual {
			t.Errorf("expected %v, got %v", expected, actual)
		}
	} else {
		t.FailNow()
	}

	searchBucket := fmt.Sprintf("%s_search", testBucketName)
	b2 := NewStoreBucketPropsCommandBuilder()
	cmd, err = b2.WithBucket(searchBucket).WithSearchIndex(indexName).Build()
	if err != nil {
		t.Fatal(err.Error())
	}
	if err = cluster.Execute(cmd); err != nil {
		t.Fatal(err.Error())
	}

	stuffToStore := [...]string{
		"{ \"content_s\":\"Alice was beginning to get very tired of sitting by her sister on the bank, and of having nothing to do: once or twice she had peeped into the book her sister was reading, but it had no pictures or conversations in it, 'and what is the use of a book,' thought Alice 'without pictures or conversation?'\"}",
		"{ \"content_s\":\"So she was considering in her own mind (as well as she could, for the hot day made her feel very sleepy and stupid), whether the pleasure of making a daisy-chain would be worth the trouble of getting up and picking the daisies, when suddenly a White Rabbit with pink eyes ran close by her.\", \"multi_ss\":[\"this\",\"that\"]}",
		"{ \"content_s\":\"The rabbit-hole went straight on like a tunnel for some way, and then dipped suddenly down, so suddenly that Alice had not a moment to think about stopping herself before she found herself falling down a very deep well.\"}",
	}

	for _, s := range stuffToStore {
		b3 := NewStoreValueCommandBuilder()
		obj := &Object{
			ContentType: "application/json", // NB: *very* important for extractor
			Value:       []byte(s),
		}
		cmd, err = b3.WithBucket(searchBucket).WithContent(obj).Build()
		if err != nil {
			t.Fatal(err.Error())
		}
		if err = cluster.Execute(cmd); err != nil {
			t.Fatal(err.Error())
		}
	}

	for count := 0; count < 10; count++ {
		time.Sleep(time.Millisecond * 500)
		b4 := NewSearchCommandBuilder()
		cmd, err = b4.WithIndexName(indexName).WithQuery("multi_ss:t*").Build()
		if err != nil {
			t.Fatal(err.Error())
		}
		if err = cluster.Execute(cmd); err != nil {
			t.Fatal(err.Error())
		}
		if scmd, ok := cmd.(*SearchCommand); ok {
			resp := scmd.Response
			if expected, actual := 1, len(resp.Docs); expected != actual {
				if count < 10 {
					t.Logf("expected %v, got %v - RETRYING", expected, actual)
				} else {
					t.Errorf("expected %v, got %v - DONE", expected, actual)
				}
			}
		} else {
			t.FailNow()
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"io/ioutil"
	"time"
)

const NCONCURRENTMAX = 20
const NREQS = 5
const DSTURL = "http://www.aftonbladet.se/"

type httpStats struct {
	nsucceded int
	nfailed   int
}

func sendHttpRequests(url string, nreqs int, statsChan chan httpStats) {
	var stats httpStats
	for i := 0; i < nreqs; i++ {
		resp, err := http.Get(url)
		if err == nil {
			if resp.StatusCode == 200 {
				stats.nsucceded += 1
			} else {
				stats.nfailed += 1
			}

			// read entire response
			ioutil.ReadAll(resp.Body)
			resp.Body.Close()
		} else {
			stats.nfailed += 1
		}
	}

	statsChan <- stats
}

func dispatchHttpRequesters(nreqs, nconcurr int, dsturl string, stats *httpStats) {
	requesterChan := make(chan httpStats)
	for i := 0; i < nconcurr; i++ {
		go sendHttpRequests(dsturl, nreqs, requesterChan)
	}

	for i := 0; i < nconcurr; i++ {
		rstat := <-requesterChan
		if stats != nil {
			stats.nsucceded += rstat.nsucceded
			stats.nfailed += rstat.nfailed
		}
	}
}

func main() {
	var stats httpStats

	// set up connection cache, &c
	dispatchHttpRequesters(NREQS, 1, DSTURL, nil)

	for i := 1; i < NCONCURRENTMAX; i++ {
		stats.nsucceded = 0
		stats.nfailed = 0
		startTime := time.Now()
		dispatchHttpRequesters(NREQS, i, DSTURL, &stats)
		duration := time.Since(startTime)
		avg := duration.Nanoseconds() / (int64(stats.nsucceded) + int64(stats.nfailed)) / 1000000;
		fmt.Printf("nconcurrent: %v nreqs: %v succeded: %v failed: %v duration: %v avg: %vms\n",
			i, NREQS, stats.nsucceded, stats.nfailed, duration, avg)
		time.Sleep(5 * time.Second)
	}
}
<commit_msg>Removed weblat.go<commit_after><|endoftext|>


<commit_before>package storebench

import (
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"testing"

	"github.com/skyec/astore"
	blobs "github.com/skyec/astore/testing"
)

var benchStore astore.WriteableStore
var benchDir string

func init() {
}
func BenchmarkDefaultWrite(b *testing.B) {
	benchDir, err := ioutil.TempDir("", "astore-benchmarking-")
	if err != nil {
		log.Fatal(err)
	}

	benchStore, err := astore.NewReadWriteableStore(benchDir)
	if err != nil {
		log.Fatal(err)

	}

	err = benchStore.Initialize()
	if err != nil {
		log.Fatal("Failed to initialze the store:", err)
	}

	blobs := blobs.GenerateBlobs(b.N)
	b.ResetTimer()
	for i := 0; i < len(blobs); i++ {
		err := benchStore.WriteToKey(strconv.Itoa(i%10), blobs[i])
		if err != nil {
			b.Fatal("Failed to write to store:", err)
		}
	}
	b.StopTimer()
	os.RemoveAll(benchDir)

}
<commit_msg>whitespace edit<commit_after>package storebench

import (
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"testing"

	"github.com/skyec/astore"
	blobs "github.com/skyec/astore/testing"
)

var benchStore astore.WriteableStore
var benchDir string

func init() {
}
func BenchmarkDefaultWrite(b *testing.B) {
	benchDir, err := ioutil.TempDir("", "astore-benchmarking-")
	if err != nil {
		log.Fatal(err)
	}

	benchStore, err := astore.NewReadWriteableStore(benchDir)
	if err != nil {
		log.Fatal(err)
	}

	err = benchStore.Initialize()
	if err != nil {
		log.Fatal("Failed to initialze the store:", err)
	}

	blobs := blobs.GenerateBlobs(b.N)
	b.ResetTimer()
	for i := 0; i < len(blobs); i++ {
		err := benchStore.WriteToKey(strconv.Itoa(i%10), blobs[i])
		if err != nil {
			b.Fatal("Failed to write to store:", err)
		}
	}
	b.StopTimer()
	os.RemoveAll(benchDir)

}
<|endoftext|>


<commit_before>package saw

import (
	"sync"

	"golang.org/x/net/context"
)

type WindowSeqFunc func(datum Datum) int
type WindowFrameFactory func(name string, seq int) (Saw, error)

type WindowSpec struct {
	Name          string
	FrameFactory  WindowFrameFactory
	SeqFunc       WindowSeqFunc
	WindowSize    int
	MaxSeqAdvance int
}

type Window struct {
	spec WindowSpec

	mu        sync.Mutex
	frames    []Saw
	startSeq  int
	latestSeq int
	startIdx  int
	hasData   bool

	finalizeWg sync.WaitGroup

	droppedCount VarInt
}

func NewWindow(spec WindowSpec) *Window {
	return &Window{
		spec:         spec,
		frames:       make([]Saw, spec.WindowSize),
		droppedCount: ReportInt(spec.Name, "droppedCount"),
	}
}

func (win *Window) asyncFinalize(seq int, frame Saw) {
	win.finalizeWg.Add(1)
	go func() {
		frame.Result(context.Background())
		win.finalizeWg.Done()
	}()
}

// returning Saw nullable indicating drop
func (win *Window) prepareFrame(datum Datum) (frame Saw, err error) {
	seq := win.spec.SeqFunc(datum)
	win.mu.Lock()
	defer win.mu.Unlock()
	if !win.hasData {
		frame, err = win.spec.FrameFactory(win.spec.Name, seq)
		if err != nil {
			return
		}
		win.startSeq = seq
		win.latestSeq = seq
		win.startIdx = 0
		win.frames[win.startIdx] = frame
		win.hasData = true
		return frame, nil
	}
	advance := seq - win.startSeq
	// Out of window
	if advance < 0 || (win.spec.MaxSeqAdvance > 0 && advance > win.spec.MaxSeqAdvance) {
		return nil, nil
	}
	if advance < len(win.frames) {
		idx := (win.startIdx + advance) % len(win.frames)
		if win.frames[idx] == nil {
			win.frames[idx], err = win.spec.FrameFactory(win.spec.Name, seq)
			if err != nil {
				return nil, err
			}
		}
		return win.frames[idx], nil
	}
	frame, err = win.spec.FrameFactory(win.spec.Name, seq)
	if err != nil {
		return
	}
	if seq > win.latestSeq {
		win.latestSeq = seq
	}
	if advance >= len(win.frames)*2 {
		for i := 0; i < len(win.frames); i++ {
			frameIdx := (win.startIdx + i) % len(win.frames)
			frame := win.frames[frameIdx]
			if frame != nil {
				win.frames[frameIdx] = nil
				win.asyncFinalize(win.startSeq+i, frame)
			}
		}
		win.startSeq = seq - len(win.frames) + 1
		win.startIdx = 0
		win.frames[len(win.frames)-1] = frame
		return frame, nil
	} else {
		for i := 0; i < advance-len(win.frames); i++ {
			if win.frames[win.startIdx] != nil {
				win.asyncFinalize(win.startSeq, win.frames[win.startIdx])
				win.frames[win.startIdx] = nil
			}
			win.startIdx = (win.startIdx + 1) % len(win.frames)
			win.startSeq++
		}
		win.frames[(win.startIdx+(seq-win.startSeq))%len(win.frames)] = frame
		return frame, nil
	}
}

func (win *Window) Emit(datum Datum) error {
	frame, err := win.prepareFrame(datum)
	if err != nil {
		return err
	}
	if frame == nil {
		win.droppedCount.Add(1)
		return nil
	}
	return frame.Emit(datum)
}

func (win *Window) Result(ctx context.Context) (result interface{}, err error) {
	win.mu.Lock()
	defer win.mu.Unlock()
	for i := 0; i < len(win.frames); i++ {
		frameIdx := (win.startIdx + i) % len(win.frames)
		frame := win.frames[frameIdx]
		if frame != nil {
			win.frames[frameIdx] = nil
			win.asyncFinalize(win.startSeq+i, frame)
		}
	}
	win.startSeq = 0
	win.startIdx = 0
	win.latestSeq = 0
	win.hasData = false
	win.finalizeWg.Wait()
	return nil, nil
}

func (win *Window) LatestFrame() (seq int, frame Saw) {
	win.mu.Lock()
	defer win.mu.Unlock()

	if !win.hasData {
		return 0, nil
	}
	return win.latestSeq, win.frames[(win.startIdx+win.latestSeq-win.startSeq)%len(win.frames)]
}

func (win *Window) AllFrames() map[int]Saw {
	win.mu.Lock()
	defer win.mu.Unlock()
	output := make(map[int]Saw)
	if !win.hasData {
		return output
	}
	for i := 0; i < len(win.frames); i++ {
		frame := win.frames[(win.startIdx+i)%len(win.frames)]
		if frame != nil {
			output[win.startSeq+i] = frame
		}
	}
	return output
}
<commit_msg>sliding window code cleanup and docs.<commit_after>package saw

import (
	"sync"

	"golang.org/x/net/context"
)

type SeqID int

func (seq SeqID) Advance(x int) SeqID {
	return SeqID(int(seq) + x)
}

func (seq SeqID) DistanceFrom(other SeqID) int {
	return int(seq - other)
}

type DatumSeqFunc func(datum Datum) SeqID
type WindowFrameFactory func(name string, seq SeqID) (Saw, error)

type WindowSpec struct {
	Name          string
	FrameFactory  WindowFrameFactory
	SeqFunc       DatumSeqFunc
	WindowSize    int
	MaxSeqAdvance int
}

// Window implements a sliding window of saws. Window keeps finite set of frames,
// each corresponded with a SeqID. Window keeps WindowSpec.WindowSize of latest
// frames with largest, contintued SeqID.
//
// For every input, it calls WindowSpec.SeqFunc to get SeqID of a datum and route
// it to the frame if it's still in sliding window. Window assumes datum's SeqID
// are dense, rougly incremental overtime, or you will get too many datums dropped.
//
// When window needed to be slided, Result() of old frames will be called in their
// seperate gorountines. In Window.Result(), all frames' Result() will be called
// in this maner as they all been slide away. Window doesn't care frame's Result()
// return.
type Window struct {
	spec WindowSpec

	mu        sync.Mutex
	frames    []Saw
	startSeq  SeqID
	latestSeq SeqID
	startIdx  int
	hasData   bool

	finalizeWg sync.WaitGroup

	droppedCount VarInt
}

func NewWindow(spec WindowSpec) *Window {
	return &Window{
		spec:         spec,
		frames:       make([]Saw, spec.WindowSize),
		droppedCount: ReportInt(spec.Name, "droppedCount"),
	}
}

func (win *Window) asyncFinalize(seq SeqID, frame Saw) {
	win.finalizeWg.Add(1)
	go func() {
		frame.Result(context.Background())
		win.finalizeWg.Done()
	}()
}

func (win *Window) indexForSeq(seq SeqID) int {
	return (win.startIdx + seq.DistanceFrom(win.startSeq)) % len(win.frames)
}

func (win *Window) indexForOffset(offset int) int {
	return (win.startIdx + offset) % len(win.frames)
}

// returning Saw nullable indicating drop
func (win *Window) prepareFrame(datum Datum) (frame Saw, err error) {
	seq := win.spec.SeqFunc(datum)
	win.mu.Lock()
	defer win.mu.Unlock()
	if !win.hasData {
		frame, err = win.spec.FrameFactory(win.spec.Name, seq)
		if err != nil {
			return
		}
		win.startSeq = seq
		win.latestSeq = seq
		win.startIdx = 0
		win.frames[win.startIdx] = frame
		win.hasData = true
		return frame, nil
	}
	offset := seq.DistanceFrom(win.startSeq)
	// Out of window, drop
	if offset < 0 || (win.spec.MaxSeqAdvance > 0 && offset > win.spec.MaxSeqAdvance) {
		win.droppedCount.Add(1)
		return nil, nil
	}
	winSize := len(win.frames)
	if offset < winSize {
		frameIdx := win.indexForOffset(offset)
		if win.frames[frameIdx] == nil {
			win.frames[frameIdx], err = win.spec.FrameFactory(win.spec.Name, seq)
			if err != nil {
				return nil, err
			}
		}
		return win.frames[frameIdx], nil
	}
	frame, err = win.spec.FrameFactory(win.spec.Name, seq)
	if err != nil {
		return
	}
	if seq > win.latestSeq {
		win.latestSeq = seq
	}
	if offset >= winSize*2 {
		for i := 0; i < winSize; i++ {
			frameIdx := win.indexForOffset(i)
			frame := win.frames[frameIdx]
			if frame != nil {
				win.frames[frameIdx] = nil
				win.asyncFinalize(win.startSeq.Advance(i), frame)
			}
		}
		win.startSeq = seq.Advance(1 - winSize)
		win.startIdx = 0
	} else {
		for i := 0; i < offset-winSize; i++ {
			if win.frames[win.startIdx] != nil {
				win.asyncFinalize(win.startSeq, win.frames[win.startIdx])
				win.frames[win.startIdx] = nil
			}
			win.startIdx = win.indexForOffset(1)
			win.startSeq++
		}
	}
	win.frames[win.indexForSeq(seq)] = frame
	return frame, nil
}

func (win *Window) Emit(datum Datum) error {
	frame, err := win.prepareFrame(datum)
	if err != nil {
		return err
	}
	if frame == nil {
		win.droppedCount.Add(1)
		return nil
	}
	return frame.Emit(datum)
}

// Result finalize all frames it's currently managing, returns after all frames
// sent for finalize finishes, including previous ones caused by sliding.
func (win *Window) Result(ctx context.Context) (result interface{}, err error) {
	win.mu.Lock()
	defer win.mu.Unlock()
	for i := 0; i < len(win.frames); i++ {
		frameIdx := win.indexForOffset(i)
		frame := win.frames[frameIdx]
		if frame != nil {
			win.frames[frameIdx] = nil
			win.asyncFinalize(win.startSeq.Advance(i), frame)
		}
	}
	win.startSeq = 0
	win.startIdx = 0
	win.latestSeq = 0
	win.hasData = false
	win.finalizeWg.Wait()
	return nil, nil
}

// Gets the latest frame or nil when there's no data yet. returned frame is not
// locked, would be Emit()-ing or even Result()-ing when caller gets the return.
func (win *Window) LatestFrame() (seq SeqID, frame Saw) {
	win.mu.Lock()
	defer win.mu.Unlock()

	if !win.hasData {
		return 0, nil
	}
	return win.latestSeq, win.frames[win.indexForSeq(win.latestSeq)]
}

// Gets the all frame the Window currently holds. returned frames are not locked,
// would be Emit()-ing or even Result()-ing when caller gets the return.
func (win *Window) AllFrames() map[SeqID]Saw {
	win.mu.Lock()
	defer win.mu.Unlock()
	output := make(map[SeqID]Saw)
	if !win.hasData {
		return output
	}
	for i := 0; i < len(win.frames); i++ {
		frame := win.frames[win.indexForOffset(i)]
		if frame != nil {
			output[win.startSeq.Advance(i)] = frame
		}
	}
	return output
}
<|endoftext|>