<commit_before>package sftp

import (
	"fmt"
	"io"
	"path/filepath"
	"reflect"
	"sync"
	"syscall"
)

var maxTxPacket uint32 = 1 << 15

type handleHandler func(string) string

type Handlers struct {
	FileGet  FileReader
	FilePut  FileWriter
	FileCmd  FileCmder
	FileInfo FileInfoer
}

// Server that abstracts the sftp protocol for a http request-like protocol
type RequestServer struct {
	serverConn
	Handlers        Handlers
	pktChan         chan packet
	openRequests    map[string]*Request
	openRequestLock sync.RWMutex
}

// simple factory function
// one server per user-session
func NewRequestServer(rwc io.ReadWriteCloser) (*RequestServer, error) {
	s := &RequestServer{
		serverConn: serverConn{
			conn: conn{
				Reader:      rwc,
				WriteCloser: rwc,
			},
		},
		pktChan:      make(chan packet, sftpServerWorkerCount),
		openRequests: make(map[string]*Request),
	}

	return s, nil
}

func (rs *RequestServer) nextRequest(r *Request) string {
	rs.openRequestLock.Lock()
	defer rs.openRequestLock.Unlock()
	rs.openRequests[r.Filepath] = r
	return r.Filepath
}

func (rs *RequestServer) getRequest(handle string) (*Request, bool) {
	rs.openRequestLock.Lock()
	defer rs.openRequestLock.Unlock()
	r, ok := rs.openRequests[handle]
	return r, ok
}

func (rs *RequestServer) closeRequest(handle string) {
	rs.openRequestLock.Lock()
	defer rs.openRequestLock.Unlock()
	if _, ok := rs.openRequests[handle]; ok {
		delete(rs.openRequests, handle)
	}
}

// start serving requests from user session
func (rs *RequestServer) Serve() error {
	var wg sync.WaitGroup
	wg.Add(sftpServerWorkerCount)
	for i := 0; i < sftpServerWorkerCount; i++ {
		go func() {
			defer wg.Done()
			if err := rs.packetWorker(); err != nil {
				rs.conn.Close() // shuts down recvPacket
			}
		}()
	}

	var err error
	var pktType uint8
	var pktBytes []byte
	for {
		pktType, pktBytes, err = rs.recvPacket()
		if err != nil {
			break
		}
		pkt, err := makePacket(rxPacket{fxp(pktType), pktBytes})
		if err != nil {
			break
		}
		rs.pktChan <- pkt
	}

	close(rs.pktChan) // shuts down sftpServerWorkers
	wg.Wait()         // wait for all workers to exit
	return err
}

func (rs *RequestServer) packetWorker() error {
	for pkt := range rs.pktChan {
		fmt.Println("Incoming Packet: ", pkt, reflect.TypeOf(pkt))
		var handle string
		var rpkt resp_packet
		var err error
		switch pkt := pkt.(type) {
		case *sshFxInitPacket:
			rpkt = sshFxVersionPacket{sftpProtocolVersion, nil}
		case *sshFxpClosePacket:
			handle = pkt.getHandle()
			rs.closeRequest(handle)
			rpkt = statusFromError(pkt, nil)
		case *sshFxpRealpathPacket:
			rpkt = cleanPath(pkt)
		case isOpener:
			handle = rs.nextRequest(newRequest(pkt.getPath()))
			rpkt = sshFxpHandlePacket{pkt.id(), handle}
		case hasPath:
			handle = rs.nextRequest(newRequest(pkt.getPath()))
			rpkt = rs.request(handle, pkt)
		case hasHandle:
			handle = pkt.getHandle()
			rpkt = rs.request(handle, pkt)
		}

		fmt.Println("Reply Packet: ", rpkt, reflect.TypeOf(rpkt))
		err = rs.sendPacket(rpkt)
		if err != nil {
			return err
		}
	}
	return nil
}

func cleanPath(pkt *sshFxpRealpathPacket) resp_packet {
	cleaned_path := filepath.Clean(pkt.getPath())
	return &sshFxpNamePacket{
		ID: pkt.id(),
		NameAttrs: []sshFxpNameAttr{{
			Name:     cleaned_path,
			LongName: cleaned_path,
			Attrs:    emptyFileStat,
		}},
	}
}

func (rs *RequestServer) request(handle string, pkt packet) resp_packet {
	var rpkt resp_packet
	var err error
	if request, ok := rs.getRequest(handle); ok {
		// called here to keep packet handling out of request for testing
		request.populate(pkt)
		fmt.Println("Request Method: ", request.Method)
		rpkt, err = request.handle(rs.Handlers)
		if err != nil {
			rpkt = statusFromError(pkt, err)
		}
	} else {
		rpkt = statusFromError(pkt, syscall.EBADF)
	}
	return rpkt
}
<commit_msg>all paths should be absolute<commit_after>package sftp

import (
	"fmt"
	"io"
	"path/filepath"
	"reflect"
	"sync"
	"syscall"
)

var maxTxPacket uint32 = 1 << 15

type handleHandler func(string) string

type Handlers struct {
	FileGet  FileReader
	FilePut  FileWriter
	FileCmd  FileCmder
	FileInfo FileInfoer
}

// Server that abstracts the sftp protocol for a http request-like protocol
type RequestServer struct {
	serverConn
	Handlers        Handlers
	pktChan         chan packet
	openRequests    map[string]*Request
	openRequestLock sync.RWMutex
}

// simple factory function
// one server per user-session
func NewRequestServer(rwc io.ReadWriteCloser) (*RequestServer, error) {
	s := &RequestServer{
		serverConn: serverConn{
			conn: conn{
				Reader:      rwc,
				WriteCloser: rwc,
			},
		},
		pktChan:      make(chan packet, sftpServerWorkerCount),
		openRequests: make(map[string]*Request),
	}

	return s, nil
}

func (rs *RequestServer) nextRequest(r *Request) string {
	rs.openRequestLock.Lock()
	defer rs.openRequestLock.Unlock()
	rs.openRequests[r.Filepath] = r
	return r.Filepath
}

func (rs *RequestServer) getRequest(handle string) (*Request, bool) {
	rs.openRequestLock.Lock()
	defer rs.openRequestLock.Unlock()
	r, ok := rs.openRequests[handle]
	return r, ok
}

func (rs *RequestServer) closeRequest(handle string) {
	rs.openRequestLock.Lock()
	defer rs.openRequestLock.Unlock()
	if _, ok := rs.openRequests[handle]; ok {
		delete(rs.openRequests, handle)
	}
}

// start serving requests from user session
func (rs *RequestServer) Serve() error {
	var wg sync.WaitGroup
	wg.Add(sftpServerWorkerCount)
	for i := 0; i < sftpServerWorkerCount; i++ {
		go func() {
			defer wg.Done()
			if err := rs.packetWorker(); err != nil {
				rs.conn.Close() // shuts down recvPacket
			}
		}()
	}

	var err error
	var pktType uint8
	var pktBytes []byte
	for {
		pktType, pktBytes, err = rs.recvPacket()
		if err != nil {
			break
		}
		pkt, err := makePacket(rxPacket{fxp(pktType), pktBytes})
		if err != nil {
			break
		}
		rs.pktChan <- pkt
	}

	close(rs.pktChan) // shuts down sftpServerWorkers
	wg.Wait()         // wait for all workers to exit
	return err
}

func (rs *RequestServer) packetWorker() error {
	for pkt := range rs.pktChan {
		fmt.Println("Incoming Packet: ", pkt, reflect.TypeOf(pkt))
		var handle string
		var rpkt resp_packet
		var err error
		switch pkt := pkt.(type) {
		case *sshFxInitPacket:
			rpkt = sshFxVersionPacket{sftpProtocolVersion, nil}
		case *sshFxpClosePacket:
			handle = pkt.getHandle()
			rs.closeRequest(handle)
			rpkt = statusFromError(pkt, nil)
		case *sshFxpRealpathPacket:
			rpkt = cleanPath(pkt)
		case isOpener:
			handle = rs.nextRequest(newRequest(pkt.getPath()))
			rpkt = sshFxpHandlePacket{pkt.id(), handle}
		case hasPath:
			handle = rs.nextRequest(newRequest(pkt.getPath()))
			rpkt = rs.request(handle, pkt)
		case hasHandle:
			handle = pkt.getHandle()
			rpkt = rs.request(handle, pkt)
		}

		fmt.Println("Reply Packet: ", rpkt, reflect.TypeOf(rpkt))
		err = rs.sendPacket(rpkt)
		if err != nil {
			return err
		}
	}
	return nil
}

func cleanPath(pkt *sshFxpRealpathPacket) resp_packet {
	path := pkt.getPath()
	if !filepath.IsAbs(path) {
		path = "/" + path // all paths are absolute
	}
	cleaned_path := filepath.Clean(path)
	return &sshFxpNamePacket{
		ID: pkt.id(),
		NameAttrs: []sshFxpNameAttr{{
			Name:     cleaned_path,
			LongName: cleaned_path,
			Attrs:    emptyFileStat,
		}},
	}
}

func (rs *RequestServer) request(handle string, pkt packet) resp_packet {
	var rpkt resp_packet
	var err error
	if request, ok := rs.getRequest(handle); ok {
		// called here to keep packet handling out of request for testing
		request.populate(pkt)
		fmt.Println("Request Method: ", request.Method)
		rpkt, err = request.handle(rs.Handlers)
		if err != nil {
			rpkt = statusFromError(pkt, err)
		}
	} else {
		rpkt = statusFromError(pkt, syscall.EBADF)
	}
	return rpkt
}
<|endoftext|>


<commit_before>package requestcontext

import (
	"io"
	"net/http"

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

type decoratedReadCloser struct {
	body    io.ReadCloser
	context context.Context
}

func (w *decoratedReadCloser) Read(p []byte) (n int, err error) {
	return w.body.Read(p)
}

func (w *decoratedReadCloser) Close() error {
	return w.body.Close()
}

// Set add this context around the current context.
func Set(r *http.Request, context context.Context) {
	drc := decoratedReadCloser{body: r.Body, context: context}
	r.Body = &drc
}

// Get returns a parent context, will make one if needed.
// always call this before Set() to preserve context chain
func Get(r *http.Request) context.Context {
	if c, ok := r.Body.(*decoratedReadCloser); ok {
		return c.context
	}
	return context.Background()
}
<commit_msg>exposition - briefly<commit_after>// Package requestcontext is a contention free http.Request adapter for golang.org/x/net/context
package requestcontext

import (
	"io"
	"net/http"

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

type decoratedReadCloser struct {
	body    io.ReadCloser
	context context.Context
}

func (w *decoratedReadCloser) Read(p []byte) (n int, err error) {
	return w.body.Read(p)
}

func (w *decoratedReadCloser) Close() error {
	return w.body.Close()
}

// Set this context
// call Get() prior and wrap the currency context with
// one of the golang.org/x/net/context methods - see example
func Set(r *http.Request, context context.Context) {
	drc := decoratedReadCloser{body: r.Body, context: context}
	r.Body = &drc
}

// Get returns a parent context, will make one if needed.
// always call this before Set() to preserve context chain
func Get(r *http.Request) context.Context {
	if c, ok := r.Body.(*decoratedReadCloser); ok {
		return c.context
	}
	return context.Background()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math/rand"
	"strings"
	"time"
)

// TimescaleDevops produces Timescale-specific queries for all the devops query types.
type TimescaleDevops struct {
	DatabaseName string
	AllInterval  TimeInterval
}

// newTimescaleDevopsCommon makes an TimescaleDevops object ready to generate Queries.
func newTimescaleDevopsCommon(dbConfig DatabaseConfig, start, end time.Time) QueryGenerator {
	if !start.Before(end) {
		panic("bad time order")
	}
	if _, ok := dbConfig["database-name"]; !ok {
		panic("need timescale database name")
	}

	return &TimescaleDevops{
		DatabaseName: dbConfig["database-name"],
		AllInterval:  NewTimeInterval(start, end),
	}
}

// Dispatch fulfills the QueryGenerator interface.
func (d *TimescaleDevops) Dispatch(i, scaleVar int) Query {
	q := NewHTTPQuery() // from pool
	devopsDispatchAll(d, i, q, scaleVar)
	return q
}

func (d *TimescaleDevops) MaxCPUUsageHourByMinuteOneHost(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 1, time.Hour)
}

func (d *TimescaleDevops) MaxCPUUsageHourByMinuteTwoHosts(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 2, time.Hour)
}

func (d *TimescaleDevops) MaxCPUUsageHourByMinuteFourHosts(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 4, time.Hour)
}

func (d *TimescaleDevops) MaxCPUUsageHourByMinuteEightHosts(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 8, time.Hour)
}

func (d *TimescaleDevops) MaxCPUUsageHourByMinuteSixteenHosts(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 16, time.Hour)
}

func (d *TimescaleDevops) MaxCPUUsageHourByMinuteThirtyTwoHosts(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 32, time.Hour)
}

func (d *TimescaleDevops) MaxCPUUsage12HoursByMinuteOneHost(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 1, 12*time.Hour)
}

// MaxCPUUsageHourByMinuteThirtyTwoHosts populates a Query with a query that looks like:
// select time_bucket(60000000000,time) as time1min,max(usage_user) from cpu where (hostname = '$HOSTNAME_1' or ... or hostname = '$HOSTNAME_N') and time >=$HOUR_START and time < $HOUR_END group by time1min order by time1min;
func (d *TimescaleDevops) maxCPUUsageHourByMinuteNHosts(qi Query, scaleVar, nhosts int, timeRange time.Duration) {
	interval := d.AllInterval.RandWindow(timeRange)
	nn := rand.Perm(scaleVar)[:nhosts]

	hostnames := []string{}
	for _, n := range nn {
		hostnames = append(hostnames, fmt.Sprintf("host_%d", n))
	}

	hostnameClauses := []string{}
	for _, s := range hostnames {
		hostnameClauses = append(hostnameClauses, fmt.Sprintf("hostname = '%s'", s))
	}

	combinedHostnameClause := strings.Join(hostnameClauses, " or ")

	humanLabel := fmt.Sprintf("Timescale max cpu, rand %4d hosts, rand %s by 1m", nhosts, timeRange)

	q := qi.(*SQLQuery)
	q.HumanLabel = []byte(humanLabel)
	q.HumanDescription = []byte(fmt.Sprintf("%s: %s", humanLabel, interval.StartString()))

	q.QuerySQL = []byte(fmt.Sprintf("select time_bucket(60000000000,time) as \"time1min\",max(usage_user) from cpu where (%s) and time >=%d and time < %d group by \"time1min\" order by \"time1min\"", combinedHostnameClause, interval.StartUnixNano(), interval.EndUnixNano()))
}

// MeanCPUUsageDayByHourAllHosts populates a Query with a query that looks like:
// SELECT mean(usage_user) from cpu where time >= '$DAY_START' and time < '$DAY_END' group by time(1h),hostname
func (d *TimescaleDevops) MeanCPUUsageDayByHourAllHostsGroupbyHost(qi Query, _ int) {
	interval := d.AllInterval.RandWindow(24 * time.Hour)

	humanLabel := "Timescale mean cpu, all hosts, rand 1day by 1hour"
	q := qi.(*SQLQuery)
	q.HumanLabel = []byte(humanLabel)
	q.HumanDescription = []byte(fmt.Sprintf("%s: %s", humanLabel, interval.StartString()))

	q.QuerySQL = []byte(fmt.Sprintf("select time_bucket(3600000000000,time) as time1hour,avg(usage_user) from cpu where time >=%d and time < %d group by time1hour,hostname order by time1hour", interval.StartUnixNano(), interval.EndUnixNano()))
}
<commit_msg>Work-arounding query problem<commit_after>package main

import (
	"fmt"
	"math/rand"
	"strings"
	"time"
)

// TimescaleDevops produces Timescale-specific queries for all the devops query types.
type TimescaleDevops struct {
	DatabaseName string
	AllInterval  TimeInterval
}

// newTimescaleDevopsCommon makes an TimescaleDevops object ready to generate Queries.
func newTimescaleDevopsCommon(dbConfig DatabaseConfig, start, end time.Time) QueryGenerator {
	if !start.Before(end) {
		panic("bad time order")
	}
	if _, ok := dbConfig["database-name"]; !ok {
		panic("need timescale database name")
	}

	return &TimescaleDevops{
		DatabaseName: dbConfig["database-name"],
		AllInterval:  NewTimeInterval(start, end),
	}
}

// Dispatch fulfills the QueryGenerator interface.
func (d *TimescaleDevops) Dispatch(i, scaleVar int) Query {
	q := NewHTTPQuery() // from pool
	devopsDispatchAll(d, i, q, scaleVar)
	return q
}

func (d *TimescaleDevops) MaxCPUUsageHourByMinuteOneHost(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 1, time.Hour)
}

func (d *TimescaleDevops) MaxCPUUsageHourByMinuteTwoHosts(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 2, time.Hour)
}

func (d *TimescaleDevops) MaxCPUUsageHourByMinuteFourHosts(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 4, time.Hour)
}

func (d *TimescaleDevops) MaxCPUUsageHourByMinuteEightHosts(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 8, time.Hour)
}

func (d *TimescaleDevops) MaxCPUUsageHourByMinuteSixteenHosts(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 16, time.Hour)
}

func (d *TimescaleDevops) MaxCPUUsageHourByMinuteThirtyTwoHosts(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 32, time.Hour)
}

func (d *TimescaleDevops) MaxCPUUsage12HoursByMinuteOneHost(q Query, scaleVar int) {
	d.maxCPUUsageHourByMinuteNHosts(q, scaleVar, 1, 12*time.Hour)
}

// MaxCPUUsageHourByMinuteThirtyTwoHosts populates a Query with a query that looks like:
// select time_bucket(60000000000,time) as time1min,max(usage_user) from cpu where (hostname = '$HOSTNAME_1' or ... or hostname = '$HOSTNAME_N') and time >=$HOUR_START and time < $HOUR_END group by time1min order by time1min;
func (d *TimescaleDevops) maxCPUUsageHourByMinuteNHosts(qi Query, scaleVar, nhosts int, timeRange time.Duration) {
	interval := d.AllInterval.RandWindow(timeRange)
	nn := rand.Perm(scaleVar)[:nhosts]

	hostnames := []string{}
	for _, n := range nn {
		hostnames = append(hostnames, fmt.Sprintf("host_%d", n))
	}

	hostnameClauses := []string{}
	for _, s := range hostnames {
		hostnameClauses = append(hostnameClauses, fmt.Sprintf("hostname = '%s'", s))
	}

	combinedHostnameClause := strings.Join(hostnameClauses, " or ")

	humanLabel := fmt.Sprintf("Timescale max cpu, rand %4d hosts, rand %s by 1m", nhosts, timeRange)

	q := qi.(*SQLQuery)
	q.HumanLabel = []byte(humanLabel)
	q.HumanDescription = []byte(fmt.Sprintf("%s: %s", humanLabel, interval.StartString()))

	q.QuerySQL = []byte(fmt.Sprintf("select time_bucket(60000000000,time) as time1min,max(usage_user) from cpu where (%s) and time >=%d and time < %d group by time1min order by time1min ", combinedHostnameClause, interval.StartUnixNano(), interval.EndUnixNano()))
}

// MeanCPUUsageDayByHourAllHosts populates a Query with a query that looks like:
// SELECT mean(usage_user) from cpu where time >= '$DAY_START' and time < '$DAY_END' group by time(1h),hostname
func (d *TimescaleDevops) MeanCPUUsageDayByHourAllHostsGroupbyHost(qi Query, _ int) {
	interval := d.AllInterval.RandWindow(24 * time.Hour)

	humanLabel := "Timescale mean cpu, all hosts, rand 1day by 1hour"
	q := qi.(*SQLQuery)
	q.HumanLabel = []byte(humanLabel)
	q.HumanDescription = []byte(fmt.Sprintf("%s: %s", humanLabel, interval.StartString()))

	q.QuerySQL = []byte(fmt.Sprintf("select time_bucket(3600000000000,time) as time1hour,avg(usage_user) from cpu where time >=%d and time < %d group by time1hour,hostname order by time1hour", interval.StartUnixNano(), interval.EndUnixNano()))
}
<|endoftext|>


<commit_before>// Copyright 2017 Jesse Allen. All rights reserved
// Released under the MIT license found in the LICENSE file.

package rest_test

import (
	"bytes"
	"math/rand"
	"mime"
	"net/http"
	"net/http/httptest"
	"reflect"
	"strings"
	"sync"
	"testing"
	"testing/quick"

	"github.com/lazyengineering/faststatus"
	"github.com/lazyengineering/faststatus/rest"
)

func TestHandlerOnlyValidPaths(t *testing.T) {
	var s, _ = rest.New()

	isNotFound := func(method, path string) bool {
		w := httptest.NewRecorder()
		r := httptest.NewRequest(method, path, nil)
		s.ServeHTTP(w, r)

		_, isValidPath := validMethodsByPath[path]
		if !isValidPath && w.Code != http.StatusNotFound {
			t.Logf("%s %q got %03d, expected 404", method, path, w.Code)
			return false
		}
		if isValidPath && w.Code == http.StatusNotFound {
			t.Logf("%s %q got %03d, expected not 404", method, path, w.Code)
			return false
		}
		return true
	}
	if err := quick.Check(isNotFound, &quick.Config{
		Values: func(args []reflect.Value, gen *rand.Rand) {
			args[0] = reflect.ValueOf(possibleMethods[gen.Intn(len(possibleMethods))])
			scratch := make([]byte, 1000)
			gen.Read(scratch)
			var path = "/" + strings.Map(func(char rune) rune {
				switch {
				case char >= '0' && char <= '9',
					char >= 'A' && char <= 'Z',
					char >= 'a' && char <= 'z',
					char == '!',
					char == '$',
					char >= 0x27 && char <= '/':
					return char
				default:
					return -1
				}
			}, string(scratch))
			path = path[:25%len(path)]
			args[1] = reflect.ValueOf(path)
		},
	}); err != nil {
		t.Fatalf("unexpected response to bad client request: %+v", err)
	}
}

func TestHandlerOnlyValidPathsAndMethods(t *testing.T) {
	var s, _ = rest.New()
	isNotAllowed := func(path string) func(string) bool {
		return func(method string) bool {
			w := httptest.NewRecorder()
			r := httptest.NewRequest(method, path, nil)
			s.ServeHTTP(w, r)

			validMethods := validMethodsByPath[path]
			var isValid bool
			for _, validMethod := range validMethods {
				if validMethod == method {
					isValid = true
					break
				}
			}
			if !isValid && w.Code != http.StatusMethodNotAllowed {
				t.Logf("%s %q got %03d, expected 405", method, path, w.Code)
				return false
			}
			if isValid && w.Code == http.StatusMethodNotAllowed {
				t.Logf("%s %q got %03d, expected not 405", method, path, w.Code)
				return false
			}
			return true
		}
	}
	genMethod := func(args []reflect.Value, gen *rand.Rand) {
		args[0] = reflect.ValueOf(possibleMethods[gen.Intn(len(possibleMethods))])
	}
	for path := range validMethodsByPath {
		if err := quick.Check(isNotAllowed(path), &quick.Config{Values: genMethod}); err != nil {
			t.Fatalf("unexpected response to bad client request: %+v", err)
		}
	}
}

var possibleMethods = []string{
	http.MethodGet,
	http.MethodHead,
	http.MethodPut,
	http.MethodPost,
	http.MethodPatch,
	http.MethodDelete,
}
var validMethodsByPath = map[string][]string{
	"/new": []string{http.MethodGet, http.MethodHead},
}

func TestHandlerGetNew(t *testing.T) {
	const defaultContentType = "text/plain"
	var s, _ = rest.New()
	var mu sync.Mutex
	var usedIDs = make(map[faststatus.ID]struct{})
	getsNewResource := func() bool {
		w := httptest.NewRecorder()
		r := httptest.NewRequest(http.MethodGet, "/new", nil)
		s.ServeHTTP(w, r)

		if w.Code != http.StatusOK {
			t.Logf(
				"returned Status Code %03d, expected %03d",
				w.Code,
				http.StatusOK,
			)
			return false
		}

		gotType, _, err := mime.ParseMediaType(w.HeaderMap.Get("Content-Type"))
		if err != nil {
			t.Logf("error parsing content type: %+v", err)
			return false
		}
		if gotType != defaultContentType {
			t.Logf("Content-Type %q, expected %q", gotType, defaultContentType)
			return false
		}

		var (
			gotResource  faststatus.Resource
			zeroResource faststatus.Resource
		)
		err = (&gotResource).UnmarshalText(w.Body.Bytes())
		if err != nil {
			t.Logf("error unmarshaling Resource from text body: %+v", err)
			return false
		}

		if bytes.Equal(gotResource.ID[:], zeroResource.ID[:]) {
			t.Logf("ID should be non-zero")
			return false
		}

		if _, exists := usedIDs[gotResource.ID]; exists {
			t.Logf("ID should be unique")
			return false
		}
		mu.Lock()
		usedIDs[gotResource.ID] = struct{}{}
		mu.Unlock()

		gotResource.ID = zeroResource.ID
		if !gotResource.Equal(zeroResource) {
			t.Logf("with ID set to zero-value got %+v, expected %+v", gotResource, zeroResource)
			return false
		}

		return true
	}
	if err := quick.Check(getsNewResource, nil); err != nil {
		t.Fatalf("GET /new does not get a new Resource: %+v", err)
	}
}
<commit_msg>Prepares not found test for extension<commit_after>// Copyright 2017 Jesse Allen. All rights reserved
// Released under the MIT license found in the LICENSE file.

package rest_test

import (
	"bytes"
	"math/rand"
	"mime"
	"net/http"
	"net/http/httptest"
	"reflect"
	"strings"
	"sync"
	"testing"
	"testing/quick"

	"github.com/lazyengineering/faststatus"
	"github.com/lazyengineering/faststatus/rest"
)

func TestHandlerOnlyValidPaths(t *testing.T) {
	var s, _ = rest.New()
	invalidPathIsNotFound := func(method, path string) bool {
		w := httptest.NewRecorder()
		r := httptest.NewRequest(method, path, nil)
		s.ServeHTTP(w, r)

		if w.Code != http.StatusNotFound {
			t.Logf("%s %q got %03d, expected 404", method, path, w.Code)
			return false
		}
		return true
	}
	err := quick.Check(invalidPathIsNotFound, &quick.Config{
		Values: func(args []reflect.Value, gen *rand.Rand) {
			args[0] = reflect.ValueOf(possibleMethods[gen.Intn(len(possibleMethods))])
			args[1] = reflect.ValueOf(genInvalidPath(gen.Intn(100)+1, gen))
		},
	})
	if err != nil {
		t.Fatalf("unexpected response to bad client request: %+v", err)
	}
}

func TestHandlerOnlyValidPathsAndMethods(t *testing.T) {
	var s, _ = rest.New()
	isNotAllowed := func(path string) func(string) bool {
		return func(method string) bool {
			w := httptest.NewRecorder()
			r := httptest.NewRequest(method, path, nil)
			s.ServeHTTP(w, r)

			validMethods := validMethodsByPath[path]
			var isValid bool
			for _, validMethod := range validMethods {
				if validMethod == method {
					isValid = true
					break
				}
			}
			if !isValid && w.Code != http.StatusMethodNotAllowed {
				t.Logf("%s %q got %03d, expected 405", method, path, w.Code)
				return false
			}
			if isValid && w.Code == http.StatusMethodNotAllowed {
				t.Logf("%s %q got %03d, expected not 405", method, path, w.Code)
				return false
			}
			return true
		}
	}
	genMethod := func(args []reflect.Value, gen *rand.Rand) {
		args[0] = reflect.ValueOf(possibleMethods[gen.Intn(len(possibleMethods))])
	}
	for path := range validMethodsByPath {
		if err := quick.Check(isNotAllowed(path), &quick.Config{Values: genMethod}); err != nil {
			t.Fatalf("unexpected response to bad client request: %+v", err)
		}
	}
}

var possibleMethods = []string{
	http.MethodGet,
	http.MethodHead,
	http.MethodPut,
	http.MethodPost,
	http.MethodPatch,
	http.MethodDelete,
}
var validMethodsByPath = map[string][]string{
	"/new": []string{http.MethodGet, http.MethodHead},
}

func TestHandlerGetNew(t *testing.T) {
	const defaultContentType = "text/plain"
	var s, _ = rest.New()
	var mu sync.Mutex
	var usedIDs = make(map[faststatus.ID]struct{})
	getsNewResource := func() bool {
		w := httptest.NewRecorder()
		r := httptest.NewRequest(http.MethodGet, "/new", nil)
		s.ServeHTTP(w, r)

		if w.Code != http.StatusOK {
			t.Logf(
				"returned Status Code %03d, expected %03d",
				w.Code,
				http.StatusOK,
			)
			return false
		}

		gotType, _, err := mime.ParseMediaType(w.HeaderMap.Get("Content-Type"))
		if err != nil {
			t.Logf("error parsing content type: %+v", err)
			return false
		}
		if gotType != defaultContentType {
			t.Logf("Content-Type %q, expected %q", gotType, defaultContentType)
			return false
		}

		var (
			gotResource  faststatus.Resource
			zeroResource faststatus.Resource
		)
		err = (&gotResource).UnmarshalText(w.Body.Bytes())
		if err != nil {
			t.Logf("error unmarshaling Resource from text body: %+v", err)
			return false
		}

		if bytes.Equal(gotResource.ID[:], zeroResource.ID[:]) {
			t.Logf("ID should be non-zero")
			return false
		}

		if _, exists := usedIDs[gotResource.ID]; exists {
			t.Logf("ID should be unique")
			return false
		}
		mu.Lock()
		usedIDs[gotResource.ID] = struct{}{}
		mu.Unlock()

		gotResource.ID = zeroResource.ID
		if !gotResource.Equal(zeroResource) {
			t.Logf("with ID set to zero-value got %+v, expected %+v", gotResource, zeroResource)
			return false
		}

		return true
	}
	if err := quick.Check(getsNewResource, nil); err != nil {
		t.Fatalf("GET /new does not get a new Resource: %+v", err)
	}
}

var possibleMethods = []string{
	http.MethodGet,
	http.MethodHead,
	http.MethodPut,
	http.MethodPost,
	http.MethodPatch,
	http.MethodDelete,
}

func validMethodsByPath(path string) ([]string, bool) {
	if path == "/new" {
		return []string{http.MethodGet, http.MethodHead}, true
	}
	return nil, false
}

func genInvalidPath(maxLen int, r *rand.Rand) string {
	for {
		path := genPath(maxLen, r)
		_, ok := validMethodsByPath(path)
		if !ok {
			return path
		}
	}
}

func genPath(maxLen int, r *rand.Rand) string {
	scratch := make([]byte, 1000)
	r.Read(scratch)
	var path = "/" + strings.Map(func(char rune) rune {
		switch {
		case char >= '0' && char <= '9',
			char >= 'A' && char <= 'Z',
			char >= 'a' && char <= 'z',
			char == '!',
			char == '$',
			char >= 0x27 && char <= '/':
			return char
		default:
			return -1
		}
	}, string(scratch))
	return path[:maxLen%len(path)]
}
<|endoftext|>


<commit_before>package nakadi

import (
	"context"
	"math/rand"
	"sync"
	"time"

	"github.com/pkg/errors"
)

// ProcessorOptions contains optional parameters that are used to create a Processor.
type ProcessorOptions struct {
	// The maximum number of Events in each chunk (and therefore per partition) of the stream (default: 1)
	BatchLimit uint
	// The number of parallel streams the Processor will use to consume events (default: 1)
	StreamCount uint
	// Limits the number of events that the processor will handle per minute. This value represents an
	// upper bound, if some streams are not healthy e.g. if StreamCount exceeds the number of partitions,
	// or the actual batch size is lower than BatchLimit the actual number of processed events can be
	// much lower. 0 is interpreted as no limit at all (default: no limit)
	EventsPerMinute uint
	// The initial (minimal) retry interval used for the exponential backoff. This value is applied for
	// stream initialization as well as for cursor commits.
	InitialRetryInterval time.Duration
	// MaxRetryInterval the maximum retry interval. Once the exponential backoff reaches this value
	// the retry intervals remain constant. This value is applied for stream initialization as well as
	// for cursor commits.
	MaxRetryInterval time.Duration
	// MaxElapsedTime is the maximum time spent on retries when committing a cursor. Once this value
	// was reached the exponential backoff is halted and the cursor will not be committed.
	CommitMaxElapsedTime time.Duration
	// NotifyErr is called when an error occurs that leads to a retry. This notify function can be used to
	// detect unhealthy streams. The first parameter indicates the stream No that encountered the error.
	NotifyErr func(uint, error, time.Duration)
	// NotifyOK is called whenever a successful operation was completed. This notify function can be used
	// to detect that a stream is healthy again. The first parameter indicates the stream No that just
	// regained health.
	NotifyOK func(uint)
}

func (o *ProcessorOptions) withDefaults() *ProcessorOptions {
	var copyOptions ProcessorOptions
	if o != nil {
		copyOptions = *o
	}
	if copyOptions.BatchLimit == 0 {
		copyOptions.BatchLimit = 1
	}
	if copyOptions.StreamCount == 0 {
		copyOptions.StreamCount = 1
	}
	if copyOptions.EventsPerMinute == 0 {
		copyOptions.EventsPerMinute = 0
	}
	if copyOptions.InitialRetryInterval == 0 {
		copyOptions.InitialRetryInterval = defaultInitialRetryInterval
	}
	if copyOptions.MaxRetryInterval == 0 {
		copyOptions.MaxRetryInterval = defaultMaxRetryInterval
	}
	if copyOptions.CommitMaxElapsedTime == 0 {
		copyOptions.CommitMaxElapsedTime = defaultMaxElapsedTime
	}
	if copyOptions.NotifyErr == nil {
		copyOptions.NotifyErr = func(_ uint, _ error, _ time.Duration) {}
	}
	if copyOptions.NotifyOK == nil {
		copyOptions.NotifyOK = func(_ uint) {}
	}
	return &copyOptions
}

// streamAPI is a contract that is used internally in order to be able to mock StreamAPI
type streamAPI interface {
	NextEvents() (Cursor, []byte, error)
	CommitCursor(cursor Cursor) error
	Close() error
}

// NewProcessor creates a new processor for a given subscription ID.  The constructor receives a
// configured Nakadi client as first parameter. Furthermore the a valid subscription ID must be
// provided. The last parameter is a struct containing only optional parameters. The options may be
// nil, in this case the processor falls back to the defaults defined in the ProcessorOptions.
func NewProcessor(client *Client, subscriptionID string, options *ProcessorOptions) *Processor {
	options = options.withDefaults()

	var timePerBatchPerStream int64
	if options.EventsPerMinute > 0 {
		timePerBatchPerStream = int64(time.Minute) / int64(options.EventsPerMinute) * int64(options.StreamCount) * int64(options.BatchLimit)
	}

	ctx, cancel := context.WithCancel(context.Background())
	processor := &Processor{
		client:         client,
		subscriptionID: subscriptionID,
		ctx:            ctx,
		cancel:         cancel,
		newStream: func(client *Client, id string, options *StreamOptions) streamAPI {
			return NewStream(client, id, options)
		},
		timePerBatchPerStream: time.Duration(timePerBatchPerStream)}

	for i := uint(0); i < options.StreamCount; i++ {
		streamOptions := StreamOptions{
			BatchLimit:           options.BatchLimit,
			InitialRetryInterval: options.InitialRetryInterval,
			MaxRetryInterval:     options.MaxRetryInterval,
			CommitMaxElapsedTime: options.CommitMaxElapsedTime,
			NotifyErr:            func(err error, duration time.Duration) { options.NotifyErr(i, err, duration) },
			NotifyOK:             func() { options.NotifyOK(i) }}

		processor.streamOptions = append(processor.streamOptions, streamOptions)
	}

	return processor
}

// A Processor for nakadi events. The Processor is a high level API for consuming events from
// Nakadi. It can process event batches from multiple partitions (streams) and can be configured
// with a certain event rate, that limits the number of processed events per minute. The cursors of
// event batches that were successfully processed are automatically committed.
type Processor struct {
	sync.Mutex
	client                *Client
	subscriptionID        string
	streamOptions         []StreamOptions
	newStream             func(*Client, string, *StreamOptions) streamAPI
	timePerBatchPerStream time.Duration
	streams               []streamAPI
	ctx                   context.Context
	cancel                context.CancelFunc
}

// Start begins event processing. All event batches received from the underlying streams are passed to
// the operation function. If the operation function terminates without error the respective cursor will
// be automatically committed to Nakadi. If the operations terminates with an error, the underlying stream
// will be halted and a new stream will continue to pass event batches to the operation function.
//
// Event processing will go on indefinitely unless the processor is stopped via its Stop method. Star will
// return an error if the processor is already running.
func (p *Processor) Start(operation func(int, []byte) error) error {
	p.Lock()
	defer p.Unlock()

	if len(p.streams) > 0 {
		return errors.New("processor was already started")
	}

	p.streams = make([]streamAPI, len(p.streamOptions))

	for streamNo, options := range p.streamOptions {
		go p.startSingleStream(operation, streamNo, options)
	}

	return nil
}

// startSingleStream starts a single stream with a given stream number / position. After the stream has been
// started it consumes events. In cases of errors the stream is closed and a new stream will be opened.
func (p *Processor) startSingleStream(operation func(int, []byte) error, streamNo int, options StreamOptions) {
	p.streams[streamNo] = p.newStream(p.client, p.subscriptionID, &options)

	if p.timePerBatchPerStream > 0 {
		initialWait := rand.Int63n(int64(p.timePerBatchPerStream))
		select {
		case <-p.ctx.Done():
			return
		case <-time.After(time.Duration(initialWait)):
			// nothing
		}
	}

	for {
		start := time.Now()

		select {
		case <-p.ctx.Done():
			return
		default:
			cursor, events, err := p.streams[streamNo].NextEvents()
			if err != nil {
				continue
			}

			err = operation(streamNo, events)
			if err != nil {
				p.Lock()
				p.streams[streamNo].Close()
				p.streams[streamNo] = p.newStream(p.client, p.subscriptionID, &options)
				p.Unlock()
				continue
			}

			p.streams[streamNo].CommitCursor(cursor)
		}

		elapsed := time.Since(start)
		if elapsed < p.timePerBatchPerStream {
			select {
			case <-p.ctx.Done():
				return
			case <-time.After(p.timePerBatchPerStream - elapsed):
				// nothing
			}
		}
	}
}

// Stop halts all steams and terminates event processing. Stop will return with an error if the processor
// is not running.
func (p *Processor) Stop() error {
	p.Lock()
	defer p.Unlock()

	if len(p.streams) == 0 {
		return errors.New("processor is not running")
	}

	p.cancel()

	var allErrors []error
	for _, s := range p.streams {
		err := s.Close()
		if err != nil {
			allErrors = append(allErrors, err)
		}
	}

	p.streams = nil

	if len(allErrors) > 0 {
		return errors.Errorf("%d streams had errors while closing the stream", len(allErrors))
	}

	return nil
}
<commit_msg>fix mistakes in comments<commit_after>package nakadi

import (
	"context"
	"math/rand"
	"sync"
	"time"

	"github.com/pkg/errors"
)

// ProcessorOptions contains optional parameters that are used to create a Processor.
type ProcessorOptions struct {
	// The maximum number of Events in each chunk (and therefore per partition) of the stream (default: 1)
	BatchLimit uint
	// The number of parallel streams the Processor will use to consume events (default: 1)
	StreamCount uint
	// Limits the number of events that the processor will handle per minute. This value represents an
	// upper bound, if some streams are not healthy e.g. if StreamCount exceeds the number of partitions,
	// or the actual batch size is lower than BatchLimit the actual number of processed events can be
	// much lower. 0 is interpreted as no limit at all (default: no limit)
	EventsPerMinute uint
	// The initial (minimal) retry interval used for the exponential backoff. This value is applied for
	// stream initialization as well as for cursor commits.
	InitialRetryInterval time.Duration
	// MaxRetryInterval the maximum retry interval. Once the exponential backoff reaches this value
	// the retry intervals remain constant. This value is applied for stream initialization as well as
	// for cursor commits.
	MaxRetryInterval time.Duration
	// CommitMaxElapsedTime is the maximum time spent on retries when committing a cursor. Once this value
	// was reached the exponential backoff is halted and the cursor will not be committed.
	CommitMaxElapsedTime time.Duration
	// NotifyErr is called when an error occurs that leads to a retry. This notify function can be used to
	// detect unhealthy streams. The first parameter indicates the stream No that encountered the error.
	NotifyErr func(uint, error, time.Duration)
	// NotifyOK is called whenever a successful operation was completed. This notify function can be used
	// to detect that a stream is healthy again. The first parameter indicates the stream No that just
	// regained health.
	NotifyOK func(uint)
}

func (o *ProcessorOptions) withDefaults() *ProcessorOptions {
	var copyOptions ProcessorOptions
	if o != nil {
		copyOptions = *o
	}
	if copyOptions.BatchLimit == 0 {
		copyOptions.BatchLimit = 1
	}
	if copyOptions.StreamCount == 0 {
		copyOptions.StreamCount = 1
	}
	if copyOptions.EventsPerMinute == 0 {
		copyOptions.EventsPerMinute = 0
	}
	if copyOptions.InitialRetryInterval == 0 {
		copyOptions.InitialRetryInterval = defaultInitialRetryInterval
	}
	if copyOptions.MaxRetryInterval == 0 {
		copyOptions.MaxRetryInterval = defaultMaxRetryInterval
	}
	if copyOptions.CommitMaxElapsedTime == 0 {
		copyOptions.CommitMaxElapsedTime = defaultMaxElapsedTime
	}
	if copyOptions.NotifyErr == nil {
		copyOptions.NotifyErr = func(_ uint, _ error, _ time.Duration) {}
	}
	if copyOptions.NotifyOK == nil {
		copyOptions.NotifyOK = func(_ uint) {}
	}
	return &copyOptions
}

// streamAPI is a contract that is used internally in order to be able to mock StreamAPI
type streamAPI interface {
	NextEvents() (Cursor, []byte, error)
	CommitCursor(cursor Cursor) error
	Close() error
}

// NewProcessor creates a new processor for a given subscription ID.  The constructor receives a
// configured Nakadi client as first parameter. Furthermore a valid subscription ID must be
// provided. The last parameter is a struct containing only optional parameters. The options may be
// nil, in this case the processor falls back to the defaults defined in the ProcessorOptions.
func NewProcessor(client *Client, subscriptionID string, options *ProcessorOptions) *Processor {
	options = options.withDefaults()

	var timePerBatchPerStream int64
	if options.EventsPerMinute > 0 {
		timePerBatchPerStream = int64(time.Minute) / int64(options.EventsPerMinute) * int64(options.StreamCount) * int64(options.BatchLimit)
	}

	ctx, cancel := context.WithCancel(context.Background())
	processor := &Processor{
		client:         client,
		subscriptionID: subscriptionID,
		ctx:            ctx,
		cancel:         cancel,
		newStream: func(client *Client, id string, options *StreamOptions) streamAPI {
			return NewStream(client, id, options)
		},
		timePerBatchPerStream: time.Duration(timePerBatchPerStream)}

	for i := uint(0); i < options.StreamCount; i++ {
		streamOptions := StreamOptions{
			BatchLimit:           options.BatchLimit,
			InitialRetryInterval: options.InitialRetryInterval,
			MaxRetryInterval:     options.MaxRetryInterval,
			CommitMaxElapsedTime: options.CommitMaxElapsedTime,
			NotifyErr:            func(err error, duration time.Duration) { options.NotifyErr(i, err, duration) },
			NotifyOK:             func() { options.NotifyOK(i) }}

		processor.streamOptions = append(processor.streamOptions, streamOptions)
	}

	return processor
}

// A Processor for nakadi events. The Processor is a high level API for consuming events from
// Nakadi. It can process event batches from multiple partitions (streams) and can be configured
// with a certain event rate, that limits the number of processed events per minute. The cursors of
// event batches that were successfully processed are automatically committed.
type Processor struct {
	sync.Mutex
	client                *Client
	subscriptionID        string
	streamOptions         []StreamOptions
	newStream             func(*Client, string, *StreamOptions) streamAPI
	timePerBatchPerStream time.Duration
	streams               []streamAPI
	ctx                   context.Context
	cancel                context.CancelFunc
}

// Start begins event processing. All event batches received from the underlying streams are passed to
// the operation function. If the operation function terminates without error the respective cursor will
// be automatically committed to Nakadi. If the operations terminates with an error, the underlying stream
// will be halted and a new stream will continue to pass event batches to the operation function.
//
// Event processing will go on indefinitely unless the processor is stopped via its Stop method. Star will
// return an error if the processor is already running.
func (p *Processor) Start(operation func(int, []byte) error) error {
	p.Lock()
	defer p.Unlock()

	if len(p.streams) > 0 {
		return errors.New("processor was already started")
	}

	p.streams = make([]streamAPI, len(p.streamOptions))

	for streamNo, options := range p.streamOptions {
		go p.startSingleStream(operation, streamNo, options)
	}

	return nil
}

// startSingleStream starts a single stream with a given stream number / position. After the stream has been
// started it consumes events. In cases of errors the stream is closed and a new stream will be opened.
func (p *Processor) startSingleStream(operation func(int, []byte) error, streamNo int, options StreamOptions) {
	p.streams[streamNo] = p.newStream(p.client, p.subscriptionID, &options)

	if p.timePerBatchPerStream > 0 {
		initialWait := rand.Int63n(int64(p.timePerBatchPerStream))
		select {
		case <-p.ctx.Done():
			return
		case <-time.After(time.Duration(initialWait)):
			// nothing
		}
	}

	for {
		start := time.Now()

		select {
		case <-p.ctx.Done():
			return
		default:
			cursor, events, err := p.streams[streamNo].NextEvents()
			if err != nil {
				continue
			}

			err = operation(streamNo, events)
			if err != nil {
				p.Lock()
				p.streams[streamNo].Close()
				p.streams[streamNo] = p.newStream(p.client, p.subscriptionID, &options)
				p.Unlock()
				continue
			}

			p.streams[streamNo].CommitCursor(cursor)
		}

		elapsed := time.Since(start)
		if elapsed < p.timePerBatchPerStream {
			select {
			case <-p.ctx.Done():
				return
			case <-time.After(p.timePerBatchPerStream - elapsed):
				// nothing
			}
		}
	}
}

// Stop halts all steams and terminates event processing. Stop will return with an error if the processor
// is not running.
func (p *Processor) Stop() error {
	p.Lock()
	defer p.Unlock()

	if len(p.streams) == 0 {
		return errors.New("processor is not running")
	}

	p.cancel()

	var allErrors []error
	for _, s := range p.streams {
		err := s.Close()
		if err != nil {
			allErrors = append(allErrors, err)
		}
	}

	p.streams = nil

	if len(allErrors) > 0 {
		return errors.Errorf("%d streams had errors while closing the stream", len(allErrors))
	}

	return nil
}
<|endoftext|>


<commit_before>package agg

import (
	"github.com/Supernomad/quantum/common"
	"time"
)

// Agg a statistics aggregation object
type Agg struct {
	cfg    *common.Config
	start  time.Time
	stop   chan struct{}
	ticker *time.Ticker

	lastIncomingStats *common.Stats
	lastOutgoingStats *common.Stats

	incomingStats []*common.Stats
	outgoingStats []*common.Stats

	sinks []StatSink
}

// StatSink to dump statistics into
type StatSink interface {
	SendStats(statsl *StatsLog) error
}

func (agg *Agg) aggData() (incomingStats *common.Stats, outgoingStats *common.Stats) {
	incomingStats = common.NewStats()
	outgoingStats = common.NewStats()

	for i := 0; i < agg.cfg.NumWorkers; i++ {
		incomingStats.Packets += agg.incomingStats[i].Packets
		incomingStats.Bytes += agg.incomingStats[i].Bytes

		for k, v := range agg.incomingStats[i].Links {
			if link, ok := incomingStats.Links[k]; !ok {
				incomingStats.Links[k] = &common.Stats{
					Packets: v.Packets,
					Bytes:   v.Bytes,
				}
			} else {
				link.Packets += v.Packets
				link.Bytes += v.Bytes
			}
		}

		outgoingStats.Packets += agg.outgoingStats[i].Packets
		outgoingStats.Bytes += agg.outgoingStats[i].Bytes

		for k, v := range agg.outgoingStats[i].Links {
			if link, ok := outgoingStats.Links[k]; !ok {
				outgoingStats.Links[k] = &common.Stats{
					Packets: v.Packets,
					Bytes:   v.Bytes,
				}
			} else {
				link.Packets += v.Packets
				link.Bytes += v.Bytes
			}
		}
	}
	return
}

func (agg *Agg) diffData(elapsed float64, incomingStats *common.Stats, outgoingStats *common.Stats) {
	incomingStats.PacketsDiff = incomingStats.Packets - agg.lastIncomingStats.Packets
	incomingStats.PPS = float64(incomingStats.PacketsDiff) / elapsed
	incomingStats.BytesDiff = incomingStats.Bytes - agg.lastIncomingStats.Bytes
	incomingStats.Bandwidth = float64(incomingStats.BytesDiff) / elapsed

	for key, incomingLink := range incomingStats.Links {
		if lastIncomingLink, exists := agg.lastIncomingStats.Links[key]; exists {
			incomingLink.PacketsDiff = incomingLink.Packets - lastIncomingLink.Packets
			incomingLink.PPS = float64(incomingLink.PacketsDiff) / elapsed
			incomingLink.BytesDiff = incomingLink.Bytes - lastIncomingLink.Bytes
			incomingLink.Bandwidth = float64(incomingLink.BytesDiff) / elapsed
		} else {
			incomingLink.PacketsDiff = incomingLink.Packets
			incomingLink.PPS = float64(incomingLink.PacketsDiff) / elapsed
			incomingLink.BytesDiff = incomingLink.Bytes
			incomingLink.Bandwidth = float64(incomingLink.BytesDiff) / elapsed
		}
	}

	outgoingStats.PacketsDiff = outgoingStats.Packets - agg.lastOutgoingStats.Packets
	outgoingStats.PPS = float64(outgoingStats.PacketsDiff) / elapsed
	outgoingStats.BytesDiff = outgoingStats.Bytes - agg.lastOutgoingStats.Bytes
	outgoingStats.Bandwidth = float64(outgoingStats.BytesDiff) / elapsed

	for key, outgoingLink := range outgoingStats.Links {
		if lastOutgoingLink, exists := agg.lastOutgoingStats.Links[key]; exists {
			outgoingLink.PacketsDiff = outgoingLink.Packets - lastOutgoingLink.Packets
			outgoingLink.PPS = float64(outgoingLink.PacketsDiff) / elapsed
			outgoingLink.BytesDiff = outgoingLink.Bytes - lastOutgoingLink.Bytes
			outgoingLink.Bandwidth = float64(outgoingLink.BytesDiff) / elapsed
		} else {
			outgoingLink.PacketsDiff = outgoingLink.Packets
			outgoingLink.PPS = float64(outgoingLink.PacketsDiff) / elapsed
			outgoingLink.BytesDiff = outgoingLink.Bytes
			outgoingLink.Bandwidth = float64(outgoingLink.BytesDiff) / elapsed
		}
	}

	return
}

func (agg *Agg) sendData(statsl *StatsLog) error {
	for i := 0; i < len(agg.sinks); i++ {
		if err := agg.sinks[i].SendStats(statsl); err != nil {
			return err
		}
	}
	return nil
}

// Start aggregating and sending stats data
func (agg *Agg) Start() {
	go func() {
		for {
			select {
			case <-agg.stop:
				return
			case <-agg.ticker.C:
				elapsed := time.Since(agg.start)
				elapsedSec := elapsed.Seconds()

				incomingStats, outgoingStats := agg.aggData()
				agg.diffData(elapsedSec, incomingStats, outgoingStats)

				statsl := &StatsLog{
					TimeSpan: elapsedSec,

					TxStats: outgoingStats,
					RxStats: incomingStats,
				}

				agg.lastIncomingStats = incomingStats
				agg.lastOutgoingStats = outgoingStats

				agg.sendData(statsl)
				agg.start = time.Now()
			}
		}
	}()
}

// Stop aggregating and sending stats data
func (agg *Agg) Stop() {
	go func() {
		agg.stop <- struct{}{}
	}()
}

// New Agg instance pointer
func New(cfg *common.Config, incomingStats []*common.Stats, outgoingStats []*common.Stats) *Agg {
	return &Agg{
		cfg:               cfg,
		start:             time.Now(),
		sinks:             []StatSink{&ConsoleSink{}},
		ticker:            time.NewTicker(cfg.StatsWindow * time.Second),
		stop:              make(chan struct{}),
		lastIncomingStats: common.NewStats(),
		lastOutgoingStats: common.NewStats(),
		incomingStats:     incomingStats,
		outgoingStats:     outgoingStats,
	}
}
<commit_msg>Added some basic unit tests for agg<commit_after>package agg

import (
	"github.com/Supernomad/quantum/common"
	"time"
)

// Agg a statistics aggregation object
type Agg struct {
	cfg    *common.Config
	start  time.Time
	stop   chan struct{}
	ticker *time.Ticker

	lastIncomingStats *common.Stats
	lastOutgoingStats *common.Stats

	incomingStats []*common.Stats
	outgoingStats []*common.Stats

	sinks []StatSink
}

// StatSink to dump statistics into
type StatSink interface {
	SendStats(statsl *StatsLog) error
}

func (agg *Agg) aggData() (incomingStats *common.Stats, outgoingStats *common.Stats) {
	incomingStats = common.NewStats()
	outgoingStats = common.NewStats()

	for i := 0; i < agg.cfg.NumWorkers; i++ {
		incomingStats.Packets += agg.incomingStats[i].Packets
		incomingStats.Bytes += agg.incomingStats[i].Bytes

		for k, v := range agg.incomingStats[i].Links {
			if link, ok := incomingStats.Links[k]; !ok {
				incomingStats.Links[k] = &common.Stats{
					Packets: v.Packets,
					Bytes:   v.Bytes,
				}
			} else {
				link.Packets += v.Packets
				link.Bytes += v.Bytes
			}
		}

		outgoingStats.Packets += agg.outgoingStats[i].Packets
		outgoingStats.Bytes += agg.outgoingStats[i].Bytes

		for k, v := range agg.outgoingStats[i].Links {
			if link, ok := outgoingStats.Links[k]; !ok {
				outgoingStats.Links[k] = &common.Stats{
					Packets: v.Packets,
					Bytes:   v.Bytes,
				}
			} else {
				link.Packets += v.Packets
				link.Bytes += v.Bytes
			}
		}
	}
	return
}

func (agg *Agg) diffData(elapsed float64, incomingStats *common.Stats, outgoingStats *common.Stats) {
	incomingStats.PacketsDiff = incomingStats.Packets - agg.lastIncomingStats.Packets
	incomingStats.PPS = float64(incomingStats.PacketsDiff) / elapsed
	incomingStats.BytesDiff = incomingStats.Bytes - agg.lastIncomingStats.Bytes
	incomingStats.Bandwidth = float64(incomingStats.BytesDiff) / elapsed

	for key, incomingLink := range incomingStats.Links {
		if lastIncomingLink, exists := agg.lastIncomingStats.Links[key]; exists {
			incomingLink.PacketsDiff = incomingLink.Packets - lastIncomingLink.Packets
			incomingLink.PPS = float64(incomingLink.PacketsDiff) / elapsed
			incomingLink.BytesDiff = incomingLink.Bytes - lastIncomingLink.Bytes
			incomingLink.Bandwidth = float64(incomingLink.BytesDiff) / elapsed
		} else {
			incomingLink.PacketsDiff = incomingLink.Packets
			incomingLink.PPS = float64(incomingLink.PacketsDiff) / elapsed
			incomingLink.BytesDiff = incomingLink.Bytes
			incomingLink.Bandwidth = float64(incomingLink.BytesDiff) / elapsed
		}
	}

	outgoingStats.PacketsDiff = outgoingStats.Packets - agg.lastOutgoingStats.Packets
	outgoingStats.PPS = float64(outgoingStats.PacketsDiff) / elapsed
	outgoingStats.BytesDiff = outgoingStats.Bytes - agg.lastOutgoingStats.Bytes
	outgoingStats.Bandwidth = float64(outgoingStats.BytesDiff) / elapsed

	for key, outgoingLink := range outgoingStats.Links {
		if lastOutgoingLink, exists := agg.lastOutgoingStats.Links[key]; exists {
			outgoingLink.PacketsDiff = outgoingLink.Packets - lastOutgoingLink.Packets
			outgoingLink.PPS = float64(outgoingLink.PacketsDiff) / elapsed
			outgoingLink.BytesDiff = outgoingLink.Bytes - lastOutgoingLink.Bytes
			outgoingLink.Bandwidth = float64(outgoingLink.BytesDiff) / elapsed
		} else {
			outgoingLink.PacketsDiff = outgoingLink.Packets
			outgoingLink.PPS = float64(outgoingLink.PacketsDiff) / elapsed
			outgoingLink.BytesDiff = outgoingLink.Bytes
			outgoingLink.Bandwidth = float64(outgoingLink.BytesDiff) / elapsed
		}
	}

	return
}

func (agg *Agg) sendData(statsl *StatsLog) error {
	for i := 0; i < len(agg.sinks); i++ {
		if err := agg.sinks[i].SendStats(statsl); err != nil {
			return err
		}
	}
	return nil
}

func (agg *Agg) pipeline() {
	elapsed := time.Since(agg.start)
	elapsedSec := elapsed.Seconds()

	incomingStats, outgoingStats := agg.aggData()
	agg.diffData(elapsedSec, incomingStats, outgoingStats)

	statsl := &StatsLog{
		TimeSpan: elapsedSec,

		TxStats: outgoingStats,
		RxStats: incomingStats,
	}

	agg.lastIncomingStats = incomingStats
	agg.lastOutgoingStats = outgoingStats

	agg.sendData(statsl)
	agg.start = time.Now()
}

// Start aggregating and sending stats data
func (agg *Agg) Start() {
	go func() {
		for {
			select {
			case <-agg.stop:
				return
			case <-agg.ticker.C:
				agg.pipeline()
			}
		}
	}()
}

// Stop aggregating and sending stats data
func (agg *Agg) Stop() {
	go func() {
		agg.stop <- struct{}{}
	}()
}

// New Agg instance pointer
func New(cfg *common.Config, incomingStats []*common.Stats, outgoingStats []*common.Stats) *Agg {
	return &Agg{
		cfg:               cfg,
		start:             time.Now(),
		sinks:             []StatSink{&ConsoleSink{}},
		ticker:            time.NewTicker(cfg.StatsWindow * time.Second),
		stop:              make(chan struct{}),
		lastIncomingStats: common.NewStats(),
		lastOutgoingStats: common.NewStats(),
		incomingStats:     incomingStats,
		outgoingStats:     outgoingStats,
	}
}
<|endoftext|>


<commit_before>// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see <http://www.gnu.org/licenses/>.

package spacestatus

import (
	"encoding/json"
	"github.com/nmeum/marvin/irc"
	"github.com/nmeum/marvin/modules"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

type spaceapi struct {
	space string `json:"space"`
	state struct {
		open bool `json:"open"`
	} `json:"state"`
}

type Module struct {
	api      *spaceapi
	URL      string `json:"urls"`
	Notify   bool   `json:"notify"`
	Interval string `json:"interval"`
}

func Init(moduleSet *modules.ModuleSet) {
	moduleSet.Register(new(Module))
}

func (m *Module) Name() string {
	return "spacestatus"
}

func (m *Module) Help() string {
	return "USAGE: !spacestatus"
}

func (m *Module) Defaults() {
	m.Notify = true
	m.Interval = "0h15m"
}

func (m *Module) Load(client *irc.Client) error {
	if len(m.URL) <= 0 {
		return nil
	}

	client.CmdHook("privmsg", m.statusCmd)
	if !m.Notify {
		duration, err := time.ParseDuration(m.Interval)
		if err != nil {
			return err
		}

		go func(c *irc.Client) {
			for {
				m.updateHandler(c)
				time.Sleep(duration)
			}
		}(client)
	}

	return nil
}

func (m *Module) updateHandler(client *irc.Client) error {
	var oldState bool
	if m.api == nil {
		oldState = false
	} else {
		oldState = m.api.state.open
	}

	if err := m.pollStatus(); err != nil {
		return err
	}

	newState := m.api.state.open
	if newState != oldState {
		m.notify(client, newState)
	}

	return nil
}

func (m *Module) pollStatus() error {
	resp, err := http.Get(m.URL)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	if err := json.Unmarshal(data, m.api); err != nil {
		return err
	}

	return nil
}

func (m *Module) statusCmd(client *irc.Client, msg irc.Message) error {
	if msg.Data != "!spacestatus" {
		return nil
	}

	var state string
	if m.api.state.open {
		state = "open"
	} else {
		state = "closed"
	}

	return client.Write("NOTICE %s: %s is currently %s",
		msg.Receiver, m.api.space, state)
}

func (m *Module) notify(client *irc.Client, open bool) {
	name := m.api.space
	for _, ch := range client.Channels {
		var oldState, newState string
		if open {
			oldState = "closed"
			newState = "open"
		} else {
			oldState = "open"
			newState = "closed"
		}

		client.Write("NOTICE %s: %s -- %s changed state from %s to %s",
			ch, strings.ToUpper(m.Name()), name, oldState, newState)
	}
}
<commit_msg>spacestatus: fix notification disablement<commit_after>// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see <http://www.gnu.org/licenses/>.

package spacestatus

import (
	"encoding/json"
	"github.com/nmeum/marvin/irc"
	"github.com/nmeum/marvin/modules"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

type spaceapi struct {
	space string `json:"space"`
	state struct {
		open bool `json:"open"`
	} `json:"state"`
}

type Module struct {
	api      *spaceapi
	URL      string `json:"urls"`
	Notify   bool   `json:"notify"`
	Interval string `json:"interval"`
}

func Init(moduleSet *modules.ModuleSet) {
	moduleSet.Register(new(Module))
}

func (m *Module) Name() string {
	return "spacestatus"
}

func (m *Module) Help() string {
	return "USAGE: !spacestatus"
}

func (m *Module) Defaults() {
	m.Notify = true
	m.Interval = "0h15m"
}

func (m *Module) Load(client *irc.Client) error {
	if len(m.URL) <= 0 {
		return nil
	}

	duration, err := time.ParseDuration(m.Interval)
	if err != nil {
		return err
	}

	go func(c *irc.Client) {
		for {
			m.updateHandler(c)
			time.Sleep(duration)
		}
	}(client)

	client.CmdHook("privmsg", m.statusCmd)
	return nil
}

func (m *Module) updateHandler(client *irc.Client) error {
	var oldState bool
	if m.api == nil {
		oldState = false
	} else {
		oldState = m.api.state.open
	}

	if err := m.pollStatus(); err != nil {
		return err
	}

	newState := m.api.state.open
	if newState != oldState && m.Notify {
		m.notify(client, newState)
	}

	return nil
}

func (m *Module) pollStatus() error {
	resp, err := http.Get(m.URL)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	if err := json.Unmarshal(data, m.api); err != nil {
		return err
	}

	return nil
}

func (m *Module) statusCmd(client *irc.Client, msg irc.Message) error {
	if msg.Data != "!spacestatus" {
		return nil
	}

	if m.api == nil {
		return client.Write("NOTICE %s: Status is currently unknown.")
	}

	var state string
	if m.api.state.open {
		state = "open"
	} else {
		state = "closed"
	}

	return client.Write("NOTICE %s: %s is currently %s",
		msg.Receiver, m.api.space, state)
}

func (m *Module) notify(client *irc.Client, open bool) {
	name := m.api.space
	for _, ch := range client.Channels {
		var oldState, newState string
		if open {
			oldState = "closed"
			newState = "open"
		} else {
			oldState = "open"
			newState = "closed"
		}

		client.Write("NOTICE %s: %s -- %s changed state from %s to %s",
			ch, strings.ToUpper(m.Name()), name, oldState, newState)
	}
}
<|endoftext|>


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

package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/brutella/dnssd"
	"github.com/google/renameio"
	"github.com/stapelberg/airscan"
	"github.com/stapelberg/airscan/preset"
	"github.com/stapelberg/scan2drive/internal/dispatch"
	"github.com/stapelberg/scan2drive/proto"
	"golang.org/x/net/context"
	"golang.org/x/net/trace"

	// We request a jpeg image from the scanner, so make sure the image package
	// can decode jpeg images in later processing stages.
	_ "image/jpeg"
)

type pushToAirscan struct {
	name    string
	host    string
	iconURL string
}

func (p *pushToAirscan) Display() dispatch.Display {
	name := p.name
	if name == "" {
		name = "AirScan: unknown target"
	}
	return dispatch.Display{
		Name:    name,
		IconURL: p.iconURL,
	}
}

func (p *pushToAirscan) Scan(user string) (string, error) {
	tr := trace.New("AirScan", p.host)
	defer tr.Finish()

	client := proto.NewScanClient(scanConn)
	resp, err := client.DefaultUser(context.Background(), &proto.DefaultUserRequest{})
	if err != nil {
		return "", err
	}

	tr.LazyPrintf("Starting AirScan at host %s.local for user %s", p.host, resp.User)

	return scanA(tr, resp.User, p.host)
}

func Airscan() {
	tr := trace.New("AirScan", "DNSSD")
	defer tr.Finish()

	log.Printf("Searching for AirScan scanners via DNSSD")

	actionByHost := make(map[string]*pushToAirscan)

	addFn := func(srv dnssd.Service) {
		log.Printf("DNSSD service discovered: %v", srv)

		// 2020/07/16 13:54:34 add: {
		// Name: Brother\ MFC-L2750DW\ series
		// Type: _uscan._tcp
		// Domain: local
		// Host: BRN3C2xxxxxxxxx
		// Text: map[
		//   UUID:yyyyyyyy-xxxx-zzzz-aaaa-bbbbbbbbbbbb
		//   adminurl:http://BRN3C2xxxxxxxxx.local./net/net/airprint.html
		//   cs:binary,grayscale,color
		//   duplex:T
		//   is:adf,platen
		//   note:
		//   pdl:application/pdf,image/jpeg
		//   representation:http://BRN3C2xxxxxxxxx.local./icons/device-icons-128.png
		//   rs:eSCL
		//   txtvers:1
		//   ty:Brother
		//   vers:2.63]
		// TTL: 1h15m0s
		// Port: 80
		// IPs: [10.0.0.12 fe80::3e2a:aaaa:bbbb:cccc]
		// IfaceIPs:map[]

		// miekg/dns escapes characters in DNS labels, which as per RFC1034 and
		// RFC1035 does not actually permit whitespace. The purpose of escaping
		// originally appears to be to use these labels in a DNS master file,
		// but for our UI, the backslashes look just wrong.
		unescapedName := strings.ReplaceAll(srv.Name, "\\", "")

		// TODO: use srv.Text["ty"] if non-empty once
		// https://github.com/brutella/dnssd/pull/19 was merged

		if action, ok := actionByHost[srv.Host]; !ok {
			action = &pushToAirscan{
				name:    unescapedName,
				host:    srv.Host,
				iconURL: srv.Text["representation"],
			}
			actionByHost[srv.Host] = action
			// Register the default scan action (physical hardware button, or web
			// interface scan button) to be scanning from this scanner via AirScan:
			dispatch.Register(action)
		}
	}
	rmvFn := func(srv dnssd.Service) {
		log.Printf("remove: %v", srv)
		if action, ok := actionByHost[srv.Host]; ok {
			dispatch.Unregister(action)
			delete(actionByHost, srv.Host)
		}
	}
	const service = "_uscan._tcp.local." // AirScan DNSSD service name
	// addFn and rmvFn are always called (sequentially) from the same goroutine,
	// i.e. no locking is required.
	if err := dnssd.LookupType(context.Background(), service, addFn, rmvFn); err != nil {
		log.Printf("DNSSD init failed: %v", err)
		return
	}

}

func scanA(tr trace.Trace, user, host string) (string, error) {
	start := time.Now()

	relName := time.Now().Format(time.RFC3339)
	scanDir := filepath.Join(*scansDir, user, relName)

	if err := os.MkdirAll(scanDir, 0700); err != nil {
		return "", err
	}

	if err := scan1(tr, host, scanDir, relName, user); err != nil {
		return "", err
	}

	tr.LazyPrintf("scan done in %v", time.Since(start))

	// NOTE: We currently call out to the generic ProcessScan() implementation,
	// whereas the Fujitsu ScanSnap iX500 (fss500) specific implementation did a
	// lot of the work itself, resulting in significant speed-ups on the
	// Raspberry Pi 3.
	//
	// We should investigate whether similar speed-ups are required/achievable
	// with the Raspberry Pi 4 and AirScan.

	tr.LazyPrintf("processing scan")
	client := proto.NewScanClient(scanConn)
	if _, err := client.ProcessScan(context.Background(), &proto.ProcessScanRequest{User: user, Dir: relName}); err != nil {
		return "", err
	}
	tr.LazyPrintf("scan processed")

	return relName, nil
}

func scan1(tr trace.Trace, host, scanDir, relName, respUser string) error {
	cl := airscan.NewClient(host)

	status, err := cl.ScannerStatus()
	if err != nil {
		return err
	}
	tr.LazyPrintf("status: %+v", status)
	if got, want := status.State, "Idle"; got != want {
		return fmt.Errorf("scanner not ready: in state %q, want %q", got, want)
	}

	// NOTE: With the Fujitsu ScanSnap iX500 (fss500), we always scanned 600
	// dpi color full-duplex and retained the highest-quality
	// originals. Other scanners (e.g. the Brother MFC-L2750DW) take A LOT
	// longer than the fss500 for scanning duplex color originals, so with
	// AirScan, we play it safe and only request what we really need:
	// grayscale at 300 dpi.

	settings := preset.GrayscaleA4ADF()
	// For the ADF, the ScanSnap is better.
	// We use the Brother for its flatbed scan only.
	settings.InputSource = "Platen"
	scan, err := cl.Scan(settings)
	if err != nil {
		return err
	}
	defer func() {
		tr.LazyPrintf("Deleting ScanJob %s", scan)
		if err := scan.Close(); err != nil {
			log.Printf("error deleting AirScan job (probably harmless): %v", err)
		}
	}()
	tr.LazyPrintf("ScanJob created: %s", scan)

	pagenum := 1
	for scan.ScanPage() {
		fn := filepath.Join(scanDir, fmt.Sprintf("page%d.jpg", pagenum))

		o, err := renameio.TempFile("", fn)
		if err != nil {
			return err
		}
		defer o.Cleanup()

		if _, err := io.Copy(o, scan.CurrentPage()); err != nil {
			return err
		}

		if err := o.CloseAtomicallyReplace(); err != nil {
			return err
		}

		pagenum++
	}
	if err := scan.Err(); err != nil {
		return err
	}

	createCompleteMarker(respUser, relName, "scan")

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

package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/brutella/dnssd"
	"github.com/google/renameio"
	"github.com/stapelberg/airscan"
	"github.com/stapelberg/airscan/preset"
	"github.com/stapelberg/scan2drive/internal/dispatch"
	"github.com/stapelberg/scan2drive/proto"
	"golang.org/x/net/context"
	"golang.org/x/net/trace"

	// We request a jpeg image from the scanner, so make sure the image package
	// can decode jpeg images in later processing stages.
	_ "image/jpeg"
)

type pushToAirscan struct {
	name    string
	host    string
	iconURL string
}

func (p *pushToAirscan) Display() dispatch.Display {
	name := p.name
	if name == "" {
		name = "AirScan: unknown target"
	}
	return dispatch.Display{
		Name:    name,
		IconURL: p.iconURL,
	}
}

func (p *pushToAirscan) Scan(user string) (string, error) {
	tr := trace.New("AirScan", p.host)
	defer tr.Finish()

	tr.LazyPrintf("Starting AirScan at host %s.local for user %s", p.host, user)

	return scanA(tr, user, p.host)
}

func Airscan() {
	tr := trace.New("AirScan", "DNSSD")
	defer tr.Finish()

	log.Printf("Searching for AirScan scanners via DNSSD")

	actionByHost := make(map[string]*pushToAirscan)

	addFn := func(srv dnssd.Service) {
		log.Printf("DNSSD service discovered: %v", srv)

		// 2020/07/16 13:54:34 add: {
		// Name: Brother\ MFC-L2750DW\ series
		// Type: _uscan._tcp
		// Domain: local
		// Host: BRN3C2xxxxxxxxx
		// Text: map[
		//   UUID:yyyyyyyy-xxxx-zzzz-aaaa-bbbbbbbbbbbb
		//   adminurl:http://BRN3C2xxxxxxxxx.local./net/net/airprint.html
		//   cs:binary,grayscale,color
		//   duplex:T
		//   is:adf,platen
		//   note:
		//   pdl:application/pdf,image/jpeg
		//   representation:http://BRN3C2xxxxxxxxx.local./icons/device-icons-128.png
		//   rs:eSCL
		//   txtvers:1
		//   ty:Brother
		//   vers:2.63]
		// TTL: 1h15m0s
		// Port: 80
		// IPs: [10.0.0.12 fe80::3e2a:aaaa:bbbb:cccc]
		// IfaceIPs:map[]

		// miekg/dns escapes characters in DNS labels, which as per RFC1034 and
		// RFC1035 does not actually permit whitespace. The purpose of escaping
		// originally appears to be to use these labels in a DNS master file,
		// but for our UI, the backslashes look just wrong.
		unescapedName := strings.ReplaceAll(srv.Name, "\\", "")

		// TODO: use srv.Text["ty"] if non-empty once
		// https://github.com/brutella/dnssd/pull/19 was merged

		if action, ok := actionByHost[srv.Host]; !ok {
			action = &pushToAirscan{
				name:    unescapedName,
				host:    srv.Host,
				iconURL: srv.Text["representation"],
			}
			actionByHost[srv.Host] = action
			// Register the default scan action (physical hardware button, or web
			// interface scan button) to be scanning from this scanner via AirScan:
			dispatch.Register(action)
		}
	}
	rmvFn := func(srv dnssd.Service) {
		log.Printf("remove: %v", srv)
		if action, ok := actionByHost[srv.Host]; ok {
			dispatch.Unregister(action)
			delete(actionByHost, srv.Host)
		}
	}
	const service = "_uscan._tcp.local." // AirScan DNSSD service name
	// addFn and rmvFn are always called (sequentially) from the same goroutine,
	// i.e. no locking is required.
	if err := dnssd.LookupType(context.Background(), service, addFn, rmvFn); err != nil {
		log.Printf("DNSSD init failed: %v", err)
		return
	}

}

func scanA(tr trace.Trace, user, host string) (string, error) {
	start := time.Now()

	relName := time.Now().Format(time.RFC3339)
	scanDir := filepath.Join(*scansDir, user, relName)

	if err := os.MkdirAll(scanDir, 0700); err != nil {
		return "", err
	}

	if err := scan1(tr, host, scanDir, relName, user); err != nil {
		return "", err
	}

	tr.LazyPrintf("scan done in %v", time.Since(start))

	// NOTE: We currently call out to the generic ProcessScan() implementation,
	// whereas the Fujitsu ScanSnap iX500 (fss500) specific implementation did a
	// lot of the work itself, resulting in significant speed-ups on the
	// Raspberry Pi 3.
	//
	// We should investigate whether similar speed-ups are required/achievable
	// with the Raspberry Pi 4 and AirScan.

	tr.LazyPrintf("processing scan")
	client := proto.NewScanClient(scanConn)
	if _, err := client.ProcessScan(context.Background(), &proto.ProcessScanRequest{User: user, Dir: relName}); err != nil {
		return "", err
	}
	tr.LazyPrintf("scan processed")

	return relName, nil
}

func scan1(tr trace.Trace, host, scanDir, relName, respUser string) error {
	cl := airscan.NewClient(host)

	status, err := cl.ScannerStatus()
	if err != nil {
		return err
	}
	tr.LazyPrintf("status: %+v", status)
	if got, want := status.State, "Idle"; got != want {
		return fmt.Errorf("scanner not ready: in state %q, want %q", got, want)
	}

	// NOTE: With the Fujitsu ScanSnap iX500 (fss500), we always scanned 600
	// dpi color full-duplex and retained the highest-quality
	// originals. Other scanners (e.g. the Brother MFC-L2750DW) take A LOT
	// longer than the fss500 for scanning duplex color originals, so with
	// AirScan, we play it safe and only request what we really need:
	// grayscale at 300 dpi.

	settings := preset.GrayscaleA4ADF()
	// For the ADF, the ScanSnap is better.
	// We use the Brother for its flatbed scan only.
	settings.InputSource = "Platen"
	scan, err := cl.Scan(settings)
	if err != nil {
		return err
	}
	defer func() {
		tr.LazyPrintf("Deleting ScanJob %s", scan)
		if err := scan.Close(); err != nil {
			log.Printf("error deleting AirScan job (probably harmless): %v", err)
		}
	}()
	tr.LazyPrintf("ScanJob created: %s", scan)

	pagenum := 1
	for scan.ScanPage() {
		fn := filepath.Join(scanDir, fmt.Sprintf("page%d.jpg", pagenum))

		o, err := renameio.TempFile("", fn)
		if err != nil {
			return err
		}
		defer o.Cleanup()

		if _, err := io.Copy(o, scan.CurrentPage()); err != nil {
			return err
		}

		if err := o.CloseAtomicallyReplace(); err != nil {
			return err
		}

		pagenum++
	}
	if err := scan.Err(); err != nil {
		return err
	}

	createCompleteMarker(respUser, relName, "scan")

	return nil
}
<|endoftext|>


<commit_before>package network

import (
	"github.com/gorilla/websocket"
	"github.com/name5566/leaf/log"
	"net"
	"net/http"
	"sync"
	"time"
)

type WSServer struct {
	Addr            string
	MaxConnNum      int
	PendingWriteNum int
	NewAgent        func(*WSConn) Agent
	HTTPTimeout     time.Duration
	ln              net.Listener
	handler         *WSHandler
}

type WSHandler struct {
	maxConnNum      int
	pendingWriteNum int
	upgrader        websocket.Upgrader
	conns           WebsocketConnSet
	mutexConns      sync.Mutex
	wg              sync.WaitGroup
}

func (handler *WSHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", 405)
		return
	}
	conn, err := handler.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Debug("upgrade error: %v", err)
		return
	}
}

func (server *WSServer) Start() {
	ln, err := net.Listen("tcp", server.Addr)
	if err != nil {
		log.Fatal("%v", err)
	}

	if server.MaxConnNum <= 0 {
		server.MaxConnNum = 100
		log.Release("invalid MaxConnNum, reset to %v", server.MaxConnNum)
	}
	if server.PendingWriteNum <= 0 {
		server.PendingWriteNum = 100
		log.Release("invalid PendingWriteNum, reset to %v", server.PendingWriteNum)
	}
	if server.NewAgent == nil {
		log.Fatal("NewAgent must not be nil")
	}
	if server.HTTPTimeout <= 0 {
		server.HTTPTimeout = 10 * time.Second
		log.Release("invalid HTTPTimeout, reset to %v", server.HTTPTimeout)
	}

	server.ln = ln
	server.handler = &WSHandler{
		maxConnNum:      server.MaxConnNum,
		pendingWriteNum: server.PendingWriteNum,
		conns:           make(WebsocketConnSet),
		upgrader: websocket.Upgrader{
			HandshakeTimeout: server.HTTPTimeout,
			CheckOrigin:      func(_ *http.Request) bool { return true },
		},
	}

	httpServer := &http.Server{
		Addr:           server.Addr,
		Handler:        server.handler,
		ReadTimeout:    server.HTTPTimeout,
		WriteTimeout:   server.HTTPTimeout,
		MaxHeaderBytes: 1024,
	}

	go httpServer.Serve(ln)
}

func (server *WSServer) Close() {
	server.ln.Close()

	server.handler.mutexConns.Lock()
	for conn := range server.handler.conns {
		conn.Close()
	}
	server.handler.conns = nil
	server.handler.mutexConns.Unlock()

	server.handler.wg.Wait()
}
<commit_msg>manage active conns<commit_after>package network

import (
	"github.com/gorilla/websocket"
	"github.com/name5566/leaf/log"
	"net"
	"net/http"
	"sync"
	"time"
)

type WSServer struct {
	Addr            string
	MaxConnNum      int
	PendingWriteNum int
	NewAgent        func(*WSConn) Agent
	HTTPTimeout     time.Duration
	ln              net.Listener
	handler         *WSHandler
}

type WSHandler struct {
	maxConnNum      int
	pendingWriteNum int
	upgrader        websocket.Upgrader
	conns           WebsocketConnSet
	mutexConns      sync.Mutex
	wg              sync.WaitGroup
}

func (handler *WSHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", 405)
		return
	}
	conn, err := handler.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Debug("upgrade error: %v", err)
		return
	}

	handler.wg.Add(1)
	defer handler.wg.Done()

	handler.mutexConns.Lock()
	if handler.conns == nil {
		handler.mutexConns.Unlock()
		conn.Close()
		return
	}
	if len(handler.conns) >= handler.maxConnNum {
		handler.mutexConns.Unlock()
		conn.Close()
		log.Debug("too many connections")
		return
	}
	handler.conns[conn] = struct{}{}
	handler.mutexConns.Unlock()
}

func (server *WSServer) Start() {
	ln, err := net.Listen("tcp", server.Addr)
	if err != nil {
		log.Fatal("%v", err)
	}

	if server.MaxConnNum <= 0 {
		server.MaxConnNum = 100
		log.Release("invalid MaxConnNum, reset to %v", server.MaxConnNum)
	}
	if server.PendingWriteNum <= 0 {
		server.PendingWriteNum = 100
		log.Release("invalid PendingWriteNum, reset to %v", server.PendingWriteNum)
	}
	if server.NewAgent == nil {
		log.Fatal("NewAgent must not be nil")
	}
	if server.HTTPTimeout <= 0 {
		server.HTTPTimeout = 10 * time.Second
		log.Release("invalid HTTPTimeout, reset to %v", server.HTTPTimeout)
	}

	server.ln = ln
	server.handler = &WSHandler{
		maxConnNum:      server.MaxConnNum,
		pendingWriteNum: server.PendingWriteNum,
		conns:           make(WebsocketConnSet),
		upgrader: websocket.Upgrader{
			HandshakeTimeout: server.HTTPTimeout,
			CheckOrigin:      func(_ *http.Request) bool { return true },
		},
	}

	httpServer := &http.Server{
		Addr:           server.Addr,
		Handler:        server.handler,
		ReadTimeout:    server.HTTPTimeout,
		WriteTimeout:   server.HTTPTimeout,
		MaxHeaderBytes: 1024,
	}

	go httpServer.Serve(ln)
}

func (server *WSServer) Close() {
	server.ln.Close()

	server.handler.mutexConns.Lock()
	for conn := range server.handler.conns {
		conn.Close()
	}
	server.handler.conns = nil
	server.handler.mutexConns.Unlock()

	server.handler.wg.Wait()
}
<|endoftext|>


<commit_before>package core

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/url"

	"github.com/FrenchBen/oracle-sdk-go/bmc"
	"github.com/Sirupsen/logrus"
)

// SecurityList contains the SecurityList reference from: https://docs.us-phoenix-1.oraclecloud.com/api/#/en/iaas/20160918/SecurityList/
type SecurityList struct {
	// The OCID of the compartment that contains the security list
	CompartmentID string `json:"compartmentId,omitempty"`
	// A user-friendly name
	DisplayName string `json:"displayName",omitempty`
	// The OCID of the SecurityList
	ID string `json:"id,omitempty"`
	// Rules for allowing egress IP packets
	EgressSecurityRules *[]EgressSecurityRule `json:"egressSecurityRules"`
	// Rules for allowing ingress IP packets
	IngressSecurityRules *[]IngressSecurityRule `json:"ingressSecurityRules"`
	// The SecurityList's current state
	LifeCycleState string `json:"lifecycleState,omitempty"`
	// The date and time the instance was created (RFC3339)
	TimeCreated string `json:"timeCreated,omitempty"`
	// The OCID of the VCN
	VcnID string `json:"vcnId",omitempty`
}

// EgressSecurityRule rule for allowing outbound IP packets
type EgressSecurityRule struct {
	Destination string      `json:"destination"`
	IcmpOptions *IcmpOption `json:"icmpOptions,omitempty"`
	IsStateless bool        `json:"isStateless,omitempty"`
	// Protocol values: all, ICMP ("1"), TCP ("6"), UDP ("17").
	Protocol   string      `json:"protocol"`
	TcpOptions *PortConfig `json:"tcpOptions,omitempty"`
	UdpOptions *PortConfig `json:"udpOptions,omitempty"`
}

// IngressSecurityRule rule for allowing inbound IP packets
type IngressSecurityRule struct {
	Source      string      `json:"source"`
	IcmpOptions *IcmpOption `json:"icmpOptions,omitempty"`
	IsStateless bool        `json:"isStateless,omitempty"`
	// Protocol values: all, ICMP ("1"), TCP ("6"), UDP ("17").
	Protocol   string      `json:"protocol"`
	TcpOptions *PortConfig `json:"tcpOptions,omitempty"`
	UdpOptions *PortConfig `json:"udpOptions,omitempty"`
}

type IcmpOption struct {
	// The ICMP code
	Code int `json:"code,omitempty"`
	// The ICMP type
	Type int `json:"type,omitempty"`
}

type PortConfig struct {
	DestinationPortRange *PortRange `json:"destinationPortRange,omitempty"`
	SourcePortRange      *PortRange `json:"sourcePortRange,omitempty"`
}

type PortRange struct {
	Min int `json:"min,omitempty"`
	Max int `json:"max,omitempty"`
}

// GetSecurityList returns a struct of a SecurityList request given an SecurityList ID
func (c *Client) GetSecurityList(securityListID string) (SecurityList, *bmc.Error) {
	securityList := SecurityList{}
	queryString := url.QueryEscape(securityListID)
	resp, err := c.Request("GET", "/securityLists/"+queryString, nil)
	if err != nil {
		logrus.Error(err)
		bmcError := bmc.Error{Code: string(resp.StatusCode), Message: err.Error()}
		return securityList, &bmcError
	}
	logrus.Debug("StatusCode: ", resp.StatusCode)
	if resp.StatusCode != 200 {
		return securityList, bmc.NewError(*resp)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	logrus.Debug("Body: ", string(body))
	if err != nil {
		logrus.Fatalf("Could not read JSON response: %s", err)
	}
	if err = json.Unmarshal(body, &securityList); err != nil {
		logrus.Fatalf("Unmarshal impossible: %s", err)
	}
	return securityList, nil
}

// ListSecurityLists returns a slice struct of all securityList
func (c *Client) ListSecurityLists(vcnID string) ([]SecurityList, *bmc.Error) {
	securityLists := []SecurityList{}
	queryString := url.QueryEscape(c.CompartmentID)
	if vcnID != "" {
		queryString = queryString + "&vcnId=" + url.QueryEscape(vcnID)
	}
	resp, err := c.Request("GET", fmt.Sprintf("/securityLists?compartmentId=%s", queryString), nil)
	if err != nil {
		logrus.Error(err)
		bmcError := bmc.Error{Code: string(resp.StatusCode), Message: err.Error()}
		return securityLists, &bmcError
	}
	logrus.Debug("StatusCode: ", resp.StatusCode)
	if resp.StatusCode != 200 {
		return securityLists, bmc.NewError(*resp)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		logrus.Fatalf("Could not read JSON response: %s", err)
	}

	logrus.Debug("Body: ", string(body))
	if err = json.Unmarshal(body, &securityLists); err != nil {
		logrus.Fatalf("Unmarshal impossible: %s", err)
	}
	return securityLists, nil
}

// CreateSecurityList creates a new security list for the specified VCN
func (c *Client) CreateSecurityList(securityList *SecurityList) (bool, *bmc.Error) {
	resp, err := c.Request("POST", fmt.Sprintf("/securityLists/"), *securityList)
	if err != nil {
		logrus.Error(err)
		bmcError := bmc.Error{Code: string(resp.StatusCode), Message: err.Error()}
		return false, &bmcError
	}
	logrus.Debug("StatusCode: ", resp.StatusCode)
	if resp.StatusCode != 200 {
		return false, bmc.NewError(*resp)
	}
	return true, nil
}

// UpdateSecurityList creates a new security list for the specified VCN
func (c *Client) UpdateSecurityList(securityListID string, securityList *SecurityList) (bool, *bmc.Error) {
	securityListID = url.PathEscape(securityListID)
	resp, err := c.Request("PUT", fmt.Sprintf("/securityLists/%s", securityListID), *securityList)
	if err != nil {
		logrus.Error(err)
		bmcError := bmc.Error{Code: string(resp.StatusCode), Message: err.Error()}
		return false, &bmcError
	}
	logrus.Debug("StatusCode: ", resp.StatusCode)
	if resp.StatusCode != 200 {
		return false, bmc.NewError(*resp)
	}
	return true, nil
}

// DeleteSecurityList deletes the specified security list, but only if it's not associated with a subnet
func (c *Client) DeleteSecurityList(securityListID string) (bool, *bmc.Error) {
	securityListID = url.PathEscape(securityListID)
	resp, err := c.Request("DELETE", fmt.Sprintf("/securityLists/%s", securityListID), nil)
	if err != nil {
		logrus.Error(err)
		bmcError := bmc.Error{Code: string(resp.StatusCode), Message: err.Error()}
		return false, &bmcError
	}
	logrus.Debug("StatusCode: ", resp.StatusCode)
	if resp.StatusCode != 204 {
		return false, bmc.NewError(*resp)
	}
	return true, nil
}
<commit_msg>Added struct for update only on SecurityList<commit_after>package core

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/url"

	"github.com/FrenchBen/oracle-sdk-go/bmc"
	"github.com/Sirupsen/logrus"
)

// SecurityList contains the SecurityList reference from: https://docs.us-phoenix-1.oraclecloud.com/api/#/en/iaas/20160918/SecurityList/
type SecurityList struct {
	// The OCID of the compartment that contains the security list
	CompartmentID string `json:"compartmentId,omitempty"`
	// A user-friendly name
	DisplayName string `json:"displayName",omitempty`
	// The OCID of the SecurityList
	ID string `json:"id,omitempty"`
	// Rules for allowing egress IP packets
	EgressSecurityRules *[]EgressSecurityRule `json:"egressSecurityRules"`
	// Rules for allowing ingress IP packets
	IngressSecurityRules *[]IngressSecurityRule `json:"ingressSecurityRules"`
	// The SecurityList's current state
	LifeCycleState string `json:"lifecycleState,omitempty"`
	// The date and time the instance was created (RFC3339)
	TimeCreated string `json:"timeCreated,omitempty"`
	// The OCID of the VCN
	VcnID string `json:"vcnId",omitempty`
}

type SecurityListUpdate struct {
	// A user-friendly name
	DisplayName string `json:"displayName",omitempty`
	// Rules for allowing egress IP packets
	EgressSecurityRules *[]EgressSecurityRule `json:"egressSecurityRules"`
	// Rules for allowing ingress IP packets
	IngressSecurityRules *[]IngressSecurityRule `json:"ingressSecurityRules"`
}

// EgressSecurityRule rule for allowing outbound IP packets
type EgressSecurityRule struct {
	Destination string      `json:"destination"`
	IcmpOptions *IcmpOption `json:"icmpOptions,omitempty"`
	IsStateless bool        `json:"isStateless,omitempty"`
	// Protocol values: all, ICMP ("1"), TCP ("6"), UDP ("17").
	Protocol   string      `json:"protocol"`
	TcpOptions *PortConfig `json:"tcpOptions,omitempty"`
	UdpOptions *PortConfig `json:"udpOptions,omitempty"`
}

// IngressSecurityRule rule for allowing inbound IP packets
type IngressSecurityRule struct {
	Source      string      `json:"source"`
	IcmpOptions *IcmpOption `json:"icmpOptions,omitempty"`
	IsStateless bool        `json:"isStateless,omitempty"`
	// Protocol values: all, ICMP ("1"), TCP ("6"), UDP ("17").
	Protocol   string      `json:"protocol"`
	TcpOptions *PortConfig `json:"tcpOptions,omitempty"`
	UdpOptions *PortConfig `json:"udpOptions,omitempty"`
}

type IcmpOption struct {
	// The ICMP code
	Code int `json:"code,omitempty"`
	// The ICMP type
	Type int `json:"type,omitempty"`
}

type PortConfig struct {
	DestinationPortRange *PortRange `json:"destinationPortRange,omitempty"`
	SourcePortRange      *PortRange `json:"sourcePortRange,omitempty"`
}

type PortRange struct {
	Min int `json:"min,omitempty"`
	Max int `json:"max,omitempty"`
}

// GetSecurityList returns a struct of a SecurityList request given an SecurityList ID
func (c *Client) GetSecurityList(securityListID string) (SecurityList, *bmc.Error) {
	securityList := SecurityList{}
	queryString := url.QueryEscape(securityListID)
	resp, err := c.Request("GET", "/securityLists/"+queryString, nil)
	if err != nil {
		logrus.Error(err)
		bmcError := bmc.Error{Code: string(resp.StatusCode), Message: err.Error()}
		return securityList, &bmcError
	}
	logrus.Debug("StatusCode: ", resp.StatusCode)
	if resp.StatusCode != 200 {
		return securityList, bmc.NewError(*resp)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	logrus.Debug("Body: ", string(body))
	if err != nil {
		logrus.Fatalf("Could not read JSON response: %s", err)
	}
	if err = json.Unmarshal(body, &securityList); err != nil {
		logrus.Fatalf("Unmarshal impossible: %s", err)
	}
	return securityList, nil
}

// ListSecurityLists returns a slice struct of all securityList
func (c *Client) ListSecurityLists(vcnID string) ([]SecurityList, *bmc.Error) {
	securityLists := []SecurityList{}
	queryString := url.QueryEscape(c.CompartmentID)
	if vcnID != "" {
		queryString = queryString + "&vcnId=" + url.QueryEscape(vcnID)
	}
	resp, err := c.Request("GET", fmt.Sprintf("/securityLists?compartmentId=%s", queryString), nil)
	if err != nil {
		logrus.Error(err)
		bmcError := bmc.Error{Code: string(resp.StatusCode), Message: err.Error()}
		return securityLists, &bmcError
	}
	logrus.Debug("StatusCode: ", resp.StatusCode)
	if resp.StatusCode != 200 {
		return securityLists, bmc.NewError(*resp)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		logrus.Fatalf("Could not read JSON response: %s", err)
	}

	logrus.Debug("Body: ", string(body))
	if err = json.Unmarshal(body, &securityLists); err != nil {
		logrus.Fatalf("Unmarshal impossible: %s", err)
	}
	return securityLists, nil
}

// CreateSecurityList creates a new security list for the specified VCN
func (c *Client) CreateSecurityList(securityList *SecurityList) (bool, *bmc.Error) {
	resp, err := c.Request("POST", fmt.Sprintf("/securityLists/"), *securityList)
	if err != nil {
		logrus.Error(err)
		bmcError := bmc.Error{Code: string(resp.StatusCode), Message: err.Error()}
		return false, &bmcError
	}
	logrus.Debug("StatusCode: ", resp.StatusCode)
	if resp.StatusCode != 200 {
		return false, bmc.NewError(*resp)
	}
	return true, nil
}

// UpdateSecurityList creates a new security list for the specified VCN
func (c *Client) UpdateSecurityList(securityListID string, securityListUpdate *SecurityListUpdate) (bool, *bmc.Error) {
	securityListID = url.PathEscape(securityListID)
	resp, err := c.Request("PUT", fmt.Sprintf("/securityLists/%s", securityListID), *securityListUpdate)
	if err != nil {
		logrus.Error(err)
		bmcError := bmc.Error{Code: string(resp.StatusCode), Message: err.Error()}
		return false, &bmcError
	}
	logrus.Debug("StatusCode: ", resp.StatusCode)
	if resp.StatusCode != 200 {
		return false, bmc.NewError(*resp)
	}
	return true, nil
}

// DeleteSecurityList deletes the specified security list, but only if it's not associated with a subnet
func (c *Client) DeleteSecurityList(securityListID string) (bool, *bmc.Error) {
	securityListID = url.PathEscape(securityListID)
	resp, err := c.Request("DELETE", fmt.Sprintf("/securityLists/%s", securityListID), nil)
	if err != nil {
		logrus.Error(err)
		bmcError := bmc.Error{Code: string(resp.StatusCode), Message: err.Error()}
		return false, &bmcError
	}
	logrus.Debug("StatusCode: ", resp.StatusCode)
	if resp.StatusCode != 204 {
		return false, bmc.NewError(*resp)
	}
	return true, nil
}
<|endoftext|>


<commit_before>package imageboot

import (
	"strconv"
	"strings"
	"time"

	"github.com/GoogleCloudPlatform/guest-test-infra/imagetest"
)

// Name is the name of the test package. It must match the directory name.
var Name = "image_boot"

// TestSetup sets up the test workflow.
func TestSetup(t *imagetest.TestWorkflow) error {
	vm, err := t.CreateTestVM("vm")
	if err != nil {
		return err
	}
	if err := vm.Reboot(); err != nil {
		return err
	}
	vm.RunTests("TestGuestBoot|TestGuestReboot$")

	vm2, err := t.CreateTestVM("vm2")
	if err != nil {
		return err
	}
	vm2.RunTests("TestGuestRebootOnHost")

	if strings.Contains(t.Image, "debian-9") {
		// secure boot is not supported on Debian 9
		return nil
	}

	if strings.Contains(t.Image, "rocky-linux-8") {
		// secure boot is not supported on Rocky Linux
		return nil
	}

	vm3, err := t.CreateTestVM("vm3")
	if err != nil {
		return err
	}
	vm3.AddMetadata("start-time", strconv.Itoa(time.Now().Second()))
	vm3.EnableSecureBoot()
	vm3.RunTests("TestGuestSecureBoot|TestBootTime")
	return nil
}
<commit_msg>skip secure boot test on arm (#492)<commit_after>package imageboot

import (
	"strconv"
	"strings"
	"time"

	"github.com/GoogleCloudPlatform/guest-test-infra/imagetest"
)

// Name is the name of the test package. It must match the directory name.
var Name = "image_boot"

// TestSetup sets up the test workflow.
func TestSetup(t *imagetest.TestWorkflow) error {
	vm, err := t.CreateTestVM("vm")
	if err != nil {
		return err
	}
	if err := vm.Reboot(); err != nil {
		return err
	}
	vm.RunTests("TestGuestBoot|TestGuestReboot$")

	vm2, err := t.CreateTestVM("vm2")
	if err != nil {
		return err
	}
	vm2.RunTests("TestGuestRebootOnHost")

	if strings.Contains(t.Image, "debian-9") {
		// secure boot is not supported on Debian 9
		return nil
	}

	if strings.Contains(t.Image, "rocky-linux-8") {
		// secure boot is not supported on Rocky Linux
		return nil
	}

	if strings.Contains(t.Image, "arm64") {
		// secure boot is not supported on ARM images
		return nil
	}

	vm3, err := t.CreateTestVM("vm3")
	if err != nil {
		return err
	}
	vm3.AddMetadata("start-time", strconv.Itoa(time.Now().Second()))
	vm3.EnableSecureBoot()
	vm3.RunTests("TestGuestSecureBoot|TestBootTime")
	return nil
}
<|endoftext|>