<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"os"
	"strings"
	"time"

	"github.com/google/fleetspeak/fleetspeak/src/e2etesting/balancer/proxyproto"
)

var (
	serversFile        = flag.String("servers_file", "", "File with server hosts")
	serverFrontendAddr = flag.String("frontend_address", "", "Frontend address for clients to connect")
)

func copy(wc io.WriteCloser, r io.Reader) {
	defer wc.Close()
	io.Copy(wc, r)
}

func run() error {
	dat, err := ioutil.ReadFile(*serversFile)
	if err != nil {
		return fmt.Errorf("Failed to read serversFile: %v", err)
	}
	serverHosts := strings.Fields(string(dat))
	if len(serverHosts) == 0 {
		return fmt.Errorf("No server hosts were provided")
	}

	ln, err := net.Listen("tcp", *serverFrontendAddr)
	if err != nil {
		return fmt.Errorf("Failed to bind: %v", err)
	}
	log.Printf("Load balancer started on %v\n", *serverFrontendAddr)

	for {
		lbConn, err := ln.Accept()
		if err != nil {
			return fmt.Errorf("Failed to accept connection: %v", err)
		}
		var serverAddr string
		var serverConn net.Conn
		retriesLeft := 10
		for {
			serverAddr = serverHosts[rand.Int()%len(serverHosts)]
			serverConn, err = net.Dial("tcp", serverAddr)
			if err != nil {
				log.Printf("Failed to connect to server (%v): %v, retrying...\n", serverAddr, err)
				retriesLeft--
				if retriesLeft < 0 {
					return fmt.Errorf("Maximum number of retries exceeded - no active servers were found")
				}
				time.Sleep(time.Second * 2)
			} else {
				break
			}
		}
		log.Printf("Connection accepted, server: %v\n", serverAddr)
		err = proxyproto.WriteFirstProxyMessage(serverConn, lbConn.RemoteAddr().String(), serverAddr)
		if err != nil {
			return err
		}
		go copy(serverConn, lbConn)
		go copy(lbConn, serverConn)
	}
}

func main() {
	flag.Parse()
	err := run()
	if err != nil {
		fmt.Printf("Load balancer failed: %v\n", err)
		os.Exit(1)
	}
}
<commit_msg>In e2e testing load balancer, make usage of proxyproto optional. (#309)<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"os"
	"strings"
	"time"

	"github.com/google/fleetspeak/fleetspeak/src/e2etesting/balancer/proxyproto"
)

var (
	serversFile        = flag.String("servers_file", "", "File with server hosts")
	serverFrontendAddr = flag.String("frontend_address", "", "Frontend address for clients to connect")
	useProxyProto      = flag.Bool("use_proxy_proto", true, "Whether to forward client information using proxy proto")
)

func copy(wc io.WriteCloser, r io.Reader) {
	defer wc.Close()
	io.Copy(wc, r)
}

func run() error {
	dat, err := ioutil.ReadFile(*serversFile)
	if err != nil {
		return fmt.Errorf("Failed to read serversFile: %v", err)
	}
	serverHosts := strings.Fields(string(dat))
	if len(serverHosts) == 0 {
		return fmt.Errorf("No server hosts were provided")
	}

	ln, err := net.Listen("tcp", *serverFrontendAddr)
	if err != nil {
		return fmt.Errorf("Failed to bind: %v", err)
	}
	log.Printf("Load balancer started on %v\n", *serverFrontendAddr)

	for {
		lbConn, err := ln.Accept()
		if err != nil {
			return fmt.Errorf("Failed to accept connection: %v", err)
		}
		var serverAddr string
		var serverConn net.Conn
		retriesLeft := 10
		for {
			serverAddr = serverHosts[rand.Int()%len(serverHosts)]
			serverConn, err = net.Dial("tcp", serverAddr)
			if err != nil {
				log.Printf("Failed to connect to server (%v): %v, retrying...\n", serverAddr, err)
				retriesLeft--
				if retriesLeft < 0 {
					return fmt.Errorf("Maximum number of retries exceeded - no active servers were found")
				}
				time.Sleep(time.Second * 2)
			} else {
				break
			}
		}
		log.Printf("Connection accepted, server: %v\n", serverAddr)
		if *useProxyProto {
			err = proxyproto.WriteFirstProxyMessage(serverConn, lbConn.RemoteAddr().String(), serverAddr)
			if err != nil {
				return err
			}
		}
		go copy(serverConn, lbConn)
		go copy(lbConn, serverConn)
	}
}

func main() {
	flag.Parse()
	err := run()
	if err != nil {
		fmt.Printf("Load balancer failed: %v\n", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package respond

import (
	"sync"
	"net/http"
	"log"
	"fmt"
	"crypto/md5"
)

var (
	mutex     sync.RWMutex
	options   map[*http.Request]*Options
	responded map[*http.Request]bool
)

func init() {
	options   = make(map[*http.Request]*Options)
	responded = make(map[*http.Request]bool)
}

type Options struct {
	AllowMultiple bool

	OnError func(error)

	Encoder Encoder

	Before func(w http.ResponseWriter, r *http.Request, status int, data interface{}) (int, interface{})

	StatusFunc func(w http.ResponseWriter, r *http.Request, status int) interface{}
}

// Handler wraps an HTTP handler becoming the source of options for all
// containing With calls.
func (opts *Options) Handler(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mutex.Lock()
		options[r] = opts
		mutex.Unlock()

		defer func() {
			mutex.Lock()
			delete(options, r)
			delete(responded, r)
			mutex.Unlock()
		}()

		handler.ServeHTTP(w, r)
	})
}

func (opts *Options) log(err error) {
	if opts.OnError == nil {
		log.Println("respond: " + err.Error())
	}

	opts.OnError(err)
}

func about(r *http.Request) (opts *Options, multiple bool) {
	mutex.RLock()
	opts = options[r]
	multiple = responded[r]
	mutex.RUnlock()

	return opts, multiple
}

func with(w http.ResponseWriter, r *http.Request, status int, data interface{}, opts *Options, multiple bool) {
	if opts != nil && multiple && !opts.AllowMultiple {
		panic("respond: multiple responses")
	}

	encoder := JSON
	if opts != nil {
		if opts.Before != nil {
			status, data = opts.Before(w, r, status, data)
		}

		if opts.Encoder != nil {
			encoder = opts.Encoder
		}
	}

	out, err := encoder.Encode(data)
	if err != nil {
		opts.log(err)
		return
	}

	if status >= 200 && status < 300 {
		etag := fmt.Sprintf("\"%x\"", md5.Sum(out))

		w.Header().Set("Etag", etag)

		if match := r.Header.Get("if-none-match"); match == etag {
			w.WriteHeader(http.StatusNotModified)
			return
		}
	}

	w.Header().Set("Content-Type", encoder.ContentType())
	w.WriteHeader(status)
	w.Write(out)
}

func With(w http.ResponseWriter, r *http.Request, status int, data interface{}) {
	opts, multiple := about(r)

	with(w, r, status, data, opts, multiple)
}

func WithStatus(w http.ResponseWriter, r *http.Request, status int) {
	opts, multiple := about(r)

	var data interface{}
	if opts != nil && opts.StatusFunc != nil {
		data = opts.StatusFunc(w, r, status)
	} else {
		const (
			fieldStatus = "status"
			fieldCode   = "code"
		)
		data = map[string]interface{}{fieldStatus: http.StatusText(status), fieldCode: status}
	}

	with(w, r, status, data, opts, multiple)
}
<commit_msg>readers and custom headers<commit_after>package respond

import (
	"sync"
	"net/http"
	"log"
	"fmt"
	"crypto/md5"
	"io"
	"bytes"
)

var (
	mutex     sync.RWMutex
	options   map[*http.Request]*Options
	responded map[*http.Request]bool
)

func init() {
	options   = make(map[*http.Request]*Options)
	responded = make(map[*http.Request]bool)
}

type ReadHasher interface {
	io.Reader
	Hash() []byte
}

type rhImpl struct {
	io.Reader
	hash []byte
}

func (rh *rhImpl) Hash() []byte {
	return rh.hash
}

type Options struct {
	AllowMultiple bool

	OnError func(error)

	Encoder Encoder

	Before func(w http.ResponseWriter, r *http.Request, status int, data interface{}) (int, interface{})

	StatusFunc func(w http.ResponseWriter, r *http.Request, status int) interface{}
}

// Handler wraps an HTTP handler becoming the source of options for all
// containing With calls.
func (opts *Options) Handler(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mutex.Lock()
		options[r] = opts
		mutex.Unlock()

		defer func() {
			mutex.Lock()
			delete(options, r)
			delete(responded, r)
			mutex.Unlock()
		}()

		handler.ServeHTTP(w, r)
	})
}

func (opts *Options) log(err error) {
	if opts.OnError == nil {
		log.Println("respond: " + err.Error())
	}

	opts.OnError(err)
}

func about(r *http.Request) (opts *Options, multiple bool) {
	mutex.RLock()
	opts = options[r]
	multiple = responded[r]
	mutex.RUnlock()

	return opts, multiple
}

func with(w http.ResponseWriter, r *http.Request, status int, in io.Reader, headers []string, opts *Options, multiple bool) {
	if opts != nil && multiple && !opts.AllowMultiple {
		panic("respond: multiple responses")
	}

	if hasher, ok := in.(ReadHasher); ok {
		if status >= 200 && status < 300 {
			etag := fmt.Sprintf("\"%s\"", hasher.Hash())

			w.Header().Set("Etag", etag)

			if match := r.Header.Get("if-none-match"); match == etag {
				w.WriteHeader(http.StatusNotModified)
				return
			}
		}
	}

	for i := 0; i < len(headers)/2; i++ {
		w.Header().Set(headers[i*2], headers[i*2 + 1])
	}

	w.WriteHeader(status)
	io.Copy(w, in)
}

func With(w http.ResponseWriter, r *http.Request, status int, data interface{}, headers ...string) {
	opts, multiple := about(r)

	var out []byte
	if bs, ok := data.([]byte); ok {
		out = bs
	} else {
		encoder := JSON
		if opts != nil {
			if opts.Before != nil {
				status, data = opts.Before(w, r, status, data)
			}

			if opts.Encoder != nil {
				encoder = opts.Encoder
			}
		}

		bs, err := encoder.Encode(data)
		if err != nil {
			opts.log(err)
			return
		}

		headers = append(headers, "Content-Type", encoder.ContentType())
		out = bs
	}

	with(w, r, status, &rhImpl{bytes.NewBuffer(out), md5.New().Sum(out)}, headers, opts, multiple)
}

func WithReader(w http.ResponseWriter, r *http.Request, status int, reader io.Reader, headers ...string) {
	opts, multiple := about(r)
	with(w, r, status, reader, headers, opts, multiple)
}

func WithStatus(w http.ResponseWriter, r *http.Request, status int, headers ...string) {
	const (
		fieldStatus = "status"
		fieldCode   = "code"
	)
	With(w, r, status, map[string]interface{}{fieldStatus: http.StatusText(status), fieldCode: status}, headers...)
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2014 Jeremy Saenz
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

package hador

import (
	"bufio"
	"fmt"
	"net"
	"net/http"
)

// ResponseWriter is a wrapper around http.ResponseWriter that provides extra information about
// the response. It is recommended that middleware handlers use this construct to wrap a responsewriter
// if the functionality calls for it.
type ResponseWriter interface {
	http.ResponseWriter
	http.Flusher
	// Status returns the status code of the response or 0 if the response has not been written.
	Status() int
	// Written returns whether or not the ResponseWriter has been written.
	Written() bool
	// Size returns the size of the response body.
	Size() int
	// Before allows for a function to be called before the ResponseWriter has been written to. This is
	// useful for setting headers or any other operations that must happen before a response has been written.
	Before(func(ResponseWriter))
}

type beforeFunc func(ResponseWriter)

// NewResponseWriter creates a ResponseWriter that wraps an http.ResponseWriter
func NewResponseWriter(rw http.ResponseWriter) ResponseWriter {
	return &responseWriter{rw, 0, 0, nil}
}

type responseWriter struct {
	http.ResponseWriter
	status      int
	size        int
	beforeFuncs []beforeFunc
}

func (rw *responseWriter) reset(w http.ResponseWriter) {
	rw.ResponseWriter = w
	rw.status = 0
	rw.size = 0
	rw.beforeFuncs = nil
}

func (rw *responseWriter) WriteHeader(s int) {
	rw.callBefore()
	rw.ResponseWriter.WriteHeader(s)
	rw.status = s
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	if !rw.Written() {
		// The status will be StatusOK if WriteHeader has not been called yet
		rw.WriteHeader(http.StatusOK)
	}
	size, err := rw.ResponseWriter.Write(b)
	rw.size += size
	return size, err
}

func (rw *responseWriter) Status() int {
	return rw.status
}

func (rw *responseWriter) Size() int {
	return rw.size
}

func (rw *responseWriter) Written() bool {
	return rw.status != 0
}

func (rw *responseWriter) Before(before func(ResponseWriter)) {
	rw.beforeFuncs = append(rw.beforeFuncs, before)
}

func (rw *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hijacker, ok := rw.ResponseWriter.(http.Hijacker)
	if !ok {
		return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
	}
	return hijacker.Hijack()
}

func (rw *responseWriter) CloseNotify() <-chan bool {
	return rw.ResponseWriter.(http.CloseNotifier).CloseNotify()
}

func (rw *responseWriter) callBefore() {
	for i := len(rw.beforeFuncs) - 1; i >= 0; i-- {
		rw.beforeFuncs[i](rw)
	}
}

func (rw *responseWriter) Flush() {
	flusher, ok := rw.ResponseWriter.(http.Flusher)
	if ok {
		flusher.Flush()
	}
}
<commit_msg>Added WriteString method.<commit_after>/*
 * Copyright (c) 2014 Jeremy Saenz
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

package hador

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"net/http"
)

// ResponseWriter is a wrapper around http.ResponseWriter that provides extra information about
// the response. It is recommended that middleware handlers use this construct to wrap a responsewriter
// if the functionality calls for it.
type ResponseWriter interface {
	http.ResponseWriter
	http.Flusher
	// Status returns the status code of the response or 0 if the response has not been written.
	Status() int

	WriteString(string) (int, error)

	// Written returns whether or not the ResponseWriter has been written.
	Written() bool
	// Size returns the size of the response body.
	Size() int
	// Before allows for a function to be called before the ResponseWriter has been written to. This is
	// useful for setting headers or any other operations that must happen before a response has been written.
	Before(func(ResponseWriter))
}

type beforeFunc func(ResponseWriter)

// NewResponseWriter creates a ResponseWriter that wraps an http.ResponseWriter
func NewResponseWriter(rw http.ResponseWriter) ResponseWriter {
	return &responseWriter{rw, 0, 0, nil}
}

type responseWriter struct {
	http.ResponseWriter
	status      int
	size        int
	beforeFuncs []beforeFunc
}

func (rw *responseWriter) reset(w http.ResponseWriter) {
	rw.ResponseWriter = w
	rw.status = 0
	rw.size = 0
	rw.beforeFuncs = nil
}

func (rw *responseWriter) WriteHeader(s int) {
	rw.callBefore()
	rw.ResponseWriter.WriteHeader(s)
	rw.status = s
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	if !rw.Written() {
		// The status will be StatusOK if WriteHeader has not been called yet
		rw.WriteHeader(http.StatusOK)
	}
	size, err := rw.ResponseWriter.Write(b)
	rw.size += size
	return size, err
}

func (rw *responseWriter) WriteString(s string) (int, error) {
	if !rw.Written() {
		// The status will be StatusOK if WriteHeader has not been called yet
		rw.WriteHeader(http.StatusOK)
	}
	size, err := io.WriteString(rw.ResponseWriter, s)
	rw.size += size
	return size, err
}

func (rw *responseWriter) Status() int {
	return rw.status
}

func (rw *responseWriter) Size() int {
	return rw.size
}

func (rw *responseWriter) Written() bool {
	return rw.status != 0
}

func (rw *responseWriter) Before(before func(ResponseWriter)) {
	rw.beforeFuncs = append(rw.beforeFuncs, before)
}

func (rw *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hijacker, ok := rw.ResponseWriter.(http.Hijacker)
	if !ok {
		return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
	}
	return hijacker.Hijack()
}

func (rw *responseWriter) CloseNotify() <-chan bool {
	return rw.ResponseWriter.(http.CloseNotifier).CloseNotify()
}

func (rw *responseWriter) callBefore() {
	for i := len(rw.beforeFuncs) - 1; i >= 0; i-- {
		rw.beforeFuncs[i](rw)
	}
}

func (rw *responseWriter) Flush() {
	flusher, ok := rw.ResponseWriter.(http.Flusher)
	if ok {
		flusher.Flush()
	}
}
<|endoftext|>


<commit_before>package rockredis

import (
	"encoding/binary"
	"errors"
	"runtime"
	"sync"
	"time"

	"github.com/absolute8511/ZanRedisDB/common"
	"github.com/absolute8511/gorocksdb"
)

var (
	errExpMetaKey = errors.New("invalid expire meta key")
	errExpTimeKey = errors.New("invalid expire time key")
)

const (
	logTimeFormatStr = "2006-01-02 15:04:05"
)

var errExpType = errors.New("invalid expire type")

/*
the coded format of expire time key:
bytes:  -0-|-1-2-3-4-5-6-7-8-|----9---|-10-11--------x-|
data :  103|       when      |dataType|       key      |
*/
func expEncodeTimeKey(dataType byte, key []byte, when int64) []byte {
	buf := make([]byte, len(key)+1+8+1)

	pos := 0
	buf[pos] = ExpTimeType
	pos++

	binary.BigEndian.PutUint64(buf[pos:], uint64(when))
	pos += 8

	buf[pos] = dataType
	pos++

	copy(buf[pos:], key)

	return buf
}

/*
the coded format of expire meta key:
bytes:  -0-|----1-----|-2-3----------x--|
data :  102| dataType |       key       |
*/
func expEncodeMetaKey(dataType byte, key []byte) []byte {
	buf := make([]byte, len(key)+2)

	pos := 0

	buf[pos] = ExpMetaType
	pos++
	buf[pos] = dataType
	pos++

	copy(buf[pos:], key)

	return buf
}

//decode the expire 'meta key', the return values are: dataType, key, error
func expDecodeMetaKey(mk []byte) (byte, []byte, error) {
	pos := 0

	if pos+2 > len(mk) || mk[pos] != ExpMetaType {
		return 0, nil, errExpMetaKey
	}

	return mk[pos+1], mk[pos+2:], nil
}

//decode the expire 'time key', the return values are: dataType, key, whenToExpire, error
func expDecodeTimeKey(tk []byte) (byte, []byte, int64, error) {
	pos := 0
	if pos+10 > len(tk) || tk[pos] != ExpTimeType {
		return 0, nil, 0, errExpTimeKey
	}

	return tk[pos+9], tk[pos+10:], int64(binary.BigEndian.Uint64(tk[pos+1:])), nil
}

type expiredMeta struct {
	timeKey []byte
	metaKey []byte
	UTC     int64
}

type expiredMetaBuffer interface {
	Write(*expiredMeta) error
}

type expiration interface {
	rawExpireAt(byte, []byte, int64, *gorocksdb.WriteBatch) error
	expireAt(byte, []byte, int64) error
	ttl(byte, []byte) (int64, error)
	delExpire(byte, []byte, *gorocksdb.WriteBatch) error
	check(common.ExpiredDataBuffer, chan struct{}) error
	Start()
	Stop()
}

func (db *RockDB) expire(dataType byte, key []byte, duration int64) error {
	return db.expiration.expireAt(dataType, key, time.Now().Unix()+duration)
}

func (db *RockDB) KVTtl(key []byte) (t int64, err error) {
	return db.ttl(KVType, key)
}

func (db *RockDB) HashTtl(key []byte) (t int64, err error) {
	return db.ttl(HashType, key)
}

func (db *RockDB) ListTtl(key []byte) (t int64, err error) {
	return db.ttl(ListType, key)
}

func (db *RockDB) SetTtl(key []byte) (t int64, err error) {
	return db.ttl(SetType, key)
}

func (db *RockDB) ZSetTtl(key []byte) (t int64, err error) {
	return db.ttl(ZSetType, key)
}

type TTLChecker struct {
	sync.Mutex

	db       *RockDB
	watching int32
	wg       sync.WaitGroup

	//next check time
	nc int64
}

func newTTLChecker(db *RockDB) *TTLChecker {
	c := &TTLChecker{
		db: db,
		nc: time.Now().Unix(),
	}
	return c
}

func (c *TTLChecker) setNextCheckTime(when int64, force bool) {
	c.Lock()
	if force {
		c.nc = when
	} else if c.nc > when {
		c.nc = when
	}
	c.Unlock()
}

func (c *TTLChecker) check(expiredBuf expiredMetaBuffer, stop chan struct{}) (err error) {
	defer func() {
		if e := recover(); e != nil {
			buf := make([]byte, 4096)
			n := runtime.Stack(buf, false)
			buf = buf[0:n]
			dbLog.Errorf("check ttl panic: %s:%v", buf, e)
		}
	}()

	now := time.Now().Unix()

	c.Lock()
	nc := c.nc
	c.Unlock()

	if now < nc {
		return nil
	}

	nc = now + 3600

	minKey := expEncodeTimeKey(NoneType, nil, 0)
	maxKey := expEncodeTimeKey(maxDataType, nil, nc)

	var eCount int64
	var scanned int64
	checkStart := time.Now()

	it, err := NewDBRangeLimitIterator(c.db.eng, minKey, maxKey,
		common.RangeROpen, 0, -1, false)
	if err != nil {
		c.setNextCheckTime(now, false)
		return err
	} else if it == nil {
		c.setNextCheckTime(nc, true)
		return nil
	}
	defer it.Close()

	for ; it.Valid(); it.Next() {
		if scanned%100 == 0 {
			select {
			case <-stop:
				nc = now
				break
			default:
			}
		}
		tk := it.Key()
		mk := it.Value()

		if tk == nil {
			continue
		}

		dt, k, nt, dErr := expDecodeTimeKey(tk)
		if dErr != nil {
			continue
		}

		scanned += 1
		if scanned == 1 {
			//log the first scanned key
			dbLog.Infof("ttl check start at key:[%s] of type:%s whose expire time is: %s", string(k),
				TypeName[dt], time.Unix(nt, 0).Format(logTimeFormatStr))
		}

		if nt > now {
			//the next ttl check time is nt!
			nc = nt
			dbLog.Infof("ttl check end at key:[%s] of type:%s whose expire time is: %s", string(k),
				TypeName[dt], time.Unix(nt, 0).Format(logTimeFormatStr))
			break
		}

		eCount += 1

		err = expiredBuf.Write(&expiredMeta{timeKey: tk, metaKey: mk, UTC: nt})
		if err != nil {
			nc = now
			break
		}
	}

	c.setNextCheckTime(nc, true)

	checkCost := time.Since(checkStart).Nanoseconds() / 1000
	dbLog.Infof("[%d/%d] keys have expired during ttl checking, cost:%d us", eCount, scanned, checkCost)

	return err
}
<commit_msg>adjust log level<commit_after>package rockredis

import (
	"encoding/binary"
	"errors"
	"runtime"
	"sync"
	"time"

	"github.com/absolute8511/ZanRedisDB/common"
	"github.com/absolute8511/gorocksdb"
)

var (
	errExpMetaKey = errors.New("invalid expire meta key")
	errExpTimeKey = errors.New("invalid expire time key")
)

const (
	logTimeFormatStr = "2006-01-02 15:04:05"
)

var errExpType = errors.New("invalid expire type")

/*
the coded format of expire time key:
bytes:  -0-|-1-2-3-4-5-6-7-8-|----9---|-10-11--------x-|
data :  103|       when      |dataType|       key      |
*/
func expEncodeTimeKey(dataType byte, key []byte, when int64) []byte {
	buf := make([]byte, len(key)+1+8+1)

	pos := 0
	buf[pos] = ExpTimeType
	pos++

	binary.BigEndian.PutUint64(buf[pos:], uint64(when))
	pos += 8

	buf[pos] = dataType
	pos++

	copy(buf[pos:], key)

	return buf
}

/*
the coded format of expire meta key:
bytes:  -0-|----1-----|-2-3----------x--|
data :  102| dataType |       key       |
*/
func expEncodeMetaKey(dataType byte, key []byte) []byte {
	buf := make([]byte, len(key)+2)

	pos := 0

	buf[pos] = ExpMetaType
	pos++
	buf[pos] = dataType
	pos++

	copy(buf[pos:], key)

	return buf
}

//decode the expire 'meta key', the return values are: dataType, key, error
func expDecodeMetaKey(mk []byte) (byte, []byte, error) {
	pos := 0

	if pos+2 > len(mk) || mk[pos] != ExpMetaType {
		return 0, nil, errExpMetaKey
	}

	return mk[pos+1], mk[pos+2:], nil
}

//decode the expire 'time key', the return values are: dataType, key, whenToExpire, error
func expDecodeTimeKey(tk []byte) (byte, []byte, int64, error) {
	pos := 0
	if pos+10 > len(tk) || tk[pos] != ExpTimeType {
		return 0, nil, 0, errExpTimeKey
	}

	return tk[pos+9], tk[pos+10:], int64(binary.BigEndian.Uint64(tk[pos+1:])), nil
}

type expiredMeta struct {
	timeKey []byte
	metaKey []byte
	UTC     int64
}

type expiredMetaBuffer interface {
	Write(*expiredMeta) error
}

type expiration interface {
	rawExpireAt(byte, []byte, int64, *gorocksdb.WriteBatch) error
	expireAt(byte, []byte, int64) error
	ttl(byte, []byte) (int64, error)
	delExpire(byte, []byte, *gorocksdb.WriteBatch) error
	check(common.ExpiredDataBuffer, chan struct{}) error
	Start()
	Stop()
}

func (db *RockDB) expire(dataType byte, key []byte, duration int64) error {
	return db.expiration.expireAt(dataType, key, time.Now().Unix()+duration)
}

func (db *RockDB) KVTtl(key []byte) (t int64, err error) {
	return db.ttl(KVType, key)
}

func (db *RockDB) HashTtl(key []byte) (t int64, err error) {
	return db.ttl(HashType, key)
}

func (db *RockDB) ListTtl(key []byte) (t int64, err error) {
	return db.ttl(ListType, key)
}

func (db *RockDB) SetTtl(key []byte) (t int64, err error) {
	return db.ttl(SetType, key)
}

func (db *RockDB) ZSetTtl(key []byte) (t int64, err error) {
	return db.ttl(ZSetType, key)
}

type TTLChecker struct {
	sync.Mutex

	db       *RockDB
	watching int32
	wg       sync.WaitGroup

	//next check time
	nc int64
}

func newTTLChecker(db *RockDB) *TTLChecker {
	c := &TTLChecker{
		db: db,
		nc: time.Now().Unix(),
	}
	return c
}

func (c *TTLChecker) setNextCheckTime(when int64, force bool) {
	c.Lock()
	if force {
		c.nc = when
	} else if c.nc > when {
		c.nc = when
	}
	c.Unlock()
}

func (c *TTLChecker) check(expiredBuf expiredMetaBuffer, stop chan struct{}) (err error) {
	defer func() {
		if e := recover(); e != nil {
			buf := make([]byte, 4096)
			n := runtime.Stack(buf, false)
			buf = buf[0:n]
			dbLog.Errorf("check ttl panic: %s:%v", buf, e)
		}
	}()

	now := time.Now().Unix()

	c.Lock()
	nc := c.nc
	c.Unlock()

	if now < nc {
		return nil
	}

	nc = now + 3600

	minKey := expEncodeTimeKey(NoneType, nil, 0)
	maxKey := expEncodeTimeKey(maxDataType, nil, nc)

	var eCount int64
	var scanned int64
	checkStart := time.Now()

	it, err := NewDBRangeLimitIterator(c.db.eng, minKey, maxKey,
		common.RangeROpen, 0, -1, false)
	if err != nil {
		c.setNextCheckTime(now, false)
		return err
	} else if it == nil {
		c.setNextCheckTime(nc, true)
		return nil
	}
	defer it.Close()

	for ; it.Valid(); it.Next() {
		if scanned%100 == 0 {
			select {
			case <-stop:
				nc = now
				break
			default:
			}
		}
		tk := it.Key()
		mk := it.Value()

		if tk == nil {
			continue
		}

		dt, k, nt, dErr := expDecodeTimeKey(tk)
		if dErr != nil {
			continue
		}

		scanned += 1
		if scanned == 1 {
			//log the first scanned key
			dbLog.Debugf("ttl check start at key:[%s] of type:%s whose expire time is: %s", string(k),
				TypeName[dt], time.Unix(nt, 0).Format(logTimeFormatStr))
		}

		if nt > now {
			//the next ttl check time is nt!
			nc = nt
			dbLog.Debugf("ttl check end at key:[%s] of type:%s whose expire time is: %s", string(k),
				TypeName[dt], time.Unix(nt, 0).Format(logTimeFormatStr))
			break
		}

		eCount += 1

		err = expiredBuf.Write(&expiredMeta{timeKey: tk, metaKey: mk, UTC: nt})
		if err != nil {
			nc = now
			break
		}
	}

	c.setNextCheckTime(nc, true)

	checkCost := time.Since(checkStart).Nanoseconds() / 1000
	if dbLog.Level() >= common.LOG_DEBUG || eCount > 100 || checkCost >= time.Second.Nanoseconds()/1000 {
		dbLog.Infof("[%d/%d] keys have expired during ttl checking, cost:%d us", eCount, scanned, checkCost)
	}

	return err
}
<|endoftext|>


<commit_before>// Package routing is a wrapper on naoina/denco router.
// It uses request.Form for params instead of a separate Params
// argument. So, it requires a more memory and it is a bit slower.
// However, the downsides are an acceptable trade off for compatability
// with the standard library.
package routing

import (
	"net/http"

	"github.com/naoina/denco"
)

// Router represents a multiplexer for HTTP request.
type Router struct {
	data *denco.Router
}

// info stores information about HTTP request's handler, its
// pattern and methods.
type info struct {
	handler http.HandlerFunc // HTTP request handler.
	methods []string         // A list of allowed HTTP methods (e.g. "GET" or "POST").
	pattern string           // Pattern is a routing path for handler.
}

// New allocates and returns a new multiplexer.
func New() *Router {
	return &Router{}
}

// HasMethod checks whether specific method request is allowed.
func (t *info) HasMethod(name string) bool {
	// We are iterating through a slice of method strings
	// rather than using a map as there are only a few possible values.
	// So, hash function will require more time than a simple loop.
	for _, v := range t.methods {
		if v == name {
			return true
		}
	}
	return false
}

// ServeHTTP is used to implement http.Handler interface.
// It dispatches the request to the handler whose pattern
// most closely matches the request URL.
func (t *Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
}

// Handle registers the handler for the given pattern.
// If a handler already exists for pattern, it is being overridden.
func (t *Router) Handle(method, pattern string, handler http.Handler) {
}

// HandleFunc registers the handler function for the given pattern.
func (t *Router) HandleFunc(method, pattern string, handler http.HandlerFunc) {
}

// Handler returns the handler to use for the given request, consulting r.Method
// and r.URL.Path. It always returns a non-nil handler. If there is no registered handler
// that applies to the request, Handler returns a “page not found” handler and empty pattern.
// If there is a registered handler but requested method is not allowed,
// "method not allowed" and a pattern are returned.
func (t *Router) Handler(r *http.Request) (handler http.Handler, pattern string) {
	// Make sure we have a handler for this request.
	obj, params, found := t.data.Lookup(r.URL.Path)
	if !found {
		return http.HandlerFunc(NotFound), ""
	}

	// Check whether requested method is allowed.
	data := obj.(info)
	if !data.HasMethod(r.Method) {
		return http.HandlerFunc(MethodNotAllowed), data.pattern
	}

	// Add parameters of request to request.Form and return a handler.
	for _, param := range params {
		r.Form.Add(param.Name, param.Value)
	}
	return data.handler, data.pattern
}

// MethodNotAllowed replies to the request with an HTTP 405 method not allowed
// error. If you want to use your own MethodNotAllowed handler, please override
// this variable.
var MethodNotAllowed = func(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "405 method not allowed", http.StatusMethodNotAllowed)
}

// NotFound replies to the request with an HTTP 404 not found error.
// NotFound is called when unknown HTTP method or a handler not found.
// If you want to use the your own NotFound handler, please overwrite this variable.
var NotFound = func(w http.ResponseWriter, r *http.Request) {
	http.NotFound(w, r)
}
<commit_msg>Implemented router using denco as a base<commit_after>// Package routing is a wrapper on naoina/denco router.
// It uses request.Form for params instead of a separate Params
// argument. So, it requires a more memory and it is a bit slower.
// However, the downsides are an acceptable trade off for compatability
// with the standard library.
//
// A sample of its usage is below:
//
//	package main
//
//	import (
//		"log"
//		"net/http"
//
//		"github.com/anonx/sunplate/routing"
//	)
//
//	func main() {
//		r := routing.New()
//		err := r.Handle(routing.Routes{
//			r.Get("/profiles/:username", ShowUserHandleFunc),
//			r.Delete("/profiles/:username", DeleteUserHandleFunc),
//		}).Build()
//		if err != nil {
//			panic(err)
//		}
//		log.Fatal(http.ListenAndServe(":8080", r))
//	}
package routing

import (
	"net/http"

	"github.com/naoina/denco"
)

// Router represents a multiplexer for HTTP requests.
type Router struct {
	data    *denco.Router  // data stores denco router.
	indexes map[string]int // indexes are used to check whether a record exists.
	records []denco.Record // records is a list of handlers expected by denco router.
}

// Routes is an alias of []Route.
type Routes []Route

// Route is used to store information about HTTP request's handler
// including a list of allowed methods and pattern.
type Route struct {
	handler *http.HandlerFunc // HTTP request handler function.
	methods []string          // A list of allowed HTTP methods (e.g. "GET" or "POST").
	pattern string            // Pattern is a routing path for handler.
}

// NewRouter allocates and returns a new multiplexer.
func NewRouter() *Router {
	return &Router{
		indexes: map[string]int{},
	}
}

// HasMethod checks whether specific method request is allowed.
func (t *Route) HasMethod(name string) bool {
	// We are iterating through a slice of method strings
	// rather than using a map as there are only a few possible values.
	// So, hash function will require more time than a simple loop.
	for _, v := range t.methods {
		if v == name {
			return true
		}
	}
	return false
}

// ServeHTTP is used to implement http.Handler interface.
// It dispatches the request to the handler whose pattern
// most closely matches the request URL.
func (t *Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h, _ := t.Handler(r)
	h.ServeHTTP(w, r)
}

// Handle registers handlers for given patterns.
// If a handler already exists for pattern, it will be overridden.
// If it exists but with another method, a new method will be added.
func (t *Router) Handle(routes Routes) {
	for _, route := range routes {
		// Check whether we have already had such route.
		index, ok := t.indexes[route.pattern]
		if ok {
			// If we haven't, add it.
			t.records = append(t.records, denco.NewRecord(route.pattern, route))
			continue
		}

		// Check whether existing route has the same handler and
		// we are just trying to add a new method.
		r := t.records[index].Value.(Route)
		if r.handler != route.handler {
			// If we aren't, override an old route.
			t.records[index] = denco.NewRecord(route.pattern, route)
			continue
		}

		// Otherwise, add all methods that haven't added yet.
		for _, v := range r.methods {
			if !route.HasMethod(v) {
				route.methods = append(route.methods, v)
			}
		}
	}
}

// Build compiles registered routes. Routes that are added after building will not
// be handled. A new call to build will be required.
func (t *Router) Build() error {
	router := denco.New()
	err := router.Build(t.records)
	if err != nil {
		return err
	}
	return nil
}

// Route allocates and returns a Route struct.
func (t *Router) Route(method, pattern string, handler http.HandlerFunc) Route {
	return Route{
		handler: &handler,
		methods: []string{method},
		pattern: pattern,
	}
}

// Handler returns the handler to use for the given request, consulting r.Method
// and r.URL.Path. It always returns a non-nil handler. If there is no registered handler
// that applies to the request, Handler returns a “page not found” handler and empty pattern.
// If there is a registered handler but requested method is not allowed,
// "method not allowed" and a pattern are returned.
func (t *Router) Handler(r *http.Request) (handler http.Handler, pattern string) {
	// Make sure we have a handler for this request.
	obj, params, found := t.data.Lookup(r.URL.Path)
	if !found {
		return http.HandlerFunc(NotFound), ""
	}

	// Check whether requested method is allowed.
	route := obj.(Route)
	if !route.HasMethod(r.Method) {
		return http.HandlerFunc(MethodNotAllowed), route.pattern
	}

	// Add parameters of request to request.Form and return a handler.
	for _, param := range params {
		r.Form.Add(param.Name, param.Value)
	}
	return route.handler, route.pattern
}

// MethodNotAllowed replies to the request with an HTTP 405 method not allowed
// error. If you want to use your own MethodNotAllowed handler, please override
// this variable.
var MethodNotAllowed = func(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "405 method not allowed", http.StatusMethodNotAllowed)
}

// NotFound replies to the request with an HTTP 404 not found error.
// NotFound is called when unknown HTTP method or a handler not found.
// If you want to use the your own NotFound handler, please overwrite this variable.
var NotFound = func(w http.ResponseWriter, r *http.Request) {
	http.NotFound(w, r)
}
<|endoftext|>


<commit_before>package runtime

import (
	"fmt"
	"github.com/golang/glog"
	"io"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
	"time"
)

func SavePidFile(args ...string) (string, error) {
	cmd := filepath.Base(os.Args[0])
	pidFile, err := os.Create(fmt.Sprintf("%s-%s.pid", cmd, strings.Join(args, "-")))
	if err != nil {
		return "", err
	}

	defer pidFile.Close()
	fmt.Fprintf(pidFile, "%d", os.Getpid())
	return pidFile.Name(), nil
}

type ShutdownSequence []io.Closer

func ShutdownHook(h func() error) closeWrapper {
	return closeWrapper{run: h}
}

type closeWrapper struct {
	run func() error
}

func (w closeWrapper) Close() error {
	return w.run()
}

// Implements io.Closer
func (s ShutdownSequence) Close() (err error) {
	for _, cl := range s {
		if err1 := cl.Close(); err == nil && err1 != nil {
			err = err1
		}
	}
	return
}

func exitf(pattern string, args ...interface{}) {
	if !strings.HasSuffix(pattern, "\n") {
		pattern = pattern + "\n"
	}
	fmt.Fprintf(os.Stderr, pattern, args...)
	os.Exit(1)
}

func HandleSignals(shutdownc <-chan io.Closer) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGHUP)
	signal.Notify(c, syscall.SIGINT)
	for {
		sig := <-c
		sysSig, ok := sig.(syscall.Signal)
		if !ok {
			glog.Fatal("Not a unix signal")
		}
		switch sysSig {
		case syscall.SIGHUP:
		case syscall.SIGINT:
			glog.Warningln("Got SIGTERM: shutting down")
			donec := make(chan bool)
			go func() {
				cl := <-shutdownc
				if err := cl.Close(); err != nil {
					exitf("Error shutting down: %v", err)
				}
				donec <- true
			}()
			select {
			case <-donec:
				glog.Infoln("Shut down completed.")
				os.Exit(0)
			case <-time.After(5 * time.Second):
				exitf("Timeout shutting down. Exiting uncleanly.")
			}
		default:
			glog.Fatal("Received another signal, should not happen.")
		}
	}
}

// Runs the http server.  This server offers more control than the standard go's default http server
// in that when a 'true' is sent to the stop channel, the listener is closed to force a clean shutdown.
func RunServer(server *http.Server, stop chan bool) (stopped chan bool) {
	protocol := "tcp"
	if match, _ := regexp.MatchString("[a-zA-Z0-9\\.]*:[0-9]{4,}", server.Addr); !match {
		protocol = "unix"
	}

	listener, err := net.Listen(protocol, server.Addr)
	if err != nil {
		panic(err)
	}
	stopped = make(chan bool)

	glog.Infoln("Starting", protocol, "listener at", server.Addr)

	// This will be set to true if a shutdown signal is received. This allows us to detect
	// if the server stop is intentional or due to some error.
	fromSignal := false

	// The main goroutine where the server listens on the network connection
	go func(fromSignal *bool) {
		// Serve will block until an error (e.g. from shutdown, closed connection) occurs.
		err := server.Serve(listener)
		if !*fromSignal {
			glog.Warningln("Warning: server stops due to error", err)
		}
		stopped <- true
	}(&fromSignal)

	// Another goroutine that listens for signal to close the network connection
	// on shutdown.  This will cause the server.Serve() to return.
	go func(fromSignal *bool) {
		select {
		case <-stop:
			listener.Close()
			*fromSignal = true // Intentially stopped from signal
			return
		}
	}(&fromSignal)
	return
}
<commit_msg>Adjust file permissions for unix domain sockets<commit_after>package runtime

import (
	"fmt"
	"github.com/golang/glog"
	"io"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
	"time"
)

func SavePidFile(args ...string) (string, error) {
	cmd := filepath.Base(os.Args[0])
	pidFile, err := os.Create(fmt.Sprintf("%s-%s.pid", cmd, strings.Join(args, "-")))
	if err != nil {
		return "", err
	}

	defer pidFile.Close()
	fmt.Fprintf(pidFile, "%d", os.Getpid())
	return pidFile.Name(), nil
}

type ShutdownSequence []io.Closer

func ShutdownHook(h func() error) closeWrapper {
	return closeWrapper{run: h}
}

type closeWrapper struct {
	run func() error
}

func (w closeWrapper) Close() error {
	return w.run()
}

// Implements io.Closer
func (s ShutdownSequence) Close() (err error) {
	for _, cl := range s {
		if err1 := cl.Close(); err == nil && err1 != nil {
			err = err1
		}
	}
	return
}

func exitf(pattern string, args ...interface{}) {
	if !strings.HasSuffix(pattern, "\n") {
		pattern = pattern + "\n"
	}
	fmt.Fprintf(os.Stderr, pattern, args...)
	os.Exit(1)
}

func HandleSignals(shutdownc <-chan io.Closer) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGHUP)
	signal.Notify(c, syscall.SIGINT)
	for {
		sig := <-c
		sysSig, ok := sig.(syscall.Signal)
		if !ok {
			glog.Fatal("Not a unix signal")
		}
		switch sysSig {
		case syscall.SIGHUP:
		case syscall.SIGINT:
			glog.Warningln("Got SIGTERM: shutting down")
			donec := make(chan bool)
			go func() {
				cl := <-shutdownc
				if err := cl.Close(); err != nil {
					exitf("Error shutting down: %v", err)
				}
				donec <- true
			}()
			select {
			case <-donec:
				glog.Infoln("Shut down completed.")
				os.Exit(0)
			case <-time.After(5 * time.Second):
				exitf("Timeout shutting down. Exiting uncleanly.")
			}
		default:
			glog.Fatal("Received another signal, should not happen.")
		}
	}
}

func updateDomainSocketPermissions(filename string) (err error) {
	_, err = os.Lstat(filename)
	if err != nil {
		return
	}
	return os.Chmod(filename, 0777)
}

// Runs the http server.  This server offers more control than the standard go's default http server
// in that when a 'true' is sent to the stop channel, the listener is closed to force a clean shutdown.
func RunServer(server *http.Server, stop chan bool) (stopped chan bool) {
	protocol := "tcp"
	if match, _ := regexp.MatchString("[a-zA-Z0-9\\.]*:[0-9]{4,}", server.Addr); !match {
		protocol = "unix"
	}

	listener, err := net.Listen(protocol, server.Addr)
	if err != nil {
		panic(err)
	}
	stopped = make(chan bool)

	glog.Infoln("Starting", protocol, "listener at", server.Addr)

	if protocol == "unix" {
		updateDomainSocketPermissions(server.Addr)
	}

	// This will be set to true if a shutdown signal is received. This allows us to detect
	// if the server stop is intentional or due to some error.
	fromSignal := false

	// The main goroutine where the server listens on the network connection
	go func(fromSignal *bool) {
		// Serve will block until an error (e.g. from shutdown, closed connection) occurs.
		err := server.Serve(listener)
		if !*fromSignal {
			glog.Warningln("Warning: server stops due to error", err)
		}
		stopped <- true
	}(&fromSignal)

	// Another goroutine that listens for signal to close the network connection
	// on shutdown.  This will cause the server.Serve() to return.
	go func(fromSignal *bool) {
		select {
		case <-stop:
			listener.Close()
			*fromSignal = true // Intentially stopped from signal
			return
		}
	}(&fromSignal)
	return
}
<|endoftext|>


<commit_before>// Package slack uses Norberto Lopes' slack library to implement the bot.Connector
// interface.
package slack

import (
	"fmt"
	"log"
	"os"
	"sync"

	"github.com/lnxjedi/gopherbot/bot"
	"github.com/nlopes/slack"
)

type botDefinition struct {
	Name, ID string // e.g. 'mygit', 'BAKDBISDO'
}

type config struct {
	SlackToken      string          // the 'bot token for connecting to Slack
	MaxMessageSplit int             // the maximum # of ~4000 byte messages to split a large message into
	BotRoster       []botDefinition // roster mapping BXXX IDs to a defined username
}

var lock sync.Mutex // package var lock
var started bool    // set when connector is started

func init() {
	bot.RegisterConnector("slack", Initialize)
}

// Initialize starts the connection, sets up and returns the connector object
func Initialize(robot bot.Handler, l *log.Logger) bot.Connector {
	lock.Lock()
	if started {
		lock.Unlock()
		return nil
	}
	started = true
	lock.Unlock()

	var c config
	var tok string

	err := robot.GetProtocolConfig(&c)
	if err != nil {
		robot.Log(bot.Fatal, fmt.Errorf("Unable to retrieve protocol configuration: %v", err))
	}

	if c.MaxMessageSplit == 0 {
		c.MaxMessageSplit = 1
	}

	if len(c.SlackToken) == 0 {
		tok = os.Getenv("SLACK_TOKEN")
		if len(tok) == 0 {
			robot.Log(bot.Fatal, "No slack token found in config or env var 'SLACK_TOKEN'")
		}
		os.Unsetenv("SLACK_TOKEN")
		robot.Log(bot.Debug, "Using SLACK_TOKEN environment variable")
	} else {
		tok = c.SlackToken
	}

	api := slack.New(tok)
	// This spits out a lot of extra stuff, so we only enable it when tracing
	if robot.GetLogLevel() == bot.Trace {
		api.SetDebug(true)
	}
	slack.SetLogger(l)

	sc := &slackConnector{
		api:             api,
		conn:            api.NewRTM(),
		maxMessageSplit: c.MaxMessageSplit,
		name:            "slack",
	}
	sc.botUserID = make(map[string]string)
	sc.botIDUser = make(map[string]string)
	for _, b := range c.BotRoster {
		if len(b.ID) == 0 || len(b.Name) == 0 {
			robot.Log(bot.Error, "Zero-length Name or ID in BotRoster, skipping")
			continue
		}
		sc.botUserID[b.Name] = b.ID
		sc.botIDUser[b.ID] = b.Name
	}
	go sc.conn.ManageConnection()

	sc.Handler = robot

Loop:
	for {
		select {
		case msg := <-sc.conn.IncomingEvents:

			switch ev := msg.Data.(type) {

			case *slack.ConnectedEvent:
				sc.Log(bot.Debug, fmt.Sprintf("Infos: %T %v\n", ev, *ev.Info.User))
				sc.Log(bot.Debug, "Connection counter:", ev.ConnectionCount)
				sc.botName = ev.Info.User.Name
				sc.SetName(sc.botName)
				sc.Log(bot.Info, "Slack setting bot name to", sc.botName)
				sc.botID = ev.Info.User.ID
				sc.Log(bot.Trace, "Set bot ID to", sc.botID)
				sc.teamID = ev.Info.Team.ID
				sc.Log(bot.Info, "Set team ID to", sc.teamID)
				break Loop

			case *slack.InvalidAuthEvent:
				sc.Log(bot.Fatal, "Invalid credentials")
			}
		}
	}

	sc.updateChannelMaps("")
	sc.updateUserList("")
	sc.botFullName, _ = sc.GetProtocolUserAttribute(sc.botName, "realname")
	sc.SetFullName(sc.botFullName)
	sc.Log(bot.Debug, "Set bot full name to", sc.botFullName)
	go sc.startSendLoop()

	return bot.Connector(sc)
}

func (sc *slackConnector) Run(stop <-chan struct{}) {
	sc.Lock()
	// This should never happen, just a bit of defensive coding
	if sc.running {
		sc.Unlock()
		return
	}
	sc.running = true
	sc.Unlock()
loop:
	for {
		select {
		case <-stop:
			sc.Log(bot.Debug, "Received stop in connector")
			break loop
		case msg := <-sc.conn.IncomingEvents:
			sc.Log(bot.Trace, fmt.Sprintf("Event Received (msg, data, type): %v; %v; %T", msg, msg.Data, msg.Data))
			switch ev := msg.Data.(type) {
			case *slack.HelloEvent:
				// Ignore hello
			case *slack.ChannelArchiveEvent, *slack.ChannelCreatedEvent, *slack.ChannelDeletedEvent, *slack.ChannelRenameEvent:
				sc.updateChannelMaps("")

			case *slack.MessageEvent:
				// Message processing is done concurrently
				go sc.processMessage(ev)

			case *slack.PresenceChangeEvent:
				sc.Log(bot.Debug, fmt.Sprintf("Presence Change: %v", ev))

			case *slack.LatencyReport:
				sc.Log(bot.Debug, fmt.Sprintf("Current latency: %v", ev.Value))

			case *slack.RTMError:
				sc.Log(bot.Debug, fmt.Sprintf("Error: %s\n", ev.Error()))

			default:

				// Ignore other events..
				// robot.Debug(fmt.Sprintf("Unexpected: %v\n", msg.Data)
			}
		}
	}
}
<commit_msg>Handle more channel updates<commit_after>// Package slack uses Norberto Lopes' slack library to implement the bot.Connector
// interface.
package slack

import (
	"fmt"
	"log"
	"os"
	"sync"

	"github.com/lnxjedi/gopherbot/bot"
	"github.com/nlopes/slack"
)

type botDefinition struct {
	Name, ID string // e.g. 'mygit', 'BAKDBISDO'
}

type config struct {
	SlackToken      string          // the 'bot token for connecting to Slack
	MaxMessageSplit int             // the maximum # of ~4000 byte messages to split a large message into
	BotRoster       []botDefinition // roster mapping BXXX IDs to a defined username
}

var lock sync.Mutex // package var lock
var started bool    // set when connector is started

func init() {
	bot.RegisterConnector("slack", Initialize)
}

// Initialize starts the connection, sets up and returns the connector object
func Initialize(robot bot.Handler, l *log.Logger) bot.Connector {
	lock.Lock()
	if started {
		lock.Unlock()
		return nil
	}
	started = true
	lock.Unlock()

	var c config
	var tok string

	err := robot.GetProtocolConfig(&c)
	if err != nil {
		robot.Log(bot.Fatal, fmt.Errorf("Unable to retrieve protocol configuration: %v", err))
	}

	if c.MaxMessageSplit == 0 {
		c.MaxMessageSplit = 1
	}

	if len(c.SlackToken) == 0 {
		tok = os.Getenv("SLACK_TOKEN")
		if len(tok) == 0 {
			robot.Log(bot.Fatal, "No slack token found in config or env var 'SLACK_TOKEN'")
		}
		os.Unsetenv("SLACK_TOKEN")
		robot.Log(bot.Debug, "Using SLACK_TOKEN environment variable")
	} else {
		tok = c.SlackToken
	}

	api := slack.New(tok)
	// This spits out a lot of extra stuff, so we only enable it when tracing
	if robot.GetLogLevel() == bot.Trace {
		api.SetDebug(true)
	}
	slack.SetLogger(l)

	sc := &slackConnector{
		api:             api,
		conn:            api.NewRTM(),
		maxMessageSplit: c.MaxMessageSplit,
		name:            "slack",
	}
	sc.botUserID = make(map[string]string)
	sc.botIDUser = make(map[string]string)
	for _, b := range c.BotRoster {
		if len(b.ID) == 0 || len(b.Name) == 0 {
			robot.Log(bot.Error, "Zero-length Name or ID in BotRoster, skipping")
			continue
		}
		sc.botUserID[b.Name] = b.ID
		sc.botIDUser[b.ID] = b.Name
	}
	go sc.conn.ManageConnection()

	sc.Handler = robot

Loop:
	for {
		select {
		case msg := <-sc.conn.IncomingEvents:

			switch ev := msg.Data.(type) {

			case *slack.ConnectedEvent:
				sc.Log(bot.Debug, fmt.Sprintf("Infos: %T %v\n", ev, *ev.Info.User))
				sc.Log(bot.Debug, "Connection counter:", ev.ConnectionCount)
				sc.botName = ev.Info.User.Name
				sc.SetName(sc.botName)
				sc.Log(bot.Info, "Slack setting bot name to", sc.botName)
				sc.botID = ev.Info.User.ID
				sc.Log(bot.Trace, "Set bot ID to", sc.botID)
				sc.teamID = ev.Info.Team.ID
				sc.Log(bot.Info, "Set team ID to", sc.teamID)
				break Loop

			case *slack.InvalidAuthEvent:
				sc.Log(bot.Fatal, "Invalid credentials")
			}
		}
	}

	sc.updateChannelMaps("")
	sc.updateUserList("")
	sc.botFullName, _ = sc.GetProtocolUserAttribute(sc.botName, "realname")
	sc.SetFullName(sc.botFullName)
	sc.Log(bot.Debug, "Set bot full name to", sc.botFullName)
	go sc.startSendLoop()

	return bot.Connector(sc)
}

func (sc *slackConnector) Run(stop <-chan struct{}) {
	sc.Lock()
	// This should never happen, just a bit of defensive coding
	if sc.running {
		sc.Unlock()
		return
	}
	sc.running = true
	sc.Unlock()
loop:
	for {
		select {
		case <-stop:
			sc.Log(bot.Debug, "Received stop in connector")
			break loop
		case msg := <-sc.conn.IncomingEvents:
			sc.Log(bot.Trace, fmt.Sprintf("Event Received (msg, data, type): %v; %v; %T", msg, msg.Data, msg.Data))
			switch ev := msg.Data.(type) {
			case *slack.HelloEvent:
				// Ignore hello
			case *slack.ChannelArchiveEvent, *slack.ChannelUnarchiveEvent,
				*slack.ChannelCreatedEvent, *slack.ChannelDeletedEvent,
				*slack.ChannelRenameEvent, *slack.GroupArchiveEvent,
				*slack.GroupUnarchiveEvent, *slack.GroupCreatedEvent,
				*slack.GroupRenameEvent, *slack.IMCloseEvent,
				*slack.IMCreatedEvent, *slack.IMOpenEvent:
				sc.updateChannelMaps("")

			case *slack.MessageEvent:
				// Message processing is done concurrently
				go sc.processMessage(ev)

			case *slack.PresenceChangeEvent:
				sc.Log(bot.Debug, fmt.Sprintf("Presence Change: %v", ev))

			case *slack.LatencyReport:
				sc.Log(bot.Debug, fmt.Sprintf("Current latency: %v", ev.Value))

			case *slack.RTMError:
				sc.Log(bot.Debug, fmt.Sprintf("Error: %s\n", ev.Error()))

			default:

				// Ignore other events..
				// robot.Debug(fmt.Sprintf("Unexpected: %v\n", msg.Data)
			}
		}
	}
}
<|endoftext|>


<commit_before>package aws

import (
	"bytes"
	"fmt"
	"log"
	"sort"
	"time"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/mitchellh/goamz/ec2"
)

func resourceAwsSecurityGroup() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSecurityGroupCreate,
		Read:   resourceAwsSecurityGroupRead,
		Update: resourceAwsSecurityGroupUpdate,
		Delete: resourceAwsSecurityGroupDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"description": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			"vpc_id": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				Computed: true,
			},

			"ingress": &schema.Schema{
				Type:     schema.TypeSet,
				Required: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"from_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"to_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"protocol": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"cidr_blocks": &schema.Schema{
							Type:     schema.TypeList,
							Optional: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
						},

						"security_groups": &schema.Schema{
							Type:     schema.TypeSet,
							Optional: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
							Set: func(v interface{}) int {
								return hashcode.String(v.(string))
							},
						},

						"self": &schema.Schema{
							Type:     schema.TypeBool,
							Optional: true,
							Default:  false,
						},
					},
				},
				Set: resourceAwsSecurityGroupIngressHash,
			},

			"owner_id": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"tags": tagsSchema(),
		},
	}
}

func resourceAwsSecurityGroupCreate(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	securityGroupOpts := ec2.SecurityGroup{
		Name: d.Get("name").(string),
	}

	if v := d.Get("vpc_id"); v != nil {
		securityGroupOpts.VpcId = v.(string)
	}

	if v := d.Get("description"); v != nil {
		securityGroupOpts.Description = v.(string)
	}

	log.Printf(
		"[DEBUG] Security Group create configuration: %#v", securityGroupOpts)
	createResp, err := ec2conn.CreateSecurityGroup(securityGroupOpts)
	if err != nil {
		return fmt.Errorf("Error creating Security Group: %s", err)
	}

	d.SetId(createResp.Id)

	log.Printf("[INFO] Security Group ID: %s", d.Id())

	// Wait for the security group to truly exist
	log.Printf(
		"[DEBUG] Waiting for Security Group (%s) to exist",
		d.Id())
	stateConf := &resource.StateChangeConf{
		Pending: []string{""},
		Target:  "exists",
		Refresh: SGStateRefreshFunc(ec2conn, d.Id()),
		Timeout: 1 * time.Minute,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return fmt.Errorf(
			"Error waiting for Security Group (%s) to become available: %s",
			d.Id(), err)
	}

	return resourceAwsSecurityGroupUpdate(d, meta)
}

func resourceAwsSecurityGroupRead(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	sgRaw, _, err := SGStateRefreshFunc(ec2conn, d.Id())()
	if err != nil {
		return err
	}
	if sgRaw == nil {
		d.SetId("")
		return nil
	}

	sg := sgRaw.(*ec2.SecurityGroupInfo)

	// Gather our ingress rules
	ingressMap := make(map[string]map[string]interface{})
	for _, perm := range sg.IPPerms {
		k := fmt.Sprintf("%s-%d-%d", perm.Protocol, perm.FromPort, perm.ToPort)
		m, ok := ingressMap[k]
		if !ok {
			m = make(map[string]interface{})
			ingressMap[k] = m
		}

		m["from_port"] = perm.FromPort
		m["to_port"] = perm.ToPort
		m["protocol"] = perm.Protocol

		if len(perm.SourceIPs) > 0 {
			raw, ok := m["cidr_blocks"]
			if !ok {
				raw = make([]string, 0, len(perm.SourceIPs))
			}
			list := raw.([]string)

			list = append(list, perm.SourceIPs...)
			m["cidr_blocks"] = list
		}

		var groups []string
		if len(perm.SourceGroups) > 0 {
			groups = flattenSecurityGroups(perm.SourceGroups)
		}
		for i, id := range groups {
			if id == d.Id() {
				groups[i], groups = groups[len(groups)-1], groups[:len(groups)-1]
				m["self"] = true
			}
		}

		if len(groups) > 0 {
			raw, ok := m["security_groups"]
			if !ok {
				raw = make([]string, 0, len(groups))
			}
			list := raw.([]string)

			list = append(list, groups...)
			m["security_groups"] = list
		}
	}
	ingressRules := make([]map[string]interface{}, 0, len(ingressMap))
	for _, m := range ingressMap {
		ingressRules = append(ingressRules, m)
	}

	d.Set("description", sg.Description)
	d.Set("name", sg.Name)
	d.Set("vpc_id", sg.VpcId)
	d.Set("owner_id", sg.OwnerId)
	d.Set("ingress", ingressRules)
	d.Set("tags", tagsToMap(sg.Tags))

	return nil
}

func resourceAwsSecurityGroupUpdate(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	sgRaw, _, err := SGStateRefreshFunc(ec2conn, d.Id())()
	if err != nil {
		return err
	}
	if sgRaw == nil {
		d.SetId("")
		return nil
	}
	group := sgRaw.(*ec2.SecurityGroupInfo).SecurityGroup

	if d.HasChange("ingress") {
		o, n := d.GetChange("ingress")
		if o == nil {
			o = new(schema.Set)
		}
		if n == nil {
			n = new(schema.Set)
		}

		os := o.(*schema.Set)
		ns := n.(*schema.Set)

		remove := expandIPPerms(d.Id(), os.Difference(ns).List())
		add := expandIPPerms(d.Id(), ns.Difference(os).List())

		// TODO: We need to handle partial state better in the in-between
		// in this update.

		// TODO: It'd be nicer to authorize before removing, but then we have
		// to deal with complicated unrolling to get individual CIDR blocks
		// to avoid authorizing already authorized sources. Removing before
		// adding is easier here, and Terraform should be fast enough to
		// not have service issues.

		if len(remove) > 0 {
			// Revoke the old rules
			_, err = ec2conn.RevokeSecurityGroup(group, remove)
			if err != nil {
				return fmt.Errorf("Error authorizing security group ingress rules: %s", err)
			}
		}

		if len(add) > 0 {
			// Authorize the new rules
			_, err := ec2conn.AuthorizeSecurityGroup(group, add)
			if err != nil {
				return fmt.Errorf("Error authorizing security group ingress rules: %s", err)
			}
		}
	}

	if err := setTags(ec2conn, d); err != nil {
		return err
	} else {
		d.SetPartial("tags")
	}

	return resourceAwsSecurityGroupRead(d, meta)
}

func resourceAwsSecurityGroupDelete(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	log.Printf("[DEBUG] Security Group destroy: %v", d.Id())

	return resource.Retry(5*time.Minute, func() error {
		_, err := ec2conn.DeleteSecurityGroup(ec2.SecurityGroup{Id: d.Id()})
		if err != nil {
			ec2err, ok := err.(*ec2.Error)
			if !ok {
				return err
			}

			switch ec2err.Code {
			case "InvalidGroup.NotFound":
				return nil
			case "DependencyViolation":
				// If it is a dependency violation, we want to retry
				return err
			default:
				// Any other error, we want to quit the retry loop immediately
				return resource.RetryError{err}
			}
		}

		return nil
	})
}

func resourceAwsSecurityGroupIngressHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%d-", m["from_port"].(int)))
	buf.WriteString(fmt.Sprintf("%d-", m["to_port"].(int)))
	buf.WriteString(fmt.Sprintf("%s-", m["protocol"].(string)))

	// We need to make sure to sort the strings below so that we always
	// generate the same hash code no matter what is in the set.
	if v, ok := m["cidr_blocks"]; ok {
		vs := v.([]interface{})
		s := make([]string, len(vs))
		for i, raw := range vs {
			s[i] = raw.(string)
		}
		sort.Strings(s)

		for _, v := range s {
			buf.WriteString(fmt.Sprintf("%s-", v))
		}
	}
	if v, ok := m["security_groups"]; ok {
		vs := v.(*schema.Set).List()
		s := make([]string, len(vs))
		for i, raw := range vs {
			s[i] = raw.(string)
		}
		sort.Strings(s)

		for _, v := range s {
			buf.WriteString(fmt.Sprintf("%s-", v))
		}
	}

	return hashcode.String(buf.String())
}

// SGStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// a security group.
func SGStateRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		sgs := []ec2.SecurityGroup{ec2.SecurityGroup{Id: id}}
		resp, err := conn.SecurityGroups(sgs, nil)
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok {
				if ec2err.Code == "InvalidSecurityGroupID.NotFound" ||
					ec2err.Code == "InvalidGroup.NotFound" {
					resp = nil
					err = nil
				}
			}

			if err != nil {
				log.Printf("Error on SGStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			return nil, "", nil
		}

		group := &resp.Groups[0]
		return group, "exists", nil
	}
}
<commit_msg>aws: Making security group ingress rules optional<commit_after>package aws

import (
	"bytes"
	"fmt"
	"log"
	"sort"
	"time"

	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/mitchellh/goamz/ec2"
)

func resourceAwsSecurityGroup() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSecurityGroupCreate,
		Read:   resourceAwsSecurityGroupRead,
		Update: resourceAwsSecurityGroupUpdate,
		Delete: resourceAwsSecurityGroupDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"description": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			"vpc_id": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				Computed: true,
			},

			"ingress": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"from_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"to_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},

						"protocol": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"cidr_blocks": &schema.Schema{
							Type:     schema.TypeList,
							Optional: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
						},

						"security_groups": &schema.Schema{
							Type:     schema.TypeSet,
							Optional: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
							Set: func(v interface{}) int {
								return hashcode.String(v.(string))
							},
						},

						"self": &schema.Schema{
							Type:     schema.TypeBool,
							Optional: true,
							Default:  false,
						},
					},
				},
				Set: resourceAwsSecurityGroupIngressHash,
			},

			"owner_id": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"tags": tagsSchema(),
		},
	}
}

func resourceAwsSecurityGroupCreate(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	securityGroupOpts := ec2.SecurityGroup{
		Name: d.Get("name").(string),
	}

	if v := d.Get("vpc_id"); v != nil {
		securityGroupOpts.VpcId = v.(string)
	}

	if v := d.Get("description"); v != nil {
		securityGroupOpts.Description = v.(string)
	}

	log.Printf(
		"[DEBUG] Security Group create configuration: %#v", securityGroupOpts)
	createResp, err := ec2conn.CreateSecurityGroup(securityGroupOpts)
	if err != nil {
		return fmt.Errorf("Error creating Security Group: %s", err)
	}

	d.SetId(createResp.Id)

	log.Printf("[INFO] Security Group ID: %s", d.Id())

	// Wait for the security group to truly exist
	log.Printf(
		"[DEBUG] Waiting for Security Group (%s) to exist",
		d.Id())
	stateConf := &resource.StateChangeConf{
		Pending: []string{""},
		Target:  "exists",
		Refresh: SGStateRefreshFunc(ec2conn, d.Id()),
		Timeout: 1 * time.Minute,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return fmt.Errorf(
			"Error waiting for Security Group (%s) to become available: %s",
			d.Id(), err)
	}

	return resourceAwsSecurityGroupUpdate(d, meta)
}

func resourceAwsSecurityGroupRead(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	sgRaw, _, err := SGStateRefreshFunc(ec2conn, d.Id())()
	if err != nil {
		return err
	}
	if sgRaw == nil {
		d.SetId("")
		return nil
	}

	sg := sgRaw.(*ec2.SecurityGroupInfo)

	// Gather our ingress rules
	ingressMap := make(map[string]map[string]interface{})
	for _, perm := range sg.IPPerms {
		k := fmt.Sprintf("%s-%d-%d", perm.Protocol, perm.FromPort, perm.ToPort)
		m, ok := ingressMap[k]
		if !ok {
			m = make(map[string]interface{})
			ingressMap[k] = m
		}

		m["from_port"] = perm.FromPort
		m["to_port"] = perm.ToPort
		m["protocol"] = perm.Protocol

		if len(perm.SourceIPs) > 0 {
			raw, ok := m["cidr_blocks"]
			if !ok {
				raw = make([]string, 0, len(perm.SourceIPs))
			}
			list := raw.([]string)

			list = append(list, perm.SourceIPs...)
			m["cidr_blocks"] = list
		}

		var groups []string
		if len(perm.SourceGroups) > 0 {
			groups = flattenSecurityGroups(perm.SourceGroups)
		}
		for i, id := range groups {
			if id == d.Id() {
				groups[i], groups = groups[len(groups)-1], groups[:len(groups)-1]
				m["self"] = true
			}
		}

		if len(groups) > 0 {
			raw, ok := m["security_groups"]
			if !ok {
				raw = make([]string, 0, len(groups))
			}
			list := raw.([]string)

			list = append(list, groups...)
			m["security_groups"] = list
		}
	}
	ingressRules := make([]map[string]interface{}, 0, len(ingressMap))
	for _, m := range ingressMap {
		ingressRules = append(ingressRules, m)
	}

	d.Set("description", sg.Description)
	d.Set("name", sg.Name)
	d.Set("vpc_id", sg.VpcId)
	d.Set("owner_id", sg.OwnerId)
	d.Set("ingress", ingressRules)
	d.Set("tags", tagsToMap(sg.Tags))

	return nil
}

func resourceAwsSecurityGroupUpdate(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	sgRaw, _, err := SGStateRefreshFunc(ec2conn, d.Id())()
	if err != nil {
		return err
	}
	if sgRaw == nil {
		d.SetId("")
		return nil
	}
	group := sgRaw.(*ec2.SecurityGroupInfo).SecurityGroup

	if d.HasChange("ingress") {
		o, n := d.GetChange("ingress")
		if o == nil {
			o = new(schema.Set)
		}
		if n == nil {
			n = new(schema.Set)
		}

		os := o.(*schema.Set)
		ns := n.(*schema.Set)

		remove := expandIPPerms(d.Id(), os.Difference(ns).List())
		add := expandIPPerms(d.Id(), ns.Difference(os).List())

		// TODO: We need to handle partial state better in the in-between
		// in this update.

		// TODO: It'd be nicer to authorize before removing, but then we have
		// to deal with complicated unrolling to get individual CIDR blocks
		// to avoid authorizing already authorized sources. Removing before
		// adding is easier here, and Terraform should be fast enough to
		// not have service issues.

		if len(remove) > 0 {
			// Revoke the old rules
			_, err = ec2conn.RevokeSecurityGroup(group, remove)
			if err != nil {
				return fmt.Errorf("Error authorizing security group ingress rules: %s", err)
			}
		}

		if len(add) > 0 {
			// Authorize the new rules
			_, err := ec2conn.AuthorizeSecurityGroup(group, add)
			if err != nil {
				return fmt.Errorf("Error authorizing security group ingress rules: %s", err)
			}
		}
	}

	if err := setTags(ec2conn, d); err != nil {
		return err
	} else {
		d.SetPartial("tags")
	}

	return resourceAwsSecurityGroupRead(d, meta)
}

func resourceAwsSecurityGroupDelete(d *schema.ResourceData, meta interface{}) error {
	ec2conn := meta.(*AWSClient).ec2conn

	log.Printf("[DEBUG] Security Group destroy: %v", d.Id())

	return resource.Retry(5*time.Minute, func() error {
		_, err := ec2conn.DeleteSecurityGroup(ec2.SecurityGroup{Id: d.Id()})
		if err != nil {
			ec2err, ok := err.(*ec2.Error)
			if !ok {
				return err
			}

			switch ec2err.Code {
			case "InvalidGroup.NotFound":
				return nil
			case "DependencyViolation":
				// If it is a dependency violation, we want to retry
				return err
			default:
				// Any other error, we want to quit the retry loop immediately
				return resource.RetryError{err}
			}
		}

		return nil
	})
}

func resourceAwsSecurityGroupIngressHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%d-", m["from_port"].(int)))
	buf.WriteString(fmt.Sprintf("%d-", m["to_port"].(int)))
	buf.WriteString(fmt.Sprintf("%s-", m["protocol"].(string)))

	// We need to make sure to sort the strings below so that we always
	// generate the same hash code no matter what is in the set.
	if v, ok := m["cidr_blocks"]; ok {
		vs := v.([]interface{})
		s := make([]string, len(vs))
		for i, raw := range vs {
			s[i] = raw.(string)
		}
		sort.Strings(s)

		for _, v := range s {
			buf.WriteString(fmt.Sprintf("%s-", v))
		}
	}
	if v, ok := m["security_groups"]; ok {
		vs := v.(*schema.Set).List()
		s := make([]string, len(vs))
		for i, raw := range vs {
			s[i] = raw.(string)
		}
		sort.Strings(s)

		for _, v := range s {
			buf.WriteString(fmt.Sprintf("%s-", v))
		}
	}

	return hashcode.String(buf.String())
}

// SGStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// a security group.
func SGStateRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		sgs := []ec2.SecurityGroup{ec2.SecurityGroup{Id: id}}
		resp, err := conn.SecurityGroups(sgs, nil)
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok {
				if ec2err.Code == "InvalidSecurityGroupID.NotFound" ||
					ec2err.Code == "InvalidGroup.NotFound" {
					resp = nil
					err = nil
				}
			}

			if err != nil {
				log.Printf("Error on SGStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			return nil, "", nil
		}

		group := &resp.Groups[0]
		return group, "exists", nil
	}
}
<|endoftext|>


<commit_before>package riak_explorer

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/basho-labs/riak-mesos/cepmd/cepm"
	"github.com/basho-labs/riak-mesos/process_manager"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"text/template"
)

type RiakExplorer struct {
	tempdir string
	port    int64
	pm      *process_manager.ProcessManager
}

type templateData struct {
	HTTPPort int64
	NodeName string
}

type advancedTemplateData struct {
	CEPMDPort int
}

func (re *RiakExplorer) NewRiakExplorerClient() *RiakExplorerClient {
	return NewRiakExplorerClient(fmt.Sprintf("localhost:%d", re.port))
}
func (re *RiakExplorer) TearDown() {
	re.pm.TearDown()
}

func decompress() string {
	assetPath := fmt.Sprintf("riak_explorer_%s_%s.tar.gz", runtime.GOOS, runtime.GOARCH)

	log.Info("Decompressing Riak Explorer: ", assetPath)

	asset, err := Asset(assetPath)
	if err != nil {
		log.Fatal(err)
	}
	bytesReader := bytes.NewReader(asset)

	gzReader, err := gzip.NewReader(bytesReader)

	if err != nil {
		log.Fatal("Error encountered decompressing riak explorer: ", err)
		os.Exit(1)
	}

	tr := tar.NewReader(gzReader)
	tempdir, err := ioutil.TempDir("./", "riak_explorer.")
	if err != nil {
		log.Fatal(err)
	}
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			// end of tar archive
			break
		} else if err != nil {
			log.Fatalln(err)
		}
		filename := filepath.Join(".", tempdir, hdr.Name)
		if hdr.Typeflag == tar.TypeReg || hdr.Typeflag == tar.TypeRegA {
			file, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE, os.FileMode(hdr.Mode))
			io.Copy(file, tr)
			if err != nil {
				log.Fatal(err)
			}
			if _, err := io.Copy(file, tr); err != nil {
				log.Fatal(err)
			}
			file.Close()
		} else if hdr.Typeflag == tar.TypeDir {
			err := os.Mkdir(filename, 0777)
			if err != nil {
				log.Fatalln(err)
			}
		} else if hdr.Typeflag == tar.TypeSymlink {
			if err := os.Symlink(hdr.Linkname, filename); err != nil {
				log.Fatal(err)
			}
			// Hard link
		} else if hdr.Typeflag == tar.TypeLink {
			fmt.Printf("Encountered hardlink: %+v\n", hdr)
			linkdest := filepath.Join(".", tempdir, hdr.Linkname)
			if err := os.Link(linkdest, filename); err != nil {
				log.Fatal(err)
			}
		} else {
			log.Fatal("Experienced unknown tar file type: ", hdr.Typeflag)
		}
	}
	return tempdir
}
func (re *RiakExplorer) configure(port int64, nodename string) {
	data, err := Asset("riak_explorer.conf")
	if err != nil {
		log.Panic("Got error", err)
	}
	tmpl, err := template.New("test").Parse(string(data))

	if err != nil {
		log.Panic(err)
	}

	// Populate template data from the MesosTask
	vars := templateData{}

	vars.NodeName = nodename
	vars.HTTPPort = port
	configpath := filepath.Join(".", re.tempdir, "rex_root", "riak_explorer", "etc", "riak_explorer.conf")
	file, err := os.OpenFile(configpath, os.O_TRUNC|os.O_CREATE|os.O_RDWR, 0)

	if err != nil {
		log.Panic("Unable to open file: ", err)
	}

	err = tmpl.Execute(file, vars)

	if err != nil {
		log.Panic("Got error", err)
	}
}
func (re *RiakExplorer) configureAdvanced(cepmdPort int) {
	data, err := Asset("advanced.config")
	if err != nil {
		log.Panic("Got error", err)
	}
	tmpl, err := template.New("advanced").Parse(string(data))

	if err != nil {
		log.Panic(err)
	}

	// Populate template data from the MesosTask
	vars := advancedTemplateData{}
	vars.CEPMDPort = cepmdPort
	configpath := filepath.Join(".", re.tempdir, "rex_root", "riak_explorer", "etc", "advanced.config")
	file, err := os.OpenFile(configpath, os.O_TRUNC|os.O_CREATE|os.O_RDWR, 0)

	if err != nil {
		log.Panic("Unable to open file: ", err)
	}

	err = tmpl.Execute(file, vars)

	if err != nil {
		log.Panic("Got error", err)
	}
}

func NewRiakExplorer(port int64, nodename string, c *cepm.CEPM) (*RiakExplorer, error) {
	tempdir := decompress()
	exepath := filepath.Join(".", tempdir, "rex_root", "riak_explorer", "bin", "riak_explorer")

	args := []string{"console", "-noinput"}
	healthCheckFun := func() error {
		log.Info("Running healthcheck: ", port)
		_, err := NewRiakExplorerClient(fmt.Sprintf("localhost:%d", port)).Ping()
		log.Info("Healthcheck result ", err)
		return err
	}
	tearDownFun := func() {
		log.Info("Deleting all data in: ", tempdir)
		err := os.RemoveAll(tempdir)
		if err != nil {
			log.Error(err)
		}
	}
	re := &RiakExplorer{
		tempdir: tempdir,
		port:    port,
	}
	if c != nil {
		// This is gross -- we're passing "hidden" state by passing it through the unix environment variables.
		// Fix it -- we should convert the NewRiakExplorer into using a fluent pattern?
		libpath := filepath.Join(".", tempdir, "rex_root", "riak_explorer", "lib", "basho-patches")
		os.Mkdir(libpath, 0777)
		err := cepm.InstallInto(libpath)
		if err != nil {
			log.Panic(err)
		}
		args = append(args, "-no_epmd")
		re.configureAdvanced(c.GetPort())
	}
	re.configure(port, nodename)
	log.Debugf("Starting up Riak Explorer %v", exepath)
	var err error
	chroot := filepath.Join(".", tempdir, "rex_root", "riak_explorer")
	re.pm, err = process_manager.NewProcessManager(tearDownFun, exepath, args, healthCheckFun, &chroot)
	if err != nil {
		log.Error("Could not start Riak Explorer: ", err)
	}

	return re, err
}
<commit_msg>Alter rex to handle chroot relative paths<commit_after>package riak_explorer

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/basho-labs/riak-mesos/cepmd/cepm"
	"github.com/basho-labs/riak-mesos/process_manager"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"text/template"
)

type RiakExplorer struct {
	tempdir string
	port    int64
	pm      *process_manager.ProcessManager
}

type templateData struct {
	HTTPPort int64
	NodeName string
}

type advancedTemplateData struct {
	CEPMDPort int
}

func (re *RiakExplorer) NewRiakExplorerClient() *RiakExplorerClient {
	return NewRiakExplorerClient(fmt.Sprintf("localhost:%d", re.port))
}
func (re *RiakExplorer) TearDown() {
	re.pm.TearDown()
}

func decompress() string {
	assetPath := fmt.Sprintf("riak_explorer_%s_%s.tar.gz", runtime.GOOS, runtime.GOARCH)

	log.Info("Decompressing Riak Explorer: ", assetPath)

	asset, err := Asset(assetPath)
	if err != nil {
		log.Fatal(err)
	}
	bytesReader := bytes.NewReader(asset)

	gzReader, err := gzip.NewReader(bytesReader)

	if err != nil {
		log.Fatal("Error encountered decompressing riak explorer: ", err)
		os.Exit(1)
	}

	tr := tar.NewReader(gzReader)
	tempdir, err := ioutil.TempDir("./", "riak_explorer.")
	if err != nil {
		log.Fatal(err)
	}
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			// end of tar archive
			break
		} else if err != nil {
			log.Fatalln(err)
		}
		filename := filepath.Join(".", tempdir, hdr.Name)
		if hdr.Typeflag == tar.TypeReg || hdr.Typeflag == tar.TypeRegA {
			file, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE, os.FileMode(hdr.Mode))
			io.Copy(file, tr)
			if err != nil {
				log.Fatal(err)
			}
			if _, err := io.Copy(file, tr); err != nil {
				log.Fatal(err)
			}
			file.Close()
		} else if hdr.Typeflag == tar.TypeDir {
			err := os.Mkdir(filename, 0777)
			if err != nil {
				log.Fatalln(err)
			}
		} else if hdr.Typeflag == tar.TypeSymlink {
			if err := os.Symlink(hdr.Linkname, filename); err != nil {
				log.Fatal(err)
			}
			// Hard link
		} else if hdr.Typeflag == tar.TypeLink {
			fmt.Printf("Encountered hardlink: %+v\n", hdr)
			linkdest := filepath.Join(".", tempdir, hdr.Linkname)
			if err := os.Link(linkdest, filename); err != nil {
				log.Fatal(err)
			}
		} else {
			log.Fatal("Experienced unknown tar file type: ", hdr.Typeflag)
		}
	}
	return tempdir
}
func (re *RiakExplorer) configure(port int64, nodename string) {
	data, err := Asset("riak_explorer.conf")
	if err != nil {
		log.Panic("Got error", err)
	}
	tmpl, err := template.New("test").Parse(string(data))

	if err != nil {
		log.Panic(err)
	}

	// Populate template data from the MesosTask
	vars := templateData{}

	vars.NodeName = nodename
	vars.HTTPPort = port
	configpath := filepath.Join(".", re.tempdir, "rex_root", "riak_explorer", "etc", "riak_explorer.conf")
	file, err := os.OpenFile(configpath, os.O_TRUNC|os.O_CREATE|os.O_RDWR, 0)

	if err != nil {
		log.Panic("Unable to open file: ", err)
	}

	err = tmpl.Execute(file, vars)

	if err != nil {
		log.Panic("Got error", err)
	}
}
func (re *RiakExplorer) configureAdvanced(cepmdPort int) {
	data, err := Asset("advanced.config")
	if err != nil {
		log.Panic("Got error", err)
	}
	tmpl, err := template.New("advanced").Parse(string(data))

	if err != nil {
		log.Panic(err)
	}

	// Populate template data from the MesosTask
	vars := advancedTemplateData{}
	vars.CEPMDPort = cepmdPort
	configpath := filepath.Join(".", re.tempdir, "rex_root", "riak_explorer", "etc", "advanced.config")
	file, err := os.OpenFile(configpath, os.O_TRUNC|os.O_CREATE|os.O_RDWR, 0)

	if err != nil {
		log.Panic("Unable to open file: ", err)
	}

	err = tmpl.Execute(file, vars)

	if err != nil {
		log.Panic("Got error", err)
	}
}

func NewRiakExplorer(port int64, nodename string, c *cepm.CEPM) (*RiakExplorer, error) {
	tempdir := decompress()
	exepath := "/riak_explorer/bin/riak_explorer"

	args := []string{"console", "-noinput"}
	healthCheckFun := func() error {
		log.Info("Running healthcheck: ", port)
		_, err := NewRiakExplorerClient(fmt.Sprintf("localhost:%d", port)).Ping()
		log.Info("Healthcheck result ", err)
		return err
	}
	tearDownFun := func() {
		log.Info("Deleting all data in: ", tempdir)
		err := os.RemoveAll(tempdir)
		if err != nil {
			log.Error(err)
		}
	}
	re := &RiakExplorer{
		tempdir: tempdir,
		port:    port,
	}
	if c != nil {
		// This is gross -- we're passing "hidden" state by passing it through the unix environment variables.
		// Fix it -- we should convert the NewRiakExplorer into using a fluent pattern?
		libpath := filepath.Join(".", tempdir, "rex_root", "riak_explorer", "lib", "basho-patches")
		os.Mkdir(libpath, 0777)
		err := cepm.InstallInto(libpath)
		if err != nil {
			log.Panic(err)
		}
		args = append(args, "-no_epmd")
		re.configureAdvanced(c.GetPort())
	}
	re.configure(port, nodename)
	log.Debugf("Starting up Riak Explorer %v", exepath)
	var err error
	chroot := filepath.Join(".", tempdir, "rex_root")
	re.pm, err = process_manager.NewProcessManager(tearDownFun, exepath, args, healthCheckFun, &chroot)
	if err != nil {
		log.Error("Could not start Riak Explorer: ", err)
	}

	return re, err
}
<|endoftext|>


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

package filter

import (
	"os"
	"syscall"

	"golang.org/x/sys/unix"
	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/seccomp"
)

// allowedSyscalls is the set of syscalls executed by the gofer.
var allowedSyscalls = seccomp.SyscallRules{
	syscall.SYS_ACCEPT: {},
	syscall.SYS_SOCKET: []seccomp.Rule{
		{
			seccomp.AllowValue(syscall.AF_UNIX),
		},
	},
	syscall.SYS_CONNECT: []seccomp.Rule{
		{
			seccomp.AllowAny{},
		},
	},
	syscall.SYS_SETSOCKOPT: []seccomp.Rule{
		{
			seccomp.AllowAny{},
		},
	},
	syscall.SYS_GETSOCKNAME: []seccomp.Rule{
		{
			seccomp.AllowAny{},
		},
	},
	syscall.SYS_GETPEERNAME: []seccomp.Rule{
		{
			seccomp.AllowAny{},
		},
	},
	syscall.SYS_ARCH_PRCTL: []seccomp.Rule{
		{seccomp.AllowValue(linux.ARCH_GET_FS)},
		{seccomp.AllowValue(linux.ARCH_SET_FS)},
	},
	syscall.SYS_CLOCK_GETTIME: {},
	syscall.SYS_CLONE: []seccomp.Rule{
		{
			seccomp.AllowValue(
				syscall.CLONE_VM |
					syscall.CLONE_FS |
					syscall.CLONE_FILES |
					syscall.CLONE_SIGHAND |
					syscall.CLONE_SYSVSEM |
					syscall.CLONE_THREAD),
		},
	},
	syscall.SYS_CLOSE:     {},
	syscall.SYS_DUP:       {},
	syscall.SYS_EPOLL_CTL: {},
	syscall.SYS_EPOLL_PWAIT: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(0),
		},
	},
	syscall.SYS_EVENTFD2: []seccomp.Rule{
		{
			seccomp.AllowValue(0),
			seccomp.AllowValue(0),
		},
	},
	syscall.SYS_EXIT:       {},
	syscall.SYS_EXIT_GROUP: {},
	syscall.SYS_FALLOCATE: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowValue(0),
		},
	},
	syscall.SYS_FCHMOD:   {},
	syscall.SYS_FCHOWNAT: {},
	syscall.SYS_FCNTL: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.F_GETFL),
		},
		{
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.F_SETFL),
		},
		{
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.F_GETFD),
		},
		{
			seccomp.AllowAny{},
		},
	},
	syscall.SYS_FSTAT:     {},
	syscall.SYS_FSTATFS:   {},
	syscall.SYS_FSYNC:     {},
	syscall.SYS_FTRUNCATE: {},
	syscall.SYS_FUTEX: {
		seccomp.Rule{
			seccomp.AllowAny{},
			seccomp.AllowValue(linux.FUTEX_WAIT | linux.FUTEX_PRIVATE_FLAG),
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(0),
		},
		seccomp.Rule{
			seccomp.AllowAny{},
			seccomp.AllowValue(linux.FUTEX_WAKE | linux.FUTEX_PRIVATE_FLAG),
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(0),
		},
	},
	syscall.SYS_GETDENTS64:   {},
	syscall.SYS_GETPID:       {},
	unix.SYS_GETRANDOM:       {},
	syscall.SYS_GETTID:       {},
	syscall.SYS_GETTIMEOFDAY: {},
	syscall.SYS_LINKAT:       {},
	syscall.SYS_LSEEK:        {},
	syscall.SYS_MADVISE:      {},
	syscall.SYS_MKDIRAT:      {},
	syscall.SYS_MMAP: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.MAP_SHARED),
		},
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.MAP_PRIVATE | syscall.MAP_ANONYMOUS),
		},
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.MAP_PRIVATE | syscall.MAP_ANONYMOUS | syscall.MAP_FIXED),
		},
	},
	syscall.SYS_MPROTECT:   {},
	syscall.SYS_MUNMAP:     {},
	syscall.SYS_NANOSLEEP:  {},
	syscall.SYS_NEWFSTATAT: {},
	syscall.SYS_OPENAT:     {},
	syscall.SYS_PPOLL:      {},
	syscall.SYS_PREAD64:    {},
	syscall.SYS_PWRITE64:   {},
	syscall.SYS_READ:       {},
	syscall.SYS_READLINKAT: {},
	syscall.SYS_RECVMSG: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.MSG_DONTWAIT | syscall.MSG_TRUNC),
		},
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.MSG_DONTWAIT | syscall.MSG_TRUNC | syscall.MSG_PEEK),
		},
	},
	syscall.SYS_RENAMEAT:        {},
	syscall.SYS_RESTART_SYSCALL: {},
	syscall.SYS_RT_SIGPROCMASK:  {},
	syscall.SYS_SCHED_YIELD:     {},
	syscall.SYS_SENDMSG: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.MSG_DONTWAIT | syscall.MSG_NOSIGNAL),
		},
	},
	syscall.SYS_SHUTDOWN: []seccomp.Rule{
		{seccomp.AllowAny{}, seccomp.AllowValue(syscall.SHUT_RDWR)},
	},
	syscall.SYS_SIGALTSTACK: {},
	syscall.SYS_SYMLINKAT:   {},
	syscall.SYS_TGKILL: []seccomp.Rule{
		{
			seccomp.AllowValue(uint64(os.Getpid())),
		},
	},
	syscall.SYS_UNLINKAT:  {},
	syscall.SYS_UTIMENSAT: {},
	syscall.SYS_WRITE:     {},
}
<commit_msg>Restrict seccomp filters for UDS support.<commit_after>// Copyright 2018 The gVisor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package filter

import (
	"os"
	"syscall"

	"golang.org/x/sys/unix"
	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/seccomp"
)

// allowedSyscalls is the set of syscalls executed by the gofer.
var allowedSyscalls = seccomp.SyscallRules{
	syscall.SYS_ACCEPT: {},
	syscall.SYS_SOCKET: []seccomp.Rule{
		{
			seccomp.AllowValue(syscall.AF_UNIX),
		},
	},
	syscall.SYS_CONNECT: []seccomp.Rule{
		{
			seccomp.AllowAny{},
		},
	},
	syscall.SYS_SETSOCKOPT: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.SOL_SOCKET),
			seccomp.AllowValue(syscall.SO_BROADCAST),
		},
	},
	syscall.SYS_GETSOCKNAME: []seccomp.Rule{
		{
			seccomp.AllowAny{},
		},
	},
	syscall.SYS_GETPEERNAME: []seccomp.Rule{
		{
			seccomp.AllowAny{},
		},
	},
	syscall.SYS_ARCH_PRCTL: []seccomp.Rule{
		{seccomp.AllowValue(linux.ARCH_GET_FS)},
		{seccomp.AllowValue(linux.ARCH_SET_FS)},
	},
	syscall.SYS_CLOCK_GETTIME: {},
	syscall.SYS_CLONE: []seccomp.Rule{
		{
			seccomp.AllowValue(
				syscall.CLONE_VM |
					syscall.CLONE_FS |
					syscall.CLONE_FILES |
					syscall.CLONE_SIGHAND |
					syscall.CLONE_SYSVSEM |
					syscall.CLONE_THREAD),
		},
	},
	syscall.SYS_CLOSE:     {},
	syscall.SYS_DUP:       {},
	syscall.SYS_EPOLL_CTL: {},
	syscall.SYS_EPOLL_PWAIT: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(0),
		},
	},
	syscall.SYS_EVENTFD2: []seccomp.Rule{
		{
			seccomp.AllowValue(0),
			seccomp.AllowValue(0),
		},
	},
	syscall.SYS_EXIT:       {},
	syscall.SYS_EXIT_GROUP: {},
	syscall.SYS_FALLOCATE: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowValue(0),
		},
	},
	syscall.SYS_FCHMOD:   {},
	syscall.SYS_FCHOWNAT: {},
	syscall.SYS_FCNTL: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.F_GETFL),
		},
		{
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.F_SETFL),
		},
		{
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.F_GETFD),
		},
		{
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.F_DUPFD_CLOEXEC),
		},
	},
	syscall.SYS_FSTAT:     {},
	syscall.SYS_FSTATFS:   {},
	syscall.SYS_FSYNC:     {},
	syscall.SYS_FTRUNCATE: {},
	syscall.SYS_FUTEX: {
		seccomp.Rule{
			seccomp.AllowAny{},
			seccomp.AllowValue(linux.FUTEX_WAIT | linux.FUTEX_PRIVATE_FLAG),
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(0),
		},
		seccomp.Rule{
			seccomp.AllowAny{},
			seccomp.AllowValue(linux.FUTEX_WAKE | linux.FUTEX_PRIVATE_FLAG),
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(0),
		},
	},
	syscall.SYS_GETDENTS64:   {},
	syscall.SYS_GETPID:       {},
	unix.SYS_GETRANDOM:       {},
	syscall.SYS_GETTID:       {},
	syscall.SYS_GETTIMEOFDAY: {},
	syscall.SYS_LINKAT:       {},
	syscall.SYS_LSEEK:        {},
	syscall.SYS_MADVISE:      {},
	syscall.SYS_MKDIRAT:      {},
	syscall.SYS_MMAP: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.MAP_SHARED),
		},
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.MAP_PRIVATE | syscall.MAP_ANONYMOUS),
		},
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.MAP_PRIVATE | syscall.MAP_ANONYMOUS | syscall.MAP_FIXED),
		},
	},
	syscall.SYS_MPROTECT:   {},
	syscall.SYS_MUNMAP:     {},
	syscall.SYS_NANOSLEEP:  {},
	syscall.SYS_NEWFSTATAT: {},
	syscall.SYS_OPENAT:     {},
	syscall.SYS_PPOLL:      {},
	syscall.SYS_PREAD64:    {},
	syscall.SYS_PWRITE64:   {},
	syscall.SYS_READ:       {},
	syscall.SYS_READLINKAT: {},
	syscall.SYS_RECVMSG: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.MSG_DONTWAIT | syscall.MSG_TRUNC),
		},
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.MSG_DONTWAIT | syscall.MSG_TRUNC | syscall.MSG_PEEK),
		},
	},
	syscall.SYS_RENAMEAT:        {},
	syscall.SYS_RESTART_SYSCALL: {},
	syscall.SYS_RT_SIGPROCMASK:  {},
	syscall.SYS_SCHED_YIELD:     {},
	syscall.SYS_SENDMSG: []seccomp.Rule{
		{
			seccomp.AllowAny{},
			seccomp.AllowAny{},
			seccomp.AllowValue(syscall.MSG_DONTWAIT | syscall.MSG_NOSIGNAL),
		},
	},
	syscall.SYS_SHUTDOWN: []seccomp.Rule{
		{seccomp.AllowAny{}, seccomp.AllowValue(syscall.SHUT_RDWR)},
	},
	syscall.SYS_SIGALTSTACK: {},
	syscall.SYS_SYMLINKAT:   {},
	syscall.SYS_TGKILL: []seccomp.Rule{
		{
			seccomp.AllowValue(uint64(os.Getpid())),
		},
	},
	syscall.SYS_UNLINKAT:  {},
	syscall.SYS_UTIMENSAT: {},
	syscall.SYS_WRITE:     {},
}
<|endoftext|>


<commit_before>package sck

import (
	"bytes"
	"net"
	"testing"
	"time"
)

// Define a mock Connection for testing purposes. This will read to and write from a byte
// buffer.
type MockConn struct {
	Buffer    []byte
	SetBuffer chan []byte
}

func (c *MockConn) Read(b []byte) (int, error) {
	data := <-c.SetBuffer
	copy(b, data)
	return len(b), nil
}

func (c *MockConn) Write(b []byte) (int, error) {
	c.Buffer = make([]byte, len(b))

	// This will block until the calling thread is ready for the read to happen.
	<-c.SetBuffer

	copy(c.Buffer, b)
	return len(b), nil
}

func (c *MockConn) Close() error {
	return nil
}

func (c *MockConn) LocalAddr() net.Addr {
	return nil
}

func (c *MockConn) RemoteAddr() net.Addr {
	return nil
}

func (c *MockConn) SetDeadline(t time.Time) error {
	return nil
}

func (c *MockConn) SetReadDeadline(t time.Time) error {
	return nil
}

func (c *MockConn) SetWriteDeadline(t time.Time) error {
	return nil
}

// TestSender tests the sck.Sender() goroutine. This test appears to have unsynchronised
// access to a buffer but it actually doesn't, because the channel blocks the other goroutine.
func TestSender(t *testing.T) {
	// Set up our variables.
	dataChan := make(chan []byte)
	conn := new(MockConn)
	conn.SetBuffer = make(chan []byte)

	// Prepare the messages we're going to send.
	messages := [][]byte{
		[]byte("PASS secretpasswordhere"),
		[]byte("SERVICE dict * *.fr 0 0 :French Dictionary"),
		[]byte(":syrk!kalt@millennium.stealth.net QUIT :Gone to have lunch"),
	}

	// Start the goroutine.
	go Sender(conn, dataChan)

	for _, msg := range messages {
		dataChan <- msg
		conn.SetBuffer <- msg

		// Sleep to give the other goroutine some time to make the copy. This slows the tests,
		// but we just have to accept that.
		time.Sleep(200)

		if cmp := bytes.Compare(conn.Buffer, msg); cmp != 0 {
			t.Errorf(
				"Failed to write correctly: expected %v, got %v.",
				msg,
				conn.Buffer)
		}
	}

	// Close the channel.
	close(dataChan)
}

// TestReceiver tests the sck.Receiver() goroutine.
func TestReceiver(t *testing.T) {
	// Set up our variables.
	dataChan := make(chan []byte)
	conn := new(MockConn)
	conn.SetBuffer = make(chan []byte)

	// Prepare the messages we're going to receive.
	messages := [][]byte{
		[]byte("PASS secretpasswordhere"),
		[]byte("SERVICE dict * *.fr 0 0 :French Dictionary"),
		[]byte(":syrk!kalt@millennium.stealth.net QUIT :Gone to have lunch"),
	}

	// Start the goroutine.
	go Receiver(conn, dataChan)

	for _, msg := range messages {
		// Send the message in.
		conn.SetBuffer <- msg

		// Receieve the message back out.
		recvMsg := <-dataChan
		recvMsg = bytes.TrimRight(recvMsg, "\x00")

		// If they aren't the same, something horrible happened.
		if cmp := bytes.Compare(recvMsg, msg); cmp != 0 {
			t.Errorf(
				"Failed to read correctly: expected %v, got %v",
				msg,
				recvMsg)
		}
	}

	// Close the channel and confirm the tests don't panic.
	close(dataChan)
	time.Sleep(50 * time.Millisecond)
}
<commit_msg>Correct previous incorrect test.<commit_after>package sck

import (
	"bytes"
	"net"
	"testing"
	"time"
)

// Define a mock Connection for testing purposes. This will read to and write from a byte
// buffer.
type MockConn struct {
	Buffer    []byte
	SetBuffer chan []byte
}

func (c *MockConn) Read(b []byte) (int, error) {
	data := <-c.SetBuffer
	copy(b, data)
	return len(b), nil
}

func (c *MockConn) Write(b []byte) (int, error) {
	c.Buffer = make([]byte, len(b))

	// This will block until the calling thread is ready for the read to happen.
	<-c.SetBuffer

	copy(c.Buffer, b)
	return len(b), nil
}

func (c *MockConn) Close() error {
	return nil
}

func (c *MockConn) LocalAddr() net.Addr {
	return nil
}

func (c *MockConn) RemoteAddr() net.Addr {
	return nil
}

func (c *MockConn) SetDeadline(t time.Time) error {
	return nil
}

func (c *MockConn) SetReadDeadline(t time.Time) error {
	return nil
}

func (c *MockConn) SetWriteDeadline(t time.Time) error {
	return nil
}

// TestSender tests the sck.Sender() goroutine. This test appears to have unsynchronised
// access to a buffer but it actually doesn't, because the channel blocks the other goroutine.
func TestSender(t *testing.T) {
	// Set up our variables.
	dataChan := make(chan []byte)
	conn := new(MockConn)
	conn.SetBuffer = make(chan []byte)

	// Prepare the messages we're going to send.
	messages := [][]byte{
		[]byte("PASS secretpasswordhere"),
		[]byte("SERVICE dict * *.fr 0 0 :French Dictionary"),
		[]byte(":syrk!kalt@millennium.stealth.net QUIT :Gone to have lunch"),
	}

	// Start the goroutine.
	go Sender(conn, dataChan)

	for _, msg := range messages {
		dataChan <- msg
		conn.SetBuffer <- msg

		// Sleep to give the other goroutine some time to make the copy. This slows the tests,
		// but we just have to accept that.
		time.Sleep(200)

		if cmp := bytes.Compare(conn.Buffer, msg); cmp != 0 {
			t.Errorf(
				"Failed to write correctly: expected %v, got %v.",
				msg,
				conn.Buffer)
		}
	}

	// Close the channel.
	close(dataChan)
}

// TestReceiver tests the sck.Receiver() goroutine.
func TestReceiver(t *testing.T) {
	// Set up our variables.
	dataChan := make(chan []byte)
	conn := new(MockConn)
	conn.SetBuffer = make(chan []byte)

	// Prepare the messages we're going to receive.
	messages := [][]byte{
		[]byte("PASS secretpasswordhere"),
		[]byte("SERVICE dict * *.fr 0 0 :French Dictionary"),
		[]byte(":syrk!kalt@millennium.stealth.net QUIT :Gone to have lunch"),
	}

	// Start the goroutine.
	go Receiver(conn, dataChan)

	for _, msg := range messages {
		// Send the message in.
		conn.SetBuffer <- msg

		// Receieve the message back out.
		recvMsg := <-dataChan
		recvMsg = bytes.TrimRight(recvMsg, "\x00")

		// If they aren't the same, something horrible happened.
		if cmp := bytes.Compare(recvMsg, msg); cmp != 0 {
			t.Errorf(
				"Failed to read correctly: expected %v, got %v",
				msg,
				recvMsg)
		}
	}

	// Close the channel and confirm the tests don't panic.
	close(dataChan)
	conn.SetBuffer <- []byte("Last test")
	time.Sleep(100 * time.Millisecond)
}
<|endoftext|>