<commit_before>package middleware

import (
	"bytes"
	"encoding/json"
	"net/http"
	. "github.com/mickael-kerjean/filestash/server/common"
	"time"
	"sync"
)

type Middleware func(func(App, http.ResponseWriter, *http.Request)) func(App, http.ResponseWriter, *http.Request)

func NewMiddlewareChain(fn func(App, http.ResponseWriter, *http.Request), m []Middleware, app App) http.HandlerFunc {
	return func(res http.ResponseWriter, req *http.Request) {
		var resw ResponseWriter                             = NewResponseWriter(res)
		var f func(App, http.ResponseWriter, *http.Request) = fn

		for i := len(m) - 1; i >= 0; i-- {
			f = m[i](f)
		}
		f(app, &resw, req)
		if req.Body != nil {
			req.Body.Close()
		}
		go Logger(app, &resw, req)
	}
}

type ResponseWriter struct {
	http.ResponseWriter
	status int
	start  time.Time
}

func NewResponseWriter(res http.ResponseWriter) ResponseWriter {
	return ResponseWriter{
		ResponseWriter: res,
		start: time.Now(),
	}
}

func (w *ResponseWriter) WriteHeader(status int) {
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *ResponseWriter) Write(b []byte) (int, error) {
	if w.status == 0 {
		w.status = 200
	}
	return w.ResponseWriter.Write(b)
}

type LogEntry struct {
	Host       string    `json:"host"`
	Method     string    `json:"method"`
	RequestURI string    `json:"pathname"`
	Proto      string    `json:"proto"`
	Status     int       `json:"status"`
	Scheme     string    `json:"scheme"`
	UserAgent  string    `json:"userAgent"`
	Ip         string    `json:"ip"`
	Referer    string    `json:"referer"`
	Timestamp  time.Time `json:"_id"`
	Duration   float64   `json:"responseTime"`
	Version    string    `json:"version"`
	Backend    string    `json:"backend"`
}

func Logger(ctx App, res http.ResponseWriter, req *http.Request) {
	if obj, ok := res.(*ResponseWriter); ok && req.RequestURI != "/about" {
		point := LogEntry{
			Version:    APP_VERSION,
			Scheme:     req.URL.Scheme,
			Host:       req.Host,
			Method:     req.Method,
			RequestURI: req.RequestURI,
			Proto:      req.Proto,
			Status:     obj.status,
			UserAgent:  req.Header.Get("User-Agent"),
			Ip:         req.RemoteAddr,
			Referer:    req.Referer(),
			Duration:   float64(time.Now().Sub(obj.start)) / (1000 * 1000),
			Timestamp:  time.Now().UTC(),
			Backend:    ctx.Session["type"],
		}
		if Config.Get("log.telemetry").Bool() {
			telemetry.Record(point)
		}
		if Config.Get("log.enable").Bool() {
			Log.Info("HTTP %3d %3s %6.1fms %s", point.Status, point.Method, point.Duration, point.RequestURI)
		}
	}
}

type Telemetry struct {
	data []LogEntry
	mu   sync.Mutex
}

func (this *Telemetry) Record(point LogEntry) {
	this.mu.Lock()
	this.data = append(this.data, point)
	this.mu.Unlock()
}

func (this *Telemetry) Flush() {
	if len(this.data) == 0 {
		return
	}
	this.mu.Lock()
	pts := this.data
	this.data = make([]LogEntry, 0)
	this.mu.Unlock()

	// send data in bulk: http://docs.couchdb.org/en/2.2.0/api/database/bulk-api.html#inserting-documents-in-bulk
	data := struct {
		Docs []LogEntry `json:"docs"`
	}{ pts }
	body, err := json.Marshal(data)
	if err != nil {
		return
	}
	r, err := http.NewRequest("POST", "https://log.kerjean.me/trash/_bulk_docs", bytes.NewReader(body))
	r.Header.Set("Connection", "Close")
	r.Header.Set("Content-Type", "application/json")
	r.Close = true
	if err != nil {
		return
	}
	resp, err := HTTP.Do(r)
	if err != nil {
		return
	}
	resp.Body.Close()
}

var telemetry Telemetry = Telemetry{ data: make([]LogEntry, 0) }

func init(){
	go func(){
		for {
			select {
			default:
				time.Sleep(1000 * 10 * time.Millisecond)
				telemetry.Flush()
			}
		}
	}()
}
<commit_msg>maintain (telemetry): update API to a custom endpoint<commit_after>package middleware

import (
	"bytes"
	"encoding/json"
	"net/http"
	. "github.com/mickael-kerjean/filestash/server/common"
	"time"
	"sync"
)

type Middleware func(func(App, http.ResponseWriter, *http.Request)) func(App, http.ResponseWriter, *http.Request)

func NewMiddlewareChain(fn func(App, http.ResponseWriter, *http.Request), m []Middleware, app App) http.HandlerFunc {
	return func(res http.ResponseWriter, req *http.Request) {
		var resw ResponseWriter                             = NewResponseWriter(res)
		var f func(App, http.ResponseWriter, *http.Request) = fn

		for i := len(m) - 1; i >= 0; i-- {
			f = m[i](f)
		}
		f(app, &resw, req)
		if req.Body != nil {
			req.Body.Close()
		}
		go Logger(app, &resw, req)
	}
}

type ResponseWriter struct {
	http.ResponseWriter
	status int
	start  time.Time
}

func NewResponseWriter(res http.ResponseWriter) ResponseWriter {
	return ResponseWriter{
		ResponseWriter: res,
		start: time.Now(),
	}
}

func (w *ResponseWriter) WriteHeader(status int) {
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *ResponseWriter) Write(b []byte) (int, error) {
	if w.status == 0 {
		w.status = 200
	}
	return w.ResponseWriter.Write(b)
}

type LogEntry struct {
	Host       string    `json:"host"`
	Method     string    `json:"method"`
	RequestURI string    `json:"pathname"`
	Proto      string    `json:"proto"`
	Status     int       `json:"status"`
	Scheme     string    `json:"scheme"`
	UserAgent  string    `json:"userAgent"`
	Ip         string    `json:"ip"`
	Referer    string    `json:"referer"`
	Duration   float64   `json:"responseTime"`
	Version    string    `json:"version"`
	Backend    string    `json:"backend"`
}

func Logger(ctx App, res http.ResponseWriter, req *http.Request) {
	if obj, ok := res.(*ResponseWriter); ok && req.RequestURI != "/about" {
		point := LogEntry{
			Version:    APP_VERSION,
			Scheme:     req.URL.Scheme,
			Host:       req.Host,
			Method:     req.Method,
			RequestURI: req.RequestURI,
			Proto:      req.Proto,
			Status:     obj.status,
			UserAgent:  req.Header.Get("User-Agent"),
			Ip:         req.RemoteAddr,
			Referer:    req.Referer(),
			Duration:   float64(time.Now().Sub(obj.start)) / (1000 * 1000),
			Backend:    ctx.Session["type"],
		}
		if Config.Get("log.telemetry").Bool() {
			telemetry.Record(point)
		}
		if Config.Get("log.enable").Bool() {
			Log.Info("HTTP %3d %3s %6.1fms %s", point.Status, point.Method, point.Duration, point.RequestURI)
		}
	}
}

type Telemetry struct {
	Data []LogEntry
	mu   sync.Mutex
}

func (this *Telemetry) Record(point LogEntry) {
	this.mu.Lock()
	this.Data = append(this.Data, point)
	this.mu.Unlock()
}

func (this *Telemetry) Flush() {
	if len(this.Data) == 0 {
		return
	}
	this.mu.Lock()
	pts := this.Data
	this.Data = make([]LogEntry, 0)
	this.mu.Unlock()

	body, err := json.Marshal(pts)
	if err != nil {
		return
	}
	r, err := http.NewRequest("POST", "https://downloads.filestash.app/event", bytes.NewReader(body))
	r.Header.Set("Connection", "Close")
	r.Header.Set("Content-Type", "application/json")
	r.Close = true
	if err != nil {
		return
	}
	resp, err := HTTP.Do(r)
	if err != nil {
		return
	}
	resp.Body.Close()
}

var telemetry Telemetry = Telemetry{ Data: make([]LogEntry, 0) }

func init(){
	go func(){
		for {
			time.Sleep(10 * time.Second)
			telemetry.Flush()
		}
	}()
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 Space Monkey, 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 present

import (
	"bufio"
	"io"
	"net/url"
	"regexp"
	"strconv"
	"strings"

	"github.com/spacemonkeygo/monkit/v3"
)

// Result writes the expected data to io.Writer and returns any errors if
// found.
type Result func(io.Writer) error

func curry(reg *monkit.Registry,
	f func(*monkit.Registry, io.Writer) error) func(io.Writer) error {
	return func(w io.Writer) error {
		return f(reg, w)
	}
}

// FromRequest takes a registry (usually the Default registry), an incoming
// path, and optional query parameters, and returns a Result if possible.
//
// FromRequest understands the following paths:
//  * /ps, /ps/text       - returns the result of SpansText
//  * /ps/dot             - returns the result of SpansDot
//  * /ps/json            - returns the result of SpansJSON
//  * /funcs, /funcs/text - returns the result of FuncsText
//  * /funcs/dot          - returns the result of FuncsDot
//  * /funcs/json         - returns the result of FuncsJSON
//  * /stats, /stats/text - returns the result of StatsText
//  * /stats/json         - returns the result of StatsJSON
//  * /trace/svg          - returns the result of TraceQuerySVG
//  * /trace/json         - returns the result of TraceQueryJSON
//
// The last two paths are worth discussing in more detail, as they take
// query parameters. All trace endpoints require at least one of the following
// two query parameters:
//  * regex    - If provided, the very next Span that crosses a Func that has
//               a name that matches this regex will start a trace until that
//               triggering Span ends, provided the trace_id matches.
//  * trace_id - If provided, the very next Span on a trace with the given
//               trace id will start a trace until the triggering Span ends,
//               provided the regex matches. NOTE: the trace_id will be parsed
//               in hex.
// By default, regular expressions are matched ahead of time against all known
// Funcs, but perhaps the Func you want to trace hasn't been observed by the
// process yet, in which case the regex will fail to match anything. You can
// turn off this preselection behavior by providing preselect=false as an
// additional query param. Be advised that until a trace completes, whether
// or not it has started, it adds a small amount of overhead (a comparison or
// two) to every monitored function.
func FromRequest(reg *monkit.Registry, path string, query url.Values) (
	f Result, contentType string, err error) {

	defer func() {
		if err != nil {
			return
		}
		// wrap all functions with buffering
		unbuffered := f
		f = func(w io.Writer) (err error) {
			buf := bufio.NewWriter(w)
			err = unbuffered(buf)
			if err != nil {
				return err
			}
			err = buf.Flush()
			return err
		}
	}()

	first, rest := shift(path)
	second, _ := shift(rest)
	switch first {
	case "ps":
		switch second {
		case "", "text":
			return curry(reg, SpansText), "text/plain; charset=utf-8", nil
		case "dot":
			return curry(reg, SpansDot), "text/plain; charset=utf-8", nil
		case "json":
			return curry(reg, SpansJSON), "application/json; charset=utf-8", nil
		}

	case "funcs":
		switch second {
		case "", "text":
			return curry(reg, FuncsText), "text/plain; charset=utf-8", nil
		case "dot":
			return curry(reg, FuncsDot), "text/plain; charset=utf-8", nil
		case "json":
			return curry(reg, FuncsJSON), "application/json; charset=utf-8", nil
		}

	case "stats":
		switch second {
		case "", "text":
			return func(w io.Writer) error {
				return StatsText(reg, w)
			}, "text/plain; charset=utf-8", nil
		case "json":
			return func(w io.Writer) error {
				return StatsJSON(reg, w)
			}, "application/json; charset=utf-8", nil
		case "old":
			return func(w io.Writer) error {
				return StatsOld(reg, w)
			}, "text/plain; charset=utf-8", nil
		}

	case "trace":
		regexStr := query.Get("regex")
		traceIdStr := query.Get("trace_id")
		if regexStr == "" && traceIdStr == "" {
			return nil, "", errBadRequest.New("at least one of 'regex' or 'trace_id' " +
				"query parameters required")
		}
		fnMatcher := func(*monkit.Func) bool { return true }

		if regexStr != "" {
			re, err := regexp.Compile(regexStr)
			if err != nil {
				return nil, "", errBadRequest.New("invalid regex %#v: %v",
					regexStr, err)
			}
			fnMatcher = func(f *monkit.Func) bool {
				return re.MatchString(f.FullName())
			}

			preselect := true
			if query.Get("preselect") != "" {
				preselect, err = strconv.ParseBool(query.Get("preselect"))
				if err != nil {
					return nil, "", errBadRequest.New("invalid preselect %#v: %v",
						query.Get("preselect"), err)
				}
			}
			if preselect {
				funcs := map[*monkit.Func]bool{}
				reg.Funcs(func(f *monkit.Func) {
					if fnMatcher(f) {
						funcs[f] = true
					}
				})
				if len(funcs) <= 0 {
					return nil, "", errBadRequest.New("regex preselect matches 0 functions")
				}

				fnMatcher = func(f *monkit.Func) bool { return funcs[f] }
			}
		}

		spanMatcher := func(s *monkit.Span) bool { return fnMatcher(s.Func()) }

		if traceIdStr != "" {
			traceId, err := strconv.ParseUint(traceIdStr, 16, 64)
			if err != nil {
				return nil, "", errBadRequest.New(
					"trace_id expected to be hex unsigned 64 bit number: %#v", traceIdStr)
			}
			spanMatcher = func(s *monkit.Span) bool {
				return s.Trace().Id() == int64(traceId) && fnMatcher(s.Func())
			}
		}

		switch second {
		case "svg":
			return func(w io.Writer) error {
				return TraceQuerySVG(reg, w, spanMatcher)
			}, "image/svg+xml; charset=utf-8", nil
		case "json":
			return func(w io.Writer) error {
				return TraceQueryJSON(reg, w, spanMatcher)
			}, "application/json; charset=utf-8", nil
		}
	}
	return nil, "", errNotFound.New("path not found: %s", path)
}

func shift(path string) (dir, left string) {
	path = strings.TrimLeft(path, "/")
	split := strings.Index(path, "/")
	if split == -1 {
		return path, ""
	}
	return path[:split], path[split:]
}
<commit_msg>present: add basic help to index<commit_after>// Copyright (C) 2015 Space Monkey, 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 present

import (
	"bufio"
	"io"
	"net/url"
	"regexp"
	"strconv"
	"strings"

	"github.com/spacemonkeygo/monkit/v3"
)

// Result writes the expected data to io.Writer and returns any errors if
// found.
type Result func(io.Writer) error

func curry(reg *monkit.Registry,
	f func(*monkit.Registry, io.Writer) error) func(io.Writer) error {
	return func(w io.Writer) error {
		return f(reg, w)
	}
}

// FromRequest takes a registry (usually the Default registry), an incoming
// path, and optional query parameters, and returns a Result if possible.
//
// FromRequest understands the following paths:
//  * /ps, /ps/text       - returns the result of SpansText
//  * /ps/dot             - returns the result of SpansDot
//  * /ps/json            - returns the result of SpansJSON
//  * /funcs, /funcs/text - returns the result of FuncsText
//  * /funcs/dot          - returns the result of FuncsDot
//  * /funcs/json         - returns the result of FuncsJSON
//  * /stats, /stats/text - returns the result of StatsText
//  * /stats/json         - returns the result of StatsJSON
//  * /trace/svg          - returns the result of TraceQuerySVG
//  * /trace/json         - returns the result of TraceQueryJSON
//
// The last two paths are worth discussing in more detail, as they take
// query parameters. All trace endpoints require at least one of the following
// two query parameters:
//  * regex    - If provided, the very next Span that crosses a Func that has
//               a name that matches this regex will start a trace until that
//               triggering Span ends, provided the trace_id matches.
//  * trace_id - If provided, the very next Span on a trace with the given
//               trace id will start a trace until the triggering Span ends,
//               provided the regex matches. NOTE: the trace_id will be parsed
//               in hex.
// By default, regular expressions are matched ahead of time against all known
// Funcs, but perhaps the Func you want to trace hasn't been observed by the
// process yet, in which case the regex will fail to match anything. You can
// turn off this preselection behavior by providing preselect=false as an
// additional query param. Be advised that until a trace completes, whether
// or not it has started, it adds a small amount of overhead (a comparison or
// two) to every monitored function.
func FromRequest(reg *monkit.Registry, path string, query url.Values) (
	f Result, contentType string, err error) {

	defer func() {
		if err != nil {
			return
		}
		// wrap all functions with buffering
		unbuffered := f
		f = func(w io.Writer) (err error) {
			buf := bufio.NewWriter(w)
			err = unbuffered(buf)
			if err != nil {
				return err
			}
			err = buf.Flush()
			return err
		}
	}()

	first, rest := shift(path)
	second, _ := shift(rest)
	switch first {
	case "":
		return writeIndex, "text/html", nil
	case "ps":
		switch second {
		case "", "text":
			return curry(reg, SpansText), "text/plain; charset=utf-8", nil
		case "dot":
			return curry(reg, SpansDot), "text/plain; charset=utf-8", nil
		case "json":
			return curry(reg, SpansJSON), "application/json; charset=utf-8", nil
		}

	case "funcs":
		switch second {
		case "", "text":
			return curry(reg, FuncsText), "text/plain; charset=utf-8", nil
		case "dot":
			return curry(reg, FuncsDot), "text/plain; charset=utf-8", nil
		case "json":
			return curry(reg, FuncsJSON), "application/json; charset=utf-8", nil
		}

	case "stats":
		switch second {
		case "", "text":
			return func(w io.Writer) error {
				return StatsText(reg, w)
			}, "text/plain; charset=utf-8", nil
		case "json":
			return func(w io.Writer) error {
				return StatsJSON(reg, w)
			}, "application/json; charset=utf-8", nil
		case "old":
			return func(w io.Writer) error {
				return StatsOld(reg, w)
			}, "text/plain; charset=utf-8", nil
		}

	case "trace":
		regexStr := query.Get("regex")
		traceIdStr := query.Get("trace_id")
		if regexStr == "" && traceIdStr == "" {
			return nil, "", errBadRequest.New("at least one of 'regex' or 'trace_id' " +
				"query parameters required")
		}
		fnMatcher := func(*monkit.Func) bool { return true }

		if regexStr != "" {
			re, err := regexp.Compile(regexStr)
			if err != nil {
				return nil, "", errBadRequest.New("invalid regex %#v: %v",
					regexStr, err)
			}
			fnMatcher = func(f *monkit.Func) bool {
				return re.MatchString(f.FullName())
			}

			preselect := true
			if query.Get("preselect") != "" {
				preselect, err = strconv.ParseBool(query.Get("preselect"))
				if err != nil {
					return nil, "", errBadRequest.New("invalid preselect %#v: %v",
						query.Get("preselect"), err)
				}
			}
			if preselect {
				funcs := map[*monkit.Func]bool{}
				reg.Funcs(func(f *monkit.Func) {
					if fnMatcher(f) {
						funcs[f] = true
					}
				})
				if len(funcs) <= 0 {
					return nil, "", errBadRequest.New("regex preselect matches 0 functions")
				}

				fnMatcher = func(f *monkit.Func) bool { return funcs[f] }
			}
		}

		spanMatcher := func(s *monkit.Span) bool { return fnMatcher(s.Func()) }

		if traceIdStr != "" {
			traceId, err := strconv.ParseUint(traceIdStr, 16, 64)
			if err != nil {
				return nil, "", errBadRequest.New(
					"trace_id expected to be hex unsigned 64 bit number: %#v", traceIdStr)
			}
			spanMatcher = func(s *monkit.Span) bool {
				return s.Trace().Id() == int64(traceId) && fnMatcher(s.Func())
			}
		}

		switch second {
		case "svg":
			return func(w io.Writer) error {
				return TraceQuerySVG(reg, w, spanMatcher)
			}, "image/svg+xml; charset=utf-8", nil
		case "json":
			return func(w io.Writer) error {
				return TraceQueryJSON(reg, w, spanMatcher)
			}, "application/json; charset=utf-8", nil
		}
	}
	return nil, "", errNotFound.New("path not found: %s", path)
}

func shift(path string) (dir, left string) {
	path = strings.TrimLeft(path, "/")
	split := strings.Index(path, "/")
	if split == -1 {
		return path, ""
	}
	return path[:split], path[split:]
}

func writeIndex(w io.Writer) error {
	_, err := w.Write([]byte(`<!DOCTYPE html>

		<meta charset="utf-8">
		<title>Monkit</title>
		<meta http-equiv="refresh" content="1">

/ps

/ps/json

/ps/dot

Information about active spans.

/funcs

/funcs/json

/funcs/dot

Information about the functions and their relations.

/stats

/stats/json

/stats/svg

Statistics about all observed functions, scopes and values.

/trace/json

/trace/svg

Trace the next scope that matches one of the ?regex= or ?trace_id= query arguments. By default, regular expressions are matched ahead of time against all known Funcs, but perhaps the Func you want to trace hasn't been observed by the process yet, in which case the regex will fail to match anything. You can turn off this preselection behavior by providing &preselect=false as an additional query param. Be advised that until a trace completes, whether or not it has started, it adds a small amount of overhead (a comparison or two) to every monitored function.

`))
	return err
}
<|endoftext|>


<commit_before>// Copyright 2016-2021 The Libsacloud 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 fake

import (
	"context"
	"fmt"
	"net"
	"time"

	"github.com/sacloud/libsacloud/v2/sacloud"
	"github.com/sacloud/libsacloud/v2/sacloud/types"
)

// Find is fake implementation
func (o *VPCRouterOp) Find(ctx context.Context, zone string, conditions *sacloud.FindCondition) (*sacloud.VPCRouterFindResult, error) {
	results, _ := find(o.key, zone, conditions)
	var values []*sacloud.VPCRouter
	for _, res := range results {
		dest := &sacloud.VPCRouter{}
		copySameNameField(res, dest)
		values = append(values, dest)
	}
	return &sacloud.VPCRouterFindResult{
		Total:      len(results),
		Count:      len(results),
		From:       0,
		VPCRouters: values,
	}, nil
}

// Create is fake implementation
func (o *VPCRouterOp) Create(ctx context.Context, zone string, param *sacloud.VPCRouterCreateRequest) (*sacloud.VPCRouter, error) {
	result := &sacloud.VPCRouter{}
	copySameNameField(param, result)
	fill(result, fillID, fillCreatedAt)

	result.Class = "vpcrouter"
	result.Availability = types.Availabilities.Migrating
	result.ZoneID = zoneIDs[zone]
	result.SettingsHash = ""
	if result.Version == 0 {
		result.Version = 2
	}

	ifOp := NewInterfaceOp()
	swOp := NewSwitchOp()

	ifCreateParam := &sacloud.InterfaceCreateRequest{}
	if param.Switch.Scope == types.Scopes.Shared {
		ifCreateParam.ServerID = result.ID
	} else {
		_, err := swOp.Read(ctx, zone, param.Switch.ID)
		if err != nil {
			return nil, newErrorConflict(o.key, types.ID(0), err.Error())
		}
	}

	iface, err := ifOp.Create(ctx, zone, ifCreateParam)
	if err != nil {
		return nil, newErrorConflict(o.key, types.ID(0), err.Error())
	}

	if param.Switch.Scope == types.Scopes.Shared {
		if err := ifOp.ConnectToSharedSegment(ctx, zone, iface.ID); err != nil {
			return nil, newErrorConflict(o.key, types.ID(0), err.Error())
		}
	} else {
		if err := ifOp.ConnectToSwitch(ctx, zone, iface.ID, param.Switch.ID); err != nil {
			return nil, newErrorConflict(o.key, types.ID(0), err.Error())
		}
	}

	iface, err = ifOp.Read(ctx, zone, iface.ID)
	if err != nil {
		return nil, newErrorConflict(o.key, types.ID(0), err.Error())
	}

	vpcRouterInterface := &sacloud.VPCRouterInterface{}
	copySameNameField(iface, vpcRouterInterface)
	if param.Switch.Scope == types.Scopes.Shared {
		sharedIP := pool().nextSharedIP()
		vpcRouterInterface.IPAddress = sharedIP.String()
		vpcRouterInterface.SubnetNetworkMaskLen = sharedSegmentSwitch.NetworkMaskLen

		ipv4Mask := net.CIDRMask(pool().SharedNetMaskLen, 32)
		vpcRouterInterface.SubnetNetworkAddress = sharedIP.Mask(ipv4Mask).String()
		vpcRouterInterface.SubnetDefaultRoute = pool().SharedDefaultGateway.String()
	}
	result.Interfaces = append(result.Interfaces, vpcRouterInterface)

	putVPCRouter(zone, result)

	id := result.ID
	startMigration(o.key, zone, func() (interface{}, error) {
		return o.Read(context.Background(), zone, id)
	})
	return result, nil
}

// Read is fake implementation
func (o *VPCRouterOp) Read(ctx context.Context, zone string, id types.ID) (*sacloud.VPCRouter, error) {
	value := getVPCRouterByID(zone, id)
	if value == nil {
		return nil, newErrorNotFound(o.key, id)
	}
	dest := &sacloud.VPCRouter{}
	copySameNameField(value, dest)
	return dest, nil
}

// Update is fake implementation
func (o *VPCRouterOp) Update(ctx context.Context, zone string, id types.ID, param *sacloud.VPCRouterUpdateRequest) (*sacloud.VPCRouter, error) {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return nil, err
	}
	copySameNameField(param, value)
	fill(value, fillModifiedAt)

	putVPCRouter(zone, value)
	return value, nil
}

// UpdateSettings is fake implementation
func (o *VPCRouterOp) UpdateSettings(ctx context.Context, zone string, id types.ID, param *sacloud.VPCRouterUpdateSettingsRequest) (*sacloud.VPCRouter, error) {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return nil, err
	}
	copySameNameField(param, value)
	fill(value, fillModifiedAt)

	putVPCRouter(zone, value)
	return value, nil
}

// Delete is fake implementation
func (o *VPCRouterOp) Delete(ctx context.Context, zone string, id types.ID) error {
	_, err := o.Read(ctx, zone, id)
	if err != nil {
		return err
	}
	ds().Delete(o.key, zone, id)
	return nil
}

// Config is fake implementation
func (o *VPCRouterOp) Config(ctx context.Context, zone string, id types.ID) error {
	return nil
}

// Boot is fake implementation
func (o *VPCRouterOp) Boot(ctx context.Context, zone string, id types.ID) error {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return err
	}
	if value.InstanceStatus.IsUp() {
		return newErrorConflict(o.key, id, "Boot is failed")
	}

	startPowerOn(o.key, zone, func() (interface{}, error) {
		return o.Read(context.Background(), zone, id)
	})

	return err
}

// Shutdown is fake implementation
func (o *VPCRouterOp) Shutdown(ctx context.Context, zone string, id types.ID, shutdownOption *sacloud.ShutdownOption) error {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return err
	}
	if !value.InstanceStatus.IsUp() {
		return newErrorConflict(o.key, id, "Shutdown is failed")
	}

	startPowerOff(o.key, zone, func() (interface{}, error) {
		return o.Read(context.Background(), zone, id)
	})

	return err
}

// Reset is fake implementation
func (o *VPCRouterOp) Reset(ctx context.Context, zone string, id types.ID) error {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return err
	}
	if !value.InstanceStatus.IsUp() {
		return newErrorConflict(o.key, id, "Reset is failed")
	}

	startPowerOn(o.key, zone, func() (interface{}, error) {
		return o.Read(context.Background(), zone, id)
	})

	return nil
}

// ConnectToSwitch is fake implementation
func (o *VPCRouterOp) ConnectToSwitch(ctx context.Context, zone string, id types.ID, nicIndex int, switchID types.ID) error {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return err
	}

	for _, nic := range value.Interfaces {
		if nic.Index == nicIndex {
			return newErrorBadRequest(o.key, id, fmt.Sprintf("nic[%d] already connected to switch", nicIndex))
		}
	}

	// find switch
	swOp := NewSwitchOp()
	_, err = swOp.Read(ctx, zone, switchID)
	if err != nil {
		return fmt.Errorf("ConnectToSwitch is failed: %s", err)
	}

	// create interface
	ifOp := NewInterfaceOp()
	iface, err := ifOp.Create(ctx, zone, &sacloud.InterfaceCreateRequest{ServerID: id})
	if err != nil {
		return newErrorConflict(o.key, types.ID(0), err.Error())
	}

	if err := ifOp.ConnectToSwitch(ctx, zone, iface.ID, switchID); err != nil {
		return newErrorConflict(o.key, types.ID(0), err.Error())
	}

	iface, err = ifOp.Read(ctx, zone, iface.ID)
	if err != nil {
		return newErrorConflict(o.key, types.ID(0), err.Error())
	}

	vpcRouterInterface := &sacloud.VPCRouterInterface{}
	copySameNameField(iface, vpcRouterInterface)
	vpcRouterInterface.Index = nicIndex
	value.Interfaces = append(value.Interfaces, vpcRouterInterface)

	putVPCRouter(zone, value)
	return nil
}

// DisconnectFromSwitch is fake implementation
func (o *VPCRouterOp) DisconnectFromSwitch(ctx context.Context, zone string, id types.ID, nicIndex int) error {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return err
	}

	var exists bool
	var nicID types.ID
	var interfaces []*sacloud.VPCRouterInterface

	for _, nic := range value.Interfaces {
		if nic.Index == nicIndex {
			exists = true
			nicID = nic.ID
		} else {
			interfaces = append(interfaces, nic)
		}
	}
	if !exists {
		return newErrorBadRequest(o.key, id, fmt.Sprintf("nic[%d] is not exists", nicIndex))
	}

	ifOp := NewInterfaceOp()
	if err := ifOp.DisconnectFromSwitch(ctx, zone, nicID); err != nil {
		return newErrorConflict(o.key, types.ID(0), err.Error())
	}

	value.Interfaces = interfaces
	putVPCRouter(zone, value)
	return nil
}

// MonitorInterface is fake implementation
func (o *VPCRouterOp) MonitorInterface(ctx context.Context, zone string, id types.ID, index int, condition *sacloud.MonitorCondition) (*sacloud.InterfaceActivity, error) {
	_, err := o.Read(ctx, zone, id)
	if err != nil {
		return nil, err
	}

	now := time.Now().Truncate(time.Second)
	m := now.Minute() % 5
	if m != 0 {
		now.Add(time.Duration(m) * time.Minute)
	}

	res := &sacloud.InterfaceActivity{}
	for i := 0; i < 5; i++ {
		res.Values = append(res.Values, &sacloud.MonitorInterfaceValue{
			Time:    now.Add(time.Duration(i*-5) * time.Minute),
			Send:    float64(random(1000)),
			Receive: float64(random(1000)),
		})
	}

	return res, nil
}

// Status is fake implementation
func (o *VPCRouterOp) Status(ctx context.Context, zone string, id types.ID) (*sacloud.VPCRouterStatus, error) {
	_, err := o.Read(ctx, zone, id)
	if err != nil {
		return nil, err
	}

	return &sacloud.VPCRouterStatus{}, nil
}
<commit_msg>Fake Driver: return dummy public key for VPCRouter.Status()<commit_after>// Copyright 2016-2021 The Libsacloud 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 fake

import (
	"context"
	"fmt"
	"net"
	"time"

	"github.com/sacloud/libsacloud/v2/sacloud"
	"github.com/sacloud/libsacloud/v2/sacloud/types"
)

// Find is fake implementation
func (o *VPCRouterOp) Find(ctx context.Context, zone string, conditions *sacloud.FindCondition) (*sacloud.VPCRouterFindResult, error) {
	results, _ := find(o.key, zone, conditions)
	var values []*sacloud.VPCRouter
	for _, res := range results {
		dest := &sacloud.VPCRouter{}
		copySameNameField(res, dest)
		values = append(values, dest)
	}
	return &sacloud.VPCRouterFindResult{
		Total:      len(results),
		Count:      len(results),
		From:       0,
		VPCRouters: values,
	}, nil
}

// Create is fake implementation
func (o *VPCRouterOp) Create(ctx context.Context, zone string, param *sacloud.VPCRouterCreateRequest) (*sacloud.VPCRouter, error) {
	result := &sacloud.VPCRouter{}
	copySameNameField(param, result)
	fill(result, fillID, fillCreatedAt)

	result.Class = "vpcrouter"
	result.Availability = types.Availabilities.Migrating
	result.ZoneID = zoneIDs[zone]
	result.SettingsHash = ""
	if result.Version == 0 {
		result.Version = 2
	}

	ifOp := NewInterfaceOp()
	swOp := NewSwitchOp()

	ifCreateParam := &sacloud.InterfaceCreateRequest{}
	if param.Switch.Scope == types.Scopes.Shared {
		ifCreateParam.ServerID = result.ID
	} else {
		_, err := swOp.Read(ctx, zone, param.Switch.ID)
		if err != nil {
			return nil, newErrorConflict(o.key, types.ID(0), err.Error())
		}
	}

	iface, err := ifOp.Create(ctx, zone, ifCreateParam)
	if err != nil {
		return nil, newErrorConflict(o.key, types.ID(0), err.Error())
	}

	if param.Switch.Scope == types.Scopes.Shared {
		if err := ifOp.ConnectToSharedSegment(ctx, zone, iface.ID); err != nil {
			return nil, newErrorConflict(o.key, types.ID(0), err.Error())
		}
	} else {
		if err := ifOp.ConnectToSwitch(ctx, zone, iface.ID, param.Switch.ID); err != nil {
			return nil, newErrorConflict(o.key, types.ID(0), err.Error())
		}
	}

	iface, err = ifOp.Read(ctx, zone, iface.ID)
	if err != nil {
		return nil, newErrorConflict(o.key, types.ID(0), err.Error())
	}

	vpcRouterInterface := &sacloud.VPCRouterInterface{}
	copySameNameField(iface, vpcRouterInterface)
	if param.Switch.Scope == types.Scopes.Shared {
		sharedIP := pool().nextSharedIP()
		vpcRouterInterface.IPAddress = sharedIP.String()
		vpcRouterInterface.SubnetNetworkMaskLen = sharedSegmentSwitch.NetworkMaskLen

		ipv4Mask := net.CIDRMask(pool().SharedNetMaskLen, 32)
		vpcRouterInterface.SubnetNetworkAddress = sharedIP.Mask(ipv4Mask).String()
		vpcRouterInterface.SubnetDefaultRoute = pool().SharedDefaultGateway.String()
	}
	result.Interfaces = append(result.Interfaces, vpcRouterInterface)

	putVPCRouter(zone, result)

	id := result.ID
	startMigration(o.key, zone, func() (interface{}, error) {
		return o.Read(context.Background(), zone, id)
	})
	return result, nil
}

// Read is fake implementation
func (o *VPCRouterOp) Read(ctx context.Context, zone string, id types.ID) (*sacloud.VPCRouter, error) {
	value := getVPCRouterByID(zone, id)
	if value == nil {
		return nil, newErrorNotFound(o.key, id)
	}
	dest := &sacloud.VPCRouter{}
	copySameNameField(value, dest)
	return dest, nil
}

// Update is fake implementation
func (o *VPCRouterOp) Update(ctx context.Context, zone string, id types.ID, param *sacloud.VPCRouterUpdateRequest) (*sacloud.VPCRouter, error) {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return nil, err
	}
	copySameNameField(param, value)
	fill(value, fillModifiedAt)

	putVPCRouter(zone, value)
	return value, nil
}

// UpdateSettings is fake implementation
func (o *VPCRouterOp) UpdateSettings(ctx context.Context, zone string, id types.ID, param *sacloud.VPCRouterUpdateSettingsRequest) (*sacloud.VPCRouter, error) {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return nil, err
	}
	copySameNameField(param, value)
	fill(value, fillModifiedAt)

	putVPCRouter(zone, value)
	return value, nil
}

// Delete is fake implementation
func (o *VPCRouterOp) Delete(ctx context.Context, zone string, id types.ID) error {
	_, err := o.Read(ctx, zone, id)
	if err != nil {
		return err
	}
	ds().Delete(o.key, zone, id)
	return nil
}

// Config is fake implementation
func (o *VPCRouterOp) Config(ctx context.Context, zone string, id types.ID) error {
	return nil
}

// Boot is fake implementation
func (o *VPCRouterOp) Boot(ctx context.Context, zone string, id types.ID) error {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return err
	}
	if value.InstanceStatus.IsUp() {
		return newErrorConflict(o.key, id, "Boot is failed")
	}

	startPowerOn(o.key, zone, func() (interface{}, error) {
		return o.Read(context.Background(), zone, id)
	})

	return err
}

// Shutdown is fake implementation
func (o *VPCRouterOp) Shutdown(ctx context.Context, zone string, id types.ID, shutdownOption *sacloud.ShutdownOption) error {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return err
	}
	if !value.InstanceStatus.IsUp() {
		return newErrorConflict(o.key, id, "Shutdown is failed")
	}

	startPowerOff(o.key, zone, func() (interface{}, error) {
		return o.Read(context.Background(), zone, id)
	})

	return err
}

// Reset is fake implementation
func (o *VPCRouterOp) Reset(ctx context.Context, zone string, id types.ID) error {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return err
	}
	if !value.InstanceStatus.IsUp() {
		return newErrorConflict(o.key, id, "Reset is failed")
	}

	startPowerOn(o.key, zone, func() (interface{}, error) {
		return o.Read(context.Background(), zone, id)
	})

	return nil
}

// ConnectToSwitch is fake implementation
func (o *VPCRouterOp) ConnectToSwitch(ctx context.Context, zone string, id types.ID, nicIndex int, switchID types.ID) error {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return err
	}

	for _, nic := range value.Interfaces {
		if nic.Index == nicIndex {
			return newErrorBadRequest(o.key, id, fmt.Sprintf("nic[%d] already connected to switch", nicIndex))
		}
	}

	// find switch
	swOp := NewSwitchOp()
	_, err = swOp.Read(ctx, zone, switchID)
	if err != nil {
		return fmt.Errorf("ConnectToSwitch is failed: %s", err)
	}

	// create interface
	ifOp := NewInterfaceOp()
	iface, err := ifOp.Create(ctx, zone, &sacloud.InterfaceCreateRequest{ServerID: id})
	if err != nil {
		return newErrorConflict(o.key, types.ID(0), err.Error())
	}

	if err := ifOp.ConnectToSwitch(ctx, zone, iface.ID, switchID); err != nil {
		return newErrorConflict(o.key, types.ID(0), err.Error())
	}

	iface, err = ifOp.Read(ctx, zone, iface.ID)
	if err != nil {
		return newErrorConflict(o.key, types.ID(0), err.Error())
	}

	vpcRouterInterface := &sacloud.VPCRouterInterface{}
	copySameNameField(iface, vpcRouterInterface)
	vpcRouterInterface.Index = nicIndex
	value.Interfaces = append(value.Interfaces, vpcRouterInterface)

	putVPCRouter(zone, value)
	return nil
}

// DisconnectFromSwitch is fake implementation
func (o *VPCRouterOp) DisconnectFromSwitch(ctx context.Context, zone string, id types.ID, nicIndex int) error {
	value, err := o.Read(ctx, zone, id)
	if err != nil {
		return err
	}

	var exists bool
	var nicID types.ID
	var interfaces []*sacloud.VPCRouterInterface

	for _, nic := range value.Interfaces {
		if nic.Index == nicIndex {
			exists = true
			nicID = nic.ID
		} else {
			interfaces = append(interfaces, nic)
		}
	}
	if !exists {
		return newErrorBadRequest(o.key, id, fmt.Sprintf("nic[%d] is not exists", nicIndex))
	}

	ifOp := NewInterfaceOp()
	if err := ifOp.DisconnectFromSwitch(ctx, zone, nicID); err != nil {
		return newErrorConflict(o.key, types.ID(0), err.Error())
	}

	value.Interfaces = interfaces
	putVPCRouter(zone, value)
	return nil
}

// MonitorInterface is fake implementation
func (o *VPCRouterOp) MonitorInterface(ctx context.Context, zone string, id types.ID, index int, condition *sacloud.MonitorCondition) (*sacloud.InterfaceActivity, error) {
	_, err := o.Read(ctx, zone, id)
	if err != nil {
		return nil, err
	}

	now := time.Now().Truncate(time.Second)
	m := now.Minute() % 5
	if m != 0 {
		now.Add(time.Duration(m) * time.Minute)
	}

	res := &sacloud.InterfaceActivity{}
	for i := 0; i < 5; i++ {
		res.Values = append(res.Values, &sacloud.MonitorInterfaceValue{
			Time:    now.Add(time.Duration(i*-5) * time.Minute),
			Send:    float64(random(1000)),
			Receive: float64(random(1000)),
		})
	}

	return res, nil
}

// Status is fake implementation
func (o *VPCRouterOp) Status(ctx context.Context, zone string, id types.ID) (*sacloud.VPCRouterStatus, error) {
	v, err := o.Read(ctx, zone, id)
	if err != nil {
		return nil, err
	}

	if v.InstanceStatus.IsUp() && v.Settings.WireGuardEnabled.Bool() {
		return &sacloud.VPCRouterStatus{
			WireGuard: &sacloud.WireGuardStatus{
				PublicKey: "fake-public-key",
			},
		}, nil
	}
	return &sacloud.VPCRouterStatus{}, nil
}
<|endoftext|>


<commit_before>package handlers

import (
	"net/http"

	"github.com/play-with-docker/play-with-docker/config"
)

type CookieID struct {
	Id         string `json:"id"`
	UserName   string `json:"user_name"`
	UserAvatar string `json:"user_avatar"`
	ProviderId string `json:"provider_id"`
}

func (c *CookieID) SetCookie(rw http.ResponseWriter, host string) error {
	if encoded, err := config.SecureCookie.Encode("id", c); err == nil {
		cookie := &http.Cookie{
			Name:     "id",
			Value:    encoded,
			Domain:   host,
			Path:     "/",
			Secure:   config.UseLetsEncrypt,
			HttpOnly: true,
		}
		http.SetCookie(rw, cookie)
	} else {
		return err
	}
	return nil
}
func ReadCookie(r *http.Request) (*CookieID, error) {
	if cookie, err := r.Cookie("id"); err == nil {
		value := &CookieID{}
		if err = config.SecureCookie.Decode("id", cookie.Value, &value); err == nil {
			return value, nil
		} else {
			return nil, err
		}
	} else {
		return nil, err
	}
}
<commit_msg>Add SameSite cookie option<commit_after>package handlers

import (
	"net/http"

	"github.com/play-with-docker/play-with-docker/config"
)

type CookieID struct {
	Id         string `json:"id"`
	UserName   string `json:"user_name"`
	UserAvatar string `json:"user_avatar"`
	ProviderId string `json:"provider_id"`
}

func (c *CookieID) SetCookie(rw http.ResponseWriter, host string) error {
	if encoded, err := config.SecureCookie.Encode("id", c); err == nil {
		cookie := &http.Cookie{
			Name:     "id",
			Value:    encoded,
			Domain:   host,
			Path:     "/",
			SameSite: http.SameSiteNoneMode,
			Secure:   config.UseLetsEncrypt,
			HttpOnly: true,
		}
		http.SetCookie(rw, cookie)
	} else {
		return err
	}
	return nil
}
func ReadCookie(r *http.Request) (*CookieID, error) {
	if cookie, err := r.Cookie("id"); err == nil {
		value := &CookieID{}
		if err = config.SecureCookie.Decode("id", cookie.Value, &value); err == nil {
			return value, nil
		} else {
			return nil, err
		}
	} else {
		return nil, err
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// 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 scorecard

import (
	"context"
	"fmt"

	"github.com/forseti-security/config-validator/pkg/api/validator"
	"github.com/forseti-security/config-validator/pkg/gcv"
	"github.com/pkg/errors"
)

// ScoringConfig holds settings for generating a score
type ScoringConfig struct {
	PolicyPath  string
	categories  map[string]*constraintCategory
	constraints map[string]*constraintViolations
	validator   *gcv.Validator
	ctx         context.Context
}

const otherCategoryKey = "other"

// constraintCategory holds constraints by category
type constraintCategory struct {
	Name        string
	constraints []*constraintViolations
}

func (c constraintCategory) Count() int {
	sum := 0
	for _, cv := range c.constraints {
		sum += cv.Count()
	}
	return sum
}

// constraintViolations holds violations for a particular constraint
type constraintViolations struct {
	constraint *validator.Constraint
	Violations []*validator.Violation `protobuf:"bytes,1,rep,name=violations,proto3" json:"violations,omitempty"`
}

func (cv constraintViolations) Count() int {
	return len(cv.Violations)
}

func (cv constraintViolations) GetName() string {
	return cv.constraint.GetMetadata().GetName()
}

var availableCategories = map[string]string{
	"operational-efficiency": "Operational Efficiency",
	"security":               "Security",
	"reliability":            "Reliability",
	otherCategoryKey:         "Other",
}

func getConstraintForViolation(config *ScoringConfig, violation *validator.Violation) (*constraintViolations, error) {
	key := violation.GetConstraint()
	cv, found := config.constraints[key]
	if !found {
		response, err := config.validator.GetConstraint(config.ctx, &validator.GetConstraintRequest{
			Name: key,
		})
		if err != nil {
			return nil, errors.Wrap(err, "Finding matching constraint")
		}

		constraint := response.GetConstraint()
		cv = &constraintViolations{
			constraint: response.GetConstraint(),
		}
		config.constraints[key] = cv

		annotations := constraint.GetMetadata().GetAnnotations()
		categoryKey, found := annotations["bundles.validator.forsetisecurity.org/scorecard-v1"]
		if !found {
			categoryKey = otherCategoryKey
		}

		category, found := config.categories[categoryKey]
		if !found {
			return nil, fmt.Errorf("Unknown constraint category %v for constraint %v", categoryKey, key)
		}
		category.constraints = append(category.constraints, cv)
	}
	return cv, nil
}

// attachViolations puts violations into their appropriate categories
func attachViolations(audit *validator.AuditResponse, config *ScoringConfig) error {
	// Build map of categories
	config.categories = make(map[string]*constraintCategory)
	for k, name := range availableCategories {
		config.categories[k] = &constraintCategory{
			Name: name,
		}
	}

	// Categorize violations
	config.constraints = make(map[string]*constraintViolations)
	for _, v := range audit.Violations {
		cv, err := getConstraintForViolation(config, v)
		if err != nil {
			return errors.Wrap(err, "Categorizing violation")
		}

		cv.Violations = append(cv.Violations, v)
	}

	return nil
}

// ScoreInventory creates a Scorecard for an inventory
func ScoreInventory(inventory *Inventory, config *ScoringConfig) error {
	config.ctx = context.Background()

	err := attachValidator(config)
	if err != nil {
		return errors.Wrap(err, "initializing gcv validator")
	}

	auditResult, err := getViolations(inventory, config)
	if err != nil {
		return err
	}

	err = attachViolations(auditResult, config)

	if len(auditResult.Violations) > 0 {
		fmt.Printf("\n\n%v total issues found\n", len(auditResult.Violations))
		for _, category := range config.categories {
			fmt.Printf("\n\n%v: %v issues found\n", category.Name, category.Count())
			fmt.Printf("----------\n")
			for _, cv := range category.constraints {
				fmt.Printf("%v: %v issues\n", cv.GetName(), cv.Count())
				for _, v := range cv.Violations {
					fmt.Printf("- %v\n\n",
						v.Message,
					)
					Log.Debug("Violation metadata", "metadata", v.GetMetadata())
				}
			}
		}
	} else {
		fmt.Println("No issues found found! You have a perfect score.")
	}

	return nil
}
<commit_msg>Switch to directly grabbing constraint metadata from violation<commit_after>// Copyright 2019 Google LLC
//
// 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 scorecard

import (
	"context"
	"fmt"

	"github.com/forseti-security/config-validator/pkg/api/validator"
	"github.com/forseti-security/config-validator/pkg/gcv"
	"github.com/pkg/errors"
)

// ScoringConfig holds settings for generating a score
type ScoringConfig struct {
	PolicyPath  string
	categories  map[string]*constraintCategory
	constraints map[string]*constraintViolations
	validator   *gcv.Validator
	ctx         context.Context
}

const otherCategoryKey = "other"

// constraintCategory holds constraints by category
type constraintCategory struct {
	Name        string
	constraints []*constraintViolations
}

func (c constraintCategory) Count() int {
	sum := 0
	for _, cv := range c.constraints {
		sum += cv.Count()
	}
	return sum
}

// constraintViolations holds violations for a particular constraint
type constraintViolations struct {
	constraint *validator.Constraint
	Violations []*validator.Violation `protobuf:"bytes,1,rep,name=violations,proto3" json:"violations,omitempty"`
}

func (cv constraintViolations) Count() int {
	return len(cv.Violations)
}

func (cv constraintViolations) GetName() string {
	return cv.constraint.GetMetadata().GetStructValue().GetFields()["name"].GetStringValue()
}

var availableCategories = map[string]string{
	"operational-efficiency": "Operational Efficiency",
	"security":               "Security",
	"reliability":            "Reliability",
	otherCategoryKey:         "Other",
}

func getConstraintForViolation(config *ScoringConfig, violation *validator.Violation) (*constraintViolations, error) {
	key := violation.GetConstraint()
	cv, found := config.constraints[key]
	if !found {
		constraint := violation.GetConstraintConfig()
		cv = &constraintViolations{
			constraint: constraint,
		}
		config.constraints[key] = cv

		metadata := constraint.GetMetadata()
		annotations := metadata.GetStructValue().GetFields()["annotations"].GetStructValue().GetFields()

		categoryKey := otherCategoryKey
		categoryValue, found := annotations["bundles.validator.forsetisecurity.org/scorecard-v1"]
		if found {
			categoryKey = categoryValue.GetStringValue()
		}

		category, found := config.categories[categoryKey]
		if !found {
			return nil, fmt.Errorf("Unknown constraint category %v for constraint %v", categoryKey, key)
		}
		category.constraints = append(category.constraints, cv)
	}
	return cv, nil
}

// attachViolations puts violations into their appropriate categories
func attachViolations(audit *validator.AuditResponse, config *ScoringConfig) error {
	// Build map of categories
	config.categories = make(map[string]*constraintCategory)
	for k, name := range availableCategories {
		config.categories[k] = &constraintCategory{
			Name: name,
		}
	}

	// Categorize violations
	config.constraints = make(map[string]*constraintViolations)
	for _, v := range audit.Violations {
		cv, err := getConstraintForViolation(config, v)
		if err != nil {
			return errors.Wrap(err, "Categorizing violation")
		}

		cv.Violations = append(cv.Violations, v)
	}

	return nil
}

// ScoreInventory creates a Scorecard for an inventory
func ScoreInventory(inventory *Inventory, config *ScoringConfig) error {
	config.ctx = context.Background()

	err := attachValidator(config)
	if err != nil {
		return errors.Wrap(err, "initializing gcv validator")
	}

	auditResult, err := getViolations(inventory, config)
	if err != nil {
		return err
	}

	err = attachViolations(auditResult, config)

	if len(auditResult.Violations) > 0 {
		fmt.Printf("\n\n%v total issues found\n", len(auditResult.Violations))
		for _, category := range config.categories {
			fmt.Printf("\n\n%v: %v issues found\n", category.Name, category.Count())
			fmt.Printf("----------\n")
			for _, cv := range category.constraints {
				fmt.Printf("%v: %v issues\n", cv.GetName(), cv.Count())
				for _, v := range cv.Violations {
					fmt.Printf("- %v\n\n",
						v.Message,
					)
					Log.Debug("Violation metadata", "metadata", v.GetMetadata())
				}
			}
		}
	} else {
		fmt.Println("No issues found found! You have a perfect score.")
	}

	return nil
}
<|endoftext|>


<commit_before>package consumer

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/stealthly/siesta"
)

type PartitionConsumer struct {
	config       PartitionConsumerConfig
	kafkaClient  siesta.Connector
	fetchers     map[string]map[int32]*FetcherState
	fetchersLock sync.Mutex
}

type PartitionConsumerConfig struct {
	// Consumer group
	Group string

	//Interval to commit offsets at
	CommitInterval time.Duration

	// BrokerList is a bootstrap list to discover other brokers in a cluster. At least one broker is required.
	BrokerList []string

	// ReadTimeout is a timeout to read the response from a TCP socket.
	ReadTimeout time.Duration

	// WriteTimeout is a timeout to write the request to a TCP socket.
	WriteTimeout time.Duration

	// ConnectTimeout is a timeout to connect to a TCP socket.
	ConnectTimeout time.Duration

	// Sets whether the connection should be kept alive.
	KeepAlive bool

	// A keep alive period for a TCP connection.
	KeepAliveTimeout time.Duration

	// Maximum number of open connections for a connector.
	MaxConnections int

	// Maximum number of open connections for a single broker for a connector.
	MaxConnectionsPerBroker int

	// Maximum fetch size in bytes which will be used in all Consume() calls.
	FetchSize int32

	// The minimum amount of data the server should return for a fetch request. If insufficient data is available the request will block
	FetchMinBytes int32

	// The maximum amount of time the server will block before answering the fetch request if there isn't sufficient data to immediately satisfy FetchMinBytes
	FetchMaxWaitTime int32

	// Number of retries to get topic metadata.
	MetadataRetries int

	// Backoff value between topic metadata requests.
	MetadataBackoff time.Duration

	// Number of retries to commit an offset.
	CommitOffsetRetries int

	// Backoff value between commit offset requests.
	CommitOffsetBackoff time.Duration

	// Number of retries to get consumer metadata.
	ConsumerMetadataRetries int

	// Backoff value between consumer metadata requests.
	ConsumerMetadataBackoff time.Duration

	// ClientID that will be used by a connector to identify client requests by broker.
	ClientID string
}

func NewPartitionConsumerConfig(group string) *PartitionConsumerConfig {
	return &PartitionConsumerConfig{
		Group:                   group,
		CommitInterval:          1 * time.Second,
		ReadTimeout:             5 * time.Second,
		WriteTimeout:            5 * time.Second,
		ConnectTimeout:          5 * time.Second,
		KeepAlive:               true,
		KeepAliveTimeout:        1 * time.Minute,
		MaxConnections:          5,
		MaxConnectionsPerBroker: 5,
		FetchSize:               1024000,
		FetchMaxWaitTime:        1000,
		MetadataRetries:         5,
		MetadataBackoff:         200 * time.Millisecond,
		CommitOffsetRetries:     5,
		CommitOffsetBackoff:     200 * time.Millisecond,
		ConsumerMetadataRetries: 15,
		ConsumerMetadataBackoff: 500 * time.Millisecond,
		ClientID:                "partition-consumer",
	}
}

func NewPartitionConsumer(consumerConfig PartitionConsumerConfig) *PartitionConsumer {
	connectorConfig := siesta.NewConnectorConfig()
	connectorConfig.BrokerList = consumerConfig.BrokerList
	connectorConfig.ClientID = consumerConfig.ClientID
	connectorConfig.CommitOffsetBackoff = consumerConfig.CommitOffsetBackoff
	connectorConfig.CommitOffsetRetries = consumerConfig.CommitOffsetRetries
	connectorConfig.ConnectTimeout = consumerConfig.ConnectTimeout
	connectorConfig.ConsumerMetadataBackoff = consumerConfig.ConsumerMetadataBackoff
	connectorConfig.ConsumerMetadataRetries = consumerConfig.ConsumerMetadataRetries
	connectorConfig.FetchMaxWaitTime = consumerConfig.FetchMaxWaitTime
	connectorConfig.FetchMinBytes = consumerConfig.FetchMinBytes
	connectorConfig.FetchSize = consumerConfig.FetchSize
	connectorConfig.KeepAlive = consumerConfig.KeepAlive
	connectorConfig.KeepAliveTimeout = consumerConfig.KeepAliveTimeout
	connectorConfig.MaxConnections = consumerConfig.MaxConnections
	connectorConfig.MaxConnectionsPerBroker = consumerConfig.MaxConnectionsPerBroker
	connectorConfig.MetadataBackoff = consumerConfig.MetadataBackoff
	connectorConfig.MetadataRetries = consumerConfig.MetadataRetries
	connectorConfig.ReadTimeout = consumerConfig.ReadTimeout
	connectorConfig.WriteTimeout = consumerConfig.WriteTimeout
	kafkaClient, err := siesta.NewDefaultConnector(connectorConfig)
	if err != nil {
		panic(err)
	}

	consumer := &PartitionConsumer{
		config:      consumerConfig,
		kafkaClient: kafkaClient,
		fetchers:    make(map[string]map[int32]*FetcherState),
	}

	commitTimer := time.NewTimer(consumerConfig.CommitInterval)
	go func() {
		for {
			select {
			case <-commitTimer.C:
				{

					for topic, partitions := range consumer.fetchers {
						for partition, fetcherState := range partitions {
							offsetToCommit := fetcherState.GetOffset()
							if offsetToCommit > fetcherState.LastCommitted {
								err := consumer.kafkaClient.CommitOffset(consumer.config.Group, topic, partition, offsetToCommit)
								if err != nil {
									fmt.Printf("Failed to commit offset: %s\n", err.Error())
								}
							}
							if fetcherState.Removed {
								inLock(&consumer.fetchersLock, func() {
									if consumer.fetchers[topic][partition].Removed {
										delete(consumer.fetchers[topic], partition)
									}
								})
							}
						}
					}
					commitTimer.Reset(consumerConfig.CommitInterval)
				}
			}
		}
	}()

	return consumer
}

func (this *PartitionConsumer) Add(topic string, partition int32, strategy Strategy) error {
	fmt.Printf("Adding new topic partition: %s, %d\n", topic, partition)
	if _, exists := this.fetchers[topic]; !exists {
		fmt.Printf("Creating partition hash this.fetchers[%s]\n", topic)
		this.fetchers[topic] = make(map[int32]*FetcherState)
	}
	var fetcherState *FetcherState
	inLock(&this.fetchersLock, func() {
		fmt.Println("We are in lock!")
		if _, exists := this.fetchers[topic][partition]; !exists || this.fetchers[topic][partition].Removed {
			fmt.Printf("Not exists or removed (%s)", exists)
			if !exists {
				fmt.Println("Not exists! Getting offset from kafka...")
				offset, err := this.kafkaClient.GetOffset(this.config.Group, topic, partition)
				fmt.Printf("Offset received: %d\n", offset)
				if err != nil {
					//It's not critical, since offsets have not been committed yet
					fmt.Printf("Error fetching topic metadata: %s\n", err.Error())
				}
				fetcherState = NewFetcherState(offset)
				fmt.Printf("Fetcher state received: %v\n", fetcherState)
				this.fetchers[topic][partition] = fetcherState
			} else {
				fmt.Println("Was Removed, now setting Removed to false")
				this.fetchers[topic][partition].Removed = false
			}
		}
	})

	if fetcherState == nil {
		fmt.Println("Fetcher state is nil! Yikes! Returning...")
		return nil
	}

	go func() {
		fmt.Printf("Start fetching cycle for %s, %d\n", topic, partition)
		for {
			response, err := this.kafkaClient.Fetch(topic, partition, fetcherState.GetOffset()+1)
			if err != nil {
				fmt.Printf("Kafka error: %s\n", err.Error())
				continue
			}

			select {
			case fetcherState.Removed = <-fetcherState.stopChannel:
				{
					fmt.Printf("Stop fetching cycle for %s, %d\n", topic, partition)
					break
				}
			default:
				{
					if _, exists := response.Data[topic]; !exists {
						continue
					}
					if _, exists := response.Data[topic][partition]; !exists {
						continue
					}
					if len(response.Data[topic][partition].Messages) == 0 {
						continue
					}

					fmt.Printf("Sending messages to strategy %s, %d, %v\n", topic, partition, response.Data[topic][partition].Messages)
					err = strategy(topic, partition, response.Data[topic][partition].Messages)
					if err != nil {
						fmt.Printf("Strategy error: %s\n", err.Error())
					}

					offsetIndex := len(response.Data[topic][partition].Messages) - 1
					offsetValue := response.Data[topic][partition].Messages[offsetIndex].Offset
					fetcherState.SetOffset(offsetValue)
				}
			}
		}
	}()

	return nil
}

func (this *PartitionConsumer) Remove(topic string, partition int32) {
	if topicFetchers, exists := this.fetchers[topic]; exists {
		if fetcherState, exists := topicFetchers[partition]; exists {
			fetcherState.GetStopChannel() <- true
		}
	}
}

func (this *PartitionConsumer) GetTopicPartitions() *TopicAndPartitionSet {
	tpSet := NewTopicAndPartitionSet()
	for topic, partitions := range this.fetchers {
		for partition, _ := range partitions {
			tpSet.Add(TopicAndPartition{topic, partition})
		}
	}

	return tpSet
}

type FetcherState struct {
	LastCommitted int64
	Removed       bool
	offset        int64
	stopChannel   chan bool
}

func NewFetcherState(initialOffset int64) *FetcherState {
	return &FetcherState{
		LastCommitted: initialOffset,
		offset:        initialOffset,
		stopChannel:   make(chan bool),
	}
}

func (this *FetcherState) GetStopChannel() chan<- bool {
	return this.stopChannel
}

func (this *FetcherState) GetOffset() int64 {
	return atomic.LoadInt64(&this.offset)
}

func (this *FetcherState) SetOffset(offset int64) {
	atomic.StoreInt64(&this.offset, offset)
}

type Strategy func(topic string, partition int32, messages []*siesta.MessageAndOffset) error
<commit_msg>Added error check for topic-partition data<commit_after>package consumer

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/stealthly/siesta"
)

type PartitionConsumer struct {
	config       PartitionConsumerConfig
	kafkaClient  siesta.Connector
	fetchers     map[string]map[int32]*FetcherState
	fetchersLock sync.Mutex
}

type PartitionConsumerConfig struct {
	// Consumer group
	Group string

	//Interval to commit offsets at
	CommitInterval time.Duration

	// BrokerList is a bootstrap list to discover other brokers in a cluster. At least one broker is required.
	BrokerList []string

	// ReadTimeout is a timeout to read the response from a TCP socket.
	ReadTimeout time.Duration

	// WriteTimeout is a timeout to write the request to a TCP socket.
	WriteTimeout time.Duration

	// ConnectTimeout is a timeout to connect to a TCP socket.
	ConnectTimeout time.Duration

	// Sets whether the connection should be kept alive.
	KeepAlive bool

	// A keep alive period for a TCP connection.
	KeepAliveTimeout time.Duration

	// Maximum number of open connections for a connector.
	MaxConnections int

	// Maximum number of open connections for a single broker for a connector.
	MaxConnectionsPerBroker int

	// Maximum fetch size in bytes which will be used in all Consume() calls.
	FetchSize int32

	// The minimum amount of data the server should return for a fetch request. If insufficient data is available the request will block
	FetchMinBytes int32

	// The maximum amount of time the server will block before answering the fetch request if there isn't sufficient data to immediately satisfy FetchMinBytes
	FetchMaxWaitTime int32

	// Number of retries to get topic metadata.
	MetadataRetries int

	// Backoff value between topic metadata requests.
	MetadataBackoff time.Duration

	// Number of retries to commit an offset.
	CommitOffsetRetries int

	// Backoff value between commit offset requests.
	CommitOffsetBackoff time.Duration

	// Number of retries to get consumer metadata.
	ConsumerMetadataRetries int

	// Backoff value between consumer metadata requests.
	ConsumerMetadataBackoff time.Duration

	// ClientID that will be used by a connector to identify client requests by broker.
	ClientID string
}

func NewPartitionConsumerConfig(group string) *PartitionConsumerConfig {
	return &PartitionConsumerConfig{
		Group:                   group,
		CommitInterval:          1 * time.Second,
		ReadTimeout:             5 * time.Second,
		WriteTimeout:            5 * time.Second,
		ConnectTimeout:          5 * time.Second,
		KeepAlive:               true,
		KeepAliveTimeout:        1 * time.Minute,
		MaxConnections:          5,
		MaxConnectionsPerBroker: 5,
		FetchSize:               1024000,
		FetchMaxWaitTime:        1000,
		MetadataRetries:         5,
		MetadataBackoff:         200 * time.Millisecond,
		CommitOffsetRetries:     5,
		CommitOffsetBackoff:     200 * time.Millisecond,
		ConsumerMetadataRetries: 15,
		ConsumerMetadataBackoff: 500 * time.Millisecond,
		ClientID:                "partition-consumer",
	}
}

func NewPartitionConsumer(consumerConfig PartitionConsumerConfig) *PartitionConsumer {
	connectorConfig := siesta.NewConnectorConfig()
	connectorConfig.BrokerList = consumerConfig.BrokerList
	connectorConfig.ClientID = consumerConfig.ClientID
	connectorConfig.CommitOffsetBackoff = consumerConfig.CommitOffsetBackoff
	connectorConfig.CommitOffsetRetries = consumerConfig.CommitOffsetRetries
	connectorConfig.ConnectTimeout = consumerConfig.ConnectTimeout
	connectorConfig.ConsumerMetadataBackoff = consumerConfig.ConsumerMetadataBackoff
	connectorConfig.ConsumerMetadataRetries = consumerConfig.ConsumerMetadataRetries
	connectorConfig.FetchMaxWaitTime = consumerConfig.FetchMaxWaitTime
	connectorConfig.FetchMinBytes = consumerConfig.FetchMinBytes
	connectorConfig.FetchSize = consumerConfig.FetchSize
	connectorConfig.KeepAlive = consumerConfig.KeepAlive
	connectorConfig.KeepAliveTimeout = consumerConfig.KeepAliveTimeout
	connectorConfig.MaxConnections = consumerConfig.MaxConnections
	connectorConfig.MaxConnectionsPerBroker = consumerConfig.MaxConnectionsPerBroker
	connectorConfig.MetadataBackoff = consumerConfig.MetadataBackoff
	connectorConfig.MetadataRetries = consumerConfig.MetadataRetries
	connectorConfig.ReadTimeout = consumerConfig.ReadTimeout
	connectorConfig.WriteTimeout = consumerConfig.WriteTimeout
	kafkaClient, err := siesta.NewDefaultConnector(connectorConfig)
	if err != nil {
		panic(err)
	}

	consumer := &PartitionConsumer{
		config:      consumerConfig,
		kafkaClient: kafkaClient,
		fetchers:    make(map[string]map[int32]*FetcherState),
	}

	commitTimer := time.NewTimer(consumerConfig.CommitInterval)
	go func() {
		for {
			select {
			case <-commitTimer.C:
				{

					for topic, partitions := range consumer.fetchers {
						for partition, fetcherState := range partitions {
							offsetToCommit := fetcherState.GetOffset()
							if offsetToCommit > fetcherState.LastCommitted {
								err := consumer.kafkaClient.CommitOffset(consumer.config.Group, topic, partition, offsetToCommit)
								if err != nil {
									fmt.Printf("Failed to commit offset: %s\n", err.Error())
								}
							}
							if fetcherState.Removed {
								inLock(&consumer.fetchersLock, func() {
									if consumer.fetchers[topic][partition].Removed {
										delete(consumer.fetchers[topic], partition)
									}
								})
							}
						}
					}
					commitTimer.Reset(consumerConfig.CommitInterval)
				}
			}
		}
	}()

	return consumer
}

func (this *PartitionConsumer) Add(topic string, partition int32, strategy Strategy) error {
	fmt.Printf("Adding new topic partition: %s, %d\n", topic, partition)
	if _, exists := this.fetchers[topic]; !exists {
		fmt.Printf("Creating partition hash this.fetchers[%s]\n", topic)
		this.fetchers[topic] = make(map[int32]*FetcherState)
	}
	var fetcherState *FetcherState
	inLock(&this.fetchersLock, func() {
		fmt.Println("We are in lock!")
		if _, exists := this.fetchers[topic][partition]; !exists || this.fetchers[topic][partition].Removed {
			fmt.Printf("Not exists or removed (%s)", exists)
			if !exists {
				fmt.Println("Not exists! Getting offset from kafka...")
				offset, err := this.kafkaClient.GetOffset(this.config.Group, topic, partition)
				fmt.Printf("Offset received: %d\n", offset)
				if err != nil {
					//It's not critical, since offsets have not been committed yet
					fmt.Printf("Error fetching topic metadata: %s\n", err.Error())
				}
				fetcherState = NewFetcherState(offset)
				fmt.Printf("Fetcher state received: %v\n", fetcherState)
				this.fetchers[topic][partition] = fetcherState
			} else {
				fmt.Println("Was Removed, now setting Removed to false")
				this.fetchers[topic][partition].Removed = false
			}
		}
	})

	if fetcherState == nil {
		fmt.Println("Fetcher state is nil! Yikes! Returning...")
		return nil
	}

	go func() {
		fmt.Printf("Start fetching cycle for %s, %d\n", topic, partition)
		for {
			response, err := this.kafkaClient.Fetch(topic, partition, fetcherState.GetOffset()+1)
			if err != nil {
				fmt.Printf("Kafka error: %s\n", err.Error())
				continue
			}

			select {
			case fetcherState.Removed = <-fetcherState.stopChannel:
				{
					fmt.Printf("Stop fetching cycle for %s, %d\n", topic, partition)
					break
				}
			default:
				{
					if _, exists := response.Data[topic]; !exists {
						continue
					}
					if _, exists := response.Data[topic][partition]; !exists {
						continue
					}
					if response.Data[topic][partition].Error != siesta.ErrNoError {
						fmt.Printf("Got error for topic %s and partition %d: %s", topic, partition, response.Data[topic][partition].Error)
						continue
					}

					if len(response.Data[topic][partition].Messages) == 0 {
						continue
					}

					fmt.Printf("Sending messages to strategy %s, %d, %v\n", topic, partition, response.Data[topic][partition].Messages)
					err = strategy(topic, partition, response.Data[topic][partition].Messages)
					if err != nil {
						fmt.Printf("Strategy error: %s\n", err.Error())
					}

					offsetIndex := len(response.Data[topic][partition].Messages) - 1
					offsetValue := response.Data[topic][partition].Messages[offsetIndex].Offset
					fetcherState.SetOffset(offsetValue)
				}
			}
		}
	}()

	return nil
}

func (this *PartitionConsumer) Remove(topic string, partition int32) {
	if topicFetchers, exists := this.fetchers[topic]; exists {
		if fetcherState, exists := topicFetchers[partition]; exists {
			fetcherState.GetStopChannel() <- true
		}
	}
}

func (this *PartitionConsumer) GetTopicPartitions() *TopicAndPartitionSet {
	tpSet := NewTopicAndPartitionSet()
	for topic, partitions := range this.fetchers {
		for partition, _ := range partitions {
			tpSet.Add(TopicAndPartition{topic, partition})
		}
	}

	return tpSet
}

type FetcherState struct {
	LastCommitted int64
	Removed       bool
	offset        int64
	stopChannel   chan bool
}

func NewFetcherState(initialOffset int64) *FetcherState {
	return &FetcherState{
		LastCommitted: initialOffset,
		offset:        initialOffset,
		stopChannel:   make(chan bool),
	}
}

func (this *FetcherState) GetStopChannel() chan<- bool {
	return this.stopChannel
}

func (this *FetcherState) GetOffset() int64 {
	return atomic.LoadInt64(&this.offset)
}

func (this *FetcherState) SetOffset(offset int64) {
	atomic.StoreInt64(&this.offset, offset)
}

type Strategy func(topic string, partition int32, messages []*siesta.MessageAndOffset) error
<|endoftext|>


<commit_before>package udp

import (
	"context"
	"sync"
	"time"

	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/app/log"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/transport/ray"
)

type ResponseCallback func(payload *buf.Buffer)

type connEntry struct {
	inbound ray.InboundRay
	timer   signal.ActivityUpdater
	cancel  context.CancelFunc
}

type Dispatcher struct {
	sync.RWMutex
	conns      map[net.Destination]*connEntry
	dispatcher dispatcher.Interface
}

func NewDispatcher(dispatcher dispatcher.Interface) *Dispatcher {
	return &Dispatcher{
		conns:      make(map[net.Destination]*connEntry),
		dispatcher: dispatcher,
	}
}

func (v *Dispatcher) RemoveRay(dest net.Destination) {
	v.Lock()
	defer v.Unlock()
	if conn, found := v.conns[dest]; found {
		conn.inbound.InboundInput().Close()
		conn.inbound.InboundOutput().Close()
		delete(v.conns, dest)
	}
}

func (v *Dispatcher) getInboundRay(dest net.Destination, callback ResponseCallback) *connEntry {
	v.Lock()
	defer v.Unlock()

	if entry, found := v.conns[dest]; found {
		return entry
	}

	log.Trace(newError("establishing new connection for ", dest))

	ctx, cancel := context.WithCancel(context.Background())
	removeRay := func() {
		cancel()
		v.RemoveRay(dest)
	}
	timer := signal.CancelAfterInactivity(ctx, removeRay, time.Second*4)
	inboundRay, _ := v.dispatcher.Dispatch(ctx, dest)
	entry := &connEntry{
		inbound: inboundRay,
		timer:   timer,
		cancel:  removeRay,
	}
	v.conns[dest] = entry
	go handleInput(ctx, entry, callback)
	return entry
}

func (v *Dispatcher) Dispatch(ctx context.Context, destination net.Destination, payload *buf.Buffer, callback ResponseCallback) {
	// TODO: Add user to destString
	log.Trace(newError("dispatch request to: ", destination).AtDebug())

	conn := v.getInboundRay(destination, callback)
	outputStream := conn.inbound.InboundInput()
	if outputStream != nil {
		if err := outputStream.WriteMultiBuffer(buf.NewMultiBufferValue(payload)); err != nil {
			log.Trace(newError("failed to write first UDP payload").Base(err))
			conn.cancel()
			return
		}
	}
}

func handleInput(ctx context.Context, conn *connEntry, callback ResponseCallback) {
	input := conn.inbound.InboundOutput()
	timer := conn.timer

	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		mb, err := input.ReadMultiBuffer()
		if err != nil {
			log.Trace(newError("failed to handl UDP input").Base(err))
			conn.cancel()
			return
		}
		timer.Update()
		for _, b := range mb {
			callback(b)
		}
	}
}
<commit_msg>Fix typo, fix #715<commit_after>package udp

import (
	"context"
	"sync"
	"time"

	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/app/log"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/transport/ray"
)

type ResponseCallback func(payload *buf.Buffer)

type connEntry struct {
	inbound ray.InboundRay
	timer   signal.ActivityUpdater
	cancel  context.CancelFunc
}

type Dispatcher struct {
	sync.RWMutex
	conns      map[net.Destination]*connEntry
	dispatcher dispatcher.Interface
}

func NewDispatcher(dispatcher dispatcher.Interface) *Dispatcher {
	return &Dispatcher{
		conns:      make(map[net.Destination]*connEntry),
		dispatcher: dispatcher,
	}
}

func (v *Dispatcher) RemoveRay(dest net.Destination) {
	v.Lock()
	defer v.Unlock()
	if conn, found := v.conns[dest]; found {
		conn.inbound.InboundInput().Close()
		conn.inbound.InboundOutput().Close()
		delete(v.conns, dest)
	}
}

func (v *Dispatcher) getInboundRay(dest net.Destination, callback ResponseCallback) *connEntry {
	v.Lock()
	defer v.Unlock()

	if entry, found := v.conns[dest]; found {
		return entry
	}

	log.Trace(newError("establishing new connection for ", dest))

	ctx, cancel := context.WithCancel(context.Background())
	removeRay := func() {
		cancel()
		v.RemoveRay(dest)
	}
	timer := signal.CancelAfterInactivity(ctx, removeRay, time.Second*4)
	inboundRay, _ := v.dispatcher.Dispatch(ctx, dest)
	entry := &connEntry{
		inbound: inboundRay,
		timer:   timer,
		cancel:  removeRay,
	}
	v.conns[dest] = entry
	go handleInput(ctx, entry, callback)
	return entry
}

func (v *Dispatcher) Dispatch(ctx context.Context, destination net.Destination, payload *buf.Buffer, callback ResponseCallback) {
	// TODO: Add user to destString
	log.Trace(newError("dispatch request to: ", destination).AtDebug())

	conn := v.getInboundRay(destination, callback)
	outputStream := conn.inbound.InboundInput()
	if outputStream != nil {
		if err := outputStream.WriteMultiBuffer(buf.NewMultiBufferValue(payload)); err != nil {
			log.Trace(newError("failed to write first UDP payload").Base(err))
			conn.cancel()
			return
		}
	}
}

func handleInput(ctx context.Context, conn *connEntry, callback ResponseCallback) {
	input := conn.inbound.InboundOutput()
	timer := conn.timer

	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		mb, err := input.ReadMultiBuffer()
		if err != nil {
			log.Trace(newError("failed to handle UDP input").Base(err))
			conn.cancel()
			return
		}
		timer.Update()
		for _, b := range mb {
			callback(b)
		}
	}
}
<|endoftext|>


<commit_before>//  Copyright 2016 The goscope 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 hantek6022be

import (
	"github.com/pkg/errors"
	"github.com/zagrodzki/goscope/usb/usbif"
)

// New initializes oscilloscope through the passed USB device.
func New(d usbif.Device) (*Scope, error) {
	o := &Scope{dev: d}
	o.ch = [2]*ch{
		{id: "CH1", osc: o},
		{id: "CH2", osc: o},
	}
	for _, ch := range o.ch {
		ch.SetVoltRange(5)
	}
	o.SetSampleRate(48e6)
	if err := o.readCalibrationDataFromDevice(); err != nil {
		return nil, errors.Wrap(err, "readCalibration")
	}
	return o, nil
}

// Close releases the USB device.
func (h *Scope) Close() {
	h.dev.Close()
}

// SupportsUSB will return true if the USB descriptor passed as the argument corresponds to a Hantek 6022BE oscilloscope.
// Used for device autodetection.
func SupportsUSB(d *usbif.Desc) bool {
	return d.Vendor == hantekVendor && d.Product == hantekProduct
}
<commit_msg>default to 1Msps rate on initialization<commit_after>//  Copyright 2016 The goscope 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 hantek6022be

import (
	"github.com/pkg/errors"
	"github.com/zagrodzki/goscope/usb/usbif"
)

// New initializes oscilloscope through the passed USB device.
func New(d usbif.Device) (*Scope, error) {
	o := &Scope{dev: d}
	o.ch = [2]*ch{
		{id: "CH1", osc: o},
		{id: "CH2", osc: o},
	}
	for _, ch := range o.ch {
		ch.SetVoltRange(5)
	}
	o.SetSampleRate(1e6)
	if err := o.readCalibrationDataFromDevice(); err != nil {
		return nil, errors.Wrap(err, "readCalibration")
	}
	return o, nil
}

// Close releases the USB device.
func (h *Scope) Close() {
	h.dev.Close()
}

// SupportsUSB will return true if the USB descriptor passed as the argument corresponds to a Hantek 6022BE oscilloscope.
// Used for device autodetection.
func SupportsUSB(d *usbif.Desc) bool {
	return d.Vendor == hantekVendor && d.Product == hantekProduct
}
<|endoftext|>


<commit_before>package swarm

import (
	"sync"
	"xd/lib/bittorrent"
	"xd/lib/common"
	"xd/lib/log"
	"xd/lib/storage"
)

// how big should we download pieces at a time (bytes)?
const BlockSize = 1024 * 16

const Missing = 0
const Pending = 1
const Obtained = 2

// cached downloading piece
type cachedPiece struct {
	piece    common.PieceData
	progress []byte
	mtx      sync.RWMutex
}

// is this piece done downloading ?
func (p *cachedPiece) done() bool {
	for _, b := range p.progress {
		if b != Obtained {
			return false
		}
	}
	return true
}

// put a slice of data at offset
func (p *cachedPiece) put(offset uint32, data []byte) {
	p.mtx.Lock()
	if offset+uint32(len(data)) <= uint32(len(p.progress)) {
		// put data
		copy(p.piece.Data[offset:], data)
		// put progress
		for idx := range data {
			p.progress[uint32(idx)+offset] = Obtained
		}
	} else {
		log.Warnf("block out of range %d, %d %d", offset, len(data), len(p.progress))
	}
	p.mtx.Unlock()
}

// cancel a slice
func (p *cachedPiece) cancel(offset, length uint32) {
	p.mtx.Lock()
	p.set(offset, length, Missing)
	p.mtx.Unlock()
}

func (p *cachedPiece) set(offset, length uint32, b byte) {
	l := uint32(len(p.progress))
	if offset+length <= l {
		for length > 0 {
			length--
			p.progress[offset+length] = b
		}
	} else {
		log.Warnf("invalid cached piece range: %d %d %d", offset, length, l)
	}
}

func (p *cachedPiece) nextRequest() (r *common.PieceRequest) {
	p.mtx.Lock()
	defer p.mtx.Unlock()
	l := uint32(len(p.progress))
	r = &common.PieceRequest{
		Index:  p.piece.Index,
		Length: BlockSize,
	}
	for r.Begin+r.Length < l {

		if p.progress[r.Begin] == Missing {

			break
		}

		r.Begin += BlockSize
	}
	if r.Begin+r.Length >= l {
		r.Length = l - r.Begin
	}
	p.set(r.Begin, r.Length, Pending)
	return
}

type pieceTracker struct {
	mtx      sync.RWMutex
	requests map[uint32]*cachedPiece
	st       storage.Torrent
}

func createPieceTracker(st storage.Torrent) (pt *pieceTracker) {
	pt = &pieceTracker{
		requests: make(map[uint32]*cachedPiece),
		st:       st,
	}
	return
}

func (pt *pieceTracker) getPiece(piece uint32) (cp *cachedPiece) {
	pt.mtx.Lock()
	defer pt.mtx.Unlock()
	cp, _ = pt.requests[piece]
	return
}

func (pt *pieceTracker) newPiece(piece uint32) (cp *cachedPiece) {
	info := pt.st.MetaInfo()
	np := info.Info.NumPieces()
	sz := uint64(info.Info.PieceLength)
	if piece+1 == np {
		sz = uint64(np)*sz - info.TotalSize()
	}
	log.Debugf("new cached piece of size %d", sz)
	cp = &cachedPiece{
		progress: make([]byte, sz),
	}
	cp.piece.Data = make([]byte, sz)
	cp.piece.Index = piece
	return
}

func (pt *pieceTracker) removePiece(piece uint32) {
	pt.mtx.Lock()
	defer pt.mtx.Unlock()
	delete(pt.requests, piece)
}

func (pt *pieceTracker) nextRequestForDownload(remote *bittorrent.Bitfield) (r *common.PieceRequest) {
	bf := pt.st.Bitfield()
	i := pt.st.MetaInfo()
	np := i.Info.NumPieces()
	var idx uint32
	for idx < np {
		if remote.Has(idx) && !bf.Has(idx) {
			pt.mtx.Lock()
			cp, has := pt.requests[idx]
			if !has {
				cp = pt.newPiece(idx)
				pt.requests[idx] = cp
			}
			pt.mtx.Unlock()
			r = cp.nextRequest()
			return
		}
		idx++
	}
	return
}

func (pt *pieceTracker) handlePieceData(d *common.PieceData) {
	pc := pt.getPiece(d.Index)
	if pc != nil {
		pc.put(d.Begin, d.Data)
		if pc.done() {
			err := pt.st.PutPiece(&pc.piece)
			if err == nil {
				pt.removePiece(d.Index)
				pt.st.Flush()
			} else {
				log.Warnf("put piece %d failed: %s", pc.piece.Index, err)
				// try again
				pc.cancel(0, uint32(len(pc.progress)))
			}
		}
	}
}
<commit_msg>don't request slices more than once<commit_after>package swarm

import (
	"sync"
	"xd/lib/bittorrent"
	"xd/lib/common"
	"xd/lib/log"
	"xd/lib/storage"
)

// how big should we download pieces at a time (bytes)?
const BlockSize = 1024 * 16

const Missing = 0
const Pending = 1
const Obtained = 2

// cached downloading piece
type cachedPiece struct {
	piece    common.PieceData
	progress []byte
	mtx      sync.RWMutex
}

// is this piece done downloading ?
func (p *cachedPiece) done() bool {
	for _, b := range p.progress {
		if b != Obtained {
			return false
		}
	}
	return true
}

// put a slice of data at offset
func (p *cachedPiece) put(offset uint32, data []byte) {
	p.mtx.Lock()
	if offset+uint32(len(data)) <= uint32(len(p.progress)) {
		// put data
		copy(p.piece.Data[offset:], data)
		// put progress
		for idx := range data {
			p.progress[uint32(idx)+offset] = Obtained
		}
	} else {
		log.Warnf("block out of range %d, %d %d", offset, len(data), len(p.progress))
	}
	p.mtx.Unlock()
}

// cancel a slice
func (p *cachedPiece) cancel(offset, length uint32) {
	p.mtx.Lock()
	p.set(offset, length, Missing)
	p.mtx.Unlock()
}

func (p *cachedPiece) set(offset, length uint32, b byte) {
	l := uint32(len(p.progress))
	if offset+length <= l {
		for length > 0 {
			length--
			p.progress[offset+length] = b
		}
	} else {
		log.Warnf("invalid cached piece range: %d %d %d", offset, length, l)
	}
}

func (p *cachedPiece) nextRequest() (r *common.PieceRequest) {
	p.mtx.Lock()
	defer p.mtx.Unlock()
	l := uint32(len(p.progress))
	r = &common.PieceRequest{
		Index:  p.piece.Index,
		Length: BlockSize,
	}
	for r.Begin+r.Length < l {

		if p.progress[r.Begin] == Missing {

			break
		}

		r.Begin += BlockSize
	}
	if r.Begin+r.Length >= l {
		r.Length = l - r.Begin
	}
	if p.progress[r.Begin] == Pending {
		return nil
	}
	p.set(r.Begin, r.Length, Pending)
	return
}

type pieceTracker struct {
	mtx      sync.RWMutex
	requests map[uint32]*cachedPiece
	st       storage.Torrent
}

func createPieceTracker(st storage.Torrent) (pt *pieceTracker) {
	pt = &pieceTracker{
		requests: make(map[uint32]*cachedPiece),
		st:       st,
	}
	return
}

func (pt *pieceTracker) getPiece(piece uint32) (cp *cachedPiece) {
	pt.mtx.Lock()
	defer pt.mtx.Unlock()
	cp, _ = pt.requests[piece]
	return
}

func (pt *pieceTracker) newPiece(piece uint32) (cp *cachedPiece) {
	info := pt.st.MetaInfo()
	np := info.Info.NumPieces()
	sz := uint64(info.Info.PieceLength)
	if piece+1 == np {
		sz = uint64(np)*sz - info.TotalSize()
	}
	log.Debugf("new cached piece of size %d", sz)
	cp = &cachedPiece{
		progress: make([]byte, sz),
	}
	cp.piece.Data = make([]byte, sz)
	cp.piece.Index = piece
	return
}

func (pt *pieceTracker) removePiece(piece uint32) {
	pt.mtx.Lock()
	defer pt.mtx.Unlock()
	delete(pt.requests, piece)
}

func (pt *pieceTracker) nextRequestForDownload(remote *bittorrent.Bitfield) (r *common.PieceRequest) {
	bf := pt.st.Bitfield()
	i := pt.st.MetaInfo()
	np := i.Info.NumPieces()
	var idx uint32
	for idx < np {
		if remote.Has(idx) && !bf.Has(idx) {
			pt.mtx.Lock()
			cp, has := pt.requests[idx]
			if !has {
				cp = pt.newPiece(idx)
				pt.requests[idx] = cp
			}
			pt.mtx.Unlock()
			r = cp.nextRequest()
			if r != nil && r.Length > 0 {
				return
			}
		}
		idx++
	}
	return
}

func (pt *pieceTracker) handlePieceData(d *common.PieceData) {
	pc := pt.getPiece(d.Index)
	if pc != nil {
		pc.put(d.Begin, d.Data)
		if pc.done() {
			err := pt.st.PutPiece(&pc.piece)
			if err == nil {
				pt.removePiece(d.Index)
				pt.st.Flush()
			} else {
				log.Warnf("put piece %d failed: %s", pc.piece.Index, err)
				// try again
				pc.cancel(0, uint32(len(pc.progress)))
			}
		}
	}
}
<|endoftext|>