<commit_before>package websocket

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"github.com/op/go-logging"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

const WS_WRITE_CAPACITY = 100

func newWs(baseURL string, logTransport bool, log *logging.Logger) *ws {
	return &ws{
		BaseURL:      baseURL,
		downstream:   make(chan []byte),
		quit:         make(chan error),
		logTransport: logTransport,
		log:          log,
		lock:         &sync.RWMutex{},
		createTime:   time.Now(),
		writeChan:    make(chan []byte, WS_WRITE_CAPACITY),
	}
}

type ws struct {
	ws            *websocket.Conn
	lock          *sync.RWMutex
	BaseURL       string
	TLSSkipVerify bool
	downstream    chan []byte
	logTransport  bool
	log           *logging.Logger
	createTime    time.Time
	writeChan     chan []byte

	quit chan error    // signal to parent with error, if applicable
}

func (w *ws) Connect() error {
	if w.ws != nil {
		return nil // no op
	}
	var d = websocket.Dialer{
		Subprotocols:    []string{"p1", "p2"},
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		Proxy:           http.ProxyFromEnvironment,
	}

	d.TLSClientConfig = &tls.Config{InsecureSkipVerify: w.TLSSkipVerify}

	w.log.Info("connecting ws to %s", w.BaseURL)
	ws, resp, err := d.Dial(w.BaseURL, nil)
	if err != nil {
		if err == websocket.ErrBadHandshake {
			w.log.Errorf("bad handshake: status code %d", resp.StatusCode)
		}
		return err
	}
	w.ws = ws
	go w.listenWriteChannel()
	go w.listenWs()
	return nil
}

// Send marshals the given interface and then sends it to the API. This method
// can block so specify a context with timeout if you don't want to wait for too
// long.
func (w *ws) Send(ctx context.Context, msg interface{}) error {
	if w.ws == nil {
		return ErrWSNotConnected
	}

	bs, err := json.Marshal(msg)
	if err != nil {
		return err
	}

	select {
	case <- ctx.Done():
		return ctx.Err()
	case <- w.quit: // ws closed
		return fmt.Errorf("websocket connection closed")
	default:
	}
	w.log.Debug("ws->srv: %s", string(bs))
	// push request into writer channel
	w.writeChan <- bs
	if err != nil {
		return err
	}
	return nil
}

func (w *ws) Done() <-chan error {
	return w.quit
}

// listen for write requests and perform them
func (w *ws) listenWriteChannel() {
	for {
		select {
		case <- w.quit: // ws closed
			return
		case message := <- w.writeChan:
			wsWriter, err := w.ws.NextWriter(websocket.TextMessage)
			if err != nil {
				w.log.Error("Unable to provision ws connection writer: ", err)
				w.stop(err)
				return
			}
			_, err = wsWriter.Write(message)
			if err != nil {
				w.log.Error("Unable to write to ws: ", err)
				w.stop(err)
				return
			}
			if err := wsWriter.Close(); err != nil {
				w.log.Error("Unable to close ws connection writer: ", err)
				w.stop(err)
				return
			}
		}
	}
}

// listen on ws & fwd to listen()
func (w *ws) listenWs() {
	for {
		if w.ws == nil {
			return
		}

		select {
		case <-w.quit: // ws connection ended
			return
		default:
			_, msg, err := w.ws.ReadMessage()
			if err != nil {
				if cl, ok := err.(*websocket.CloseError); ok {
					w.log.Errorf("close error code: %d", cl.Code)
				}
				// a read during normal shutdown results in an OpError: op on closed connection
				if _, ok := err.(*net.OpError); ok {
					// general read error on a closed network connection, OK
					return
				}
				w.stop(err)
				return
			}
			w.log.Debugf("srv->ws: %s", string(msg))
			w.downstream <- msg
		}
	}
}

func (w *ws) Listen() <-chan []byte {
	return w.downstream
}

func (w *ws) stop(err error) {
	w.lock.Lock()
	defer w.lock.Unlock()
	if w.ws != nil {
		w.quit <- err	// pass error back
		close(w.quit) // signal to parent listeners
		close(w.downstream)
		if err := w.ws.Close(); err != nil {
			panic(fmt.Errorf("error closing websocket: %s", err))
		}
		w.ws = nil
	}
}

// Close the websocket connection
func (w *ws) Close() {
	w.stop(fmt.Errorf("transport connection Close called"))
}
<commit_msg>transport: lock rMutex publish to listeners<commit_after>package websocket

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"github.com/op/go-logging"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// size of channel that the websocket writer
// routine pulls from
const WS_WRITE_CAPACITY = 100
// size of channel that the websocket reader
// routine pushes websocket updates into
const WS_READ_CAPACITY = 10

func newWs(baseURL string, logTransport bool, log *logging.Logger) *ws {
	return &ws{
		BaseURL:      baseURL,
		downstream:   make(chan []byte, WS_READ_CAPACITY),
		quit:         make(chan error),
		logTransport: logTransport,
		log:          log,
		lock:         &sync.RWMutex{},
		createTime:   time.Now(),
		writeChan:    make(chan []byte, WS_WRITE_CAPACITY),
	}
}

type ws struct {
	ws            *websocket.Conn
	lock          *sync.RWMutex
	BaseURL       string
	TLSSkipVerify bool
	downstream    chan []byte
	logTransport  bool
	log           *logging.Logger
	createTime    time.Time
	writeChan     chan []byte

	quit chan error    // signal to parent with error, if applicable
}

func (w *ws) Connect() error {
	if w.ws != nil {
		return nil // no op
	}
	var d = websocket.Dialer{
		Subprotocols:    []string{"p1", "p2"},
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		Proxy:           http.ProxyFromEnvironment,
		HandshakeTimeout: time.Second * 10,
	}

	d.TLSClientConfig = &tls.Config{InsecureSkipVerify: w.TLSSkipVerify}

	w.log.Info("connecting ws to %s", w.BaseURL)
	ws, resp, err := d.Dial(w.BaseURL, nil)
	if err != nil {
		if err == websocket.ErrBadHandshake {
			w.log.Errorf("bad handshake: status code %d", resp.StatusCode)
		}
		return err
	}
	w.ws = ws
	go w.listenWriteChannel()
	go w.listenWs()
	return nil
}

// Send marshals the given interface and then sends it to the API. This method
// can block so specify a context with timeout if you don't want to wait for too
// long.
func (w *ws) Send(ctx context.Context, msg interface{}) error {
	if w.ws == nil {
		return ErrWSNotConnected
	}

	bs, err := json.Marshal(msg)
	if err != nil {
		return err
	}

	select {
	case <- ctx.Done():
		return ctx.Err()
	case <- w.quit: // ws closed
		return fmt.Errorf("websocket connection closed")
	default:
	}
	w.log.Debug("ws->srv: %s", string(bs))
	// push request into writer channel
	w.writeChan <- bs
	if err != nil {
		return err
	}
	return nil
}

func (w *ws) Done() <-chan error {
	return w.quit
}

// listen for write requests and perform them
func (w *ws) listenWriteChannel() {
	for {
		select {
		case <- w.quit: // ws closed
			return
		case message := <- w.writeChan:
			wsWriter, err := w.ws.NextWriter(websocket.TextMessage)
			if err != nil {
				w.log.Error("Unable to provision ws connection writer: ", err)
				w.stop(err)
				return
			}
			_, err = wsWriter.Write(message)
			if err != nil {
				w.log.Error("Unable to write to ws: ", err)
				w.stop(err)
				return
			}
			if err := wsWriter.Close(); err != nil {
				w.log.Error("Unable to close ws connection writer: ", err)
				w.stop(err)
				return
			}
		}
	}
}

// listen on ws & fwd to listen()
func (w *ws) listenWs() {
	for {
		if w.ws == nil {
			return
		}
		select {
		case <-w.quit: // ws connection ended
			return
		default:
			_, msg, err := w.ws.ReadMessage()
			if err != nil {
				if cl, ok := err.(*websocket.CloseError); ok {
					w.log.Errorf("close error code: %d", cl.Code)
				}
				// a read during normal shutdown results in an OpError: op on closed connection
				if _, ok := err.(*net.OpError); ok {
					// general read error on a closed network connection, OK
					return
				}
				w.stop(err)
				return
			}
			w.log.Debugf("srv->ws: %s", string(msg))
			w.lock.RLock()
			w.downstream <- msg
			w.lock.RUnlock()
		}
	}
}

func (w *ws) Listen() <-chan []byte {
	return w.downstream
}

func (w *ws) stop(err error) {
	w.lock.Lock()
	defer w.lock.Unlock()
	if w.ws != nil {
		w.quit <- err	// pass error back
		close(w.quit) // signal to parent listeners
		close(w.downstream)
		if err := w.ws.Close(); err != nil {
			panic(fmt.Errorf("error closing websocket: %s", err))
		}
		w.ws = nil
	}
}

// Close the websocket connection
func (w *ws) Close() {
	w.stop(fmt.Errorf("transport connection Close called"))
}
<|endoftext|>


<commit_before>package test

import (
	"bytes"
	"fmt"
	_ "github.com/UserStack/ustackweb/routers"
	"net/http"
	"net/http/httptest"
	"net/url"
	"path/filepath"
	"runtime"
	"strconv"
	"testing"

	"github.com/UserStack/ustackweb/backend"
	"github.com/UserStack/ustackweb/models"
	"github.com/astaxie/beego"
	. "github.com/smartystreets/goconvey/convey"
)

func init() {
	_, file, _, _ := runtime.Caller(1)
	apppath, _ := filepath.Abs(filepath.Dir(filepath.Join(file, ".."+string(filepath.Separator))))
	backend.Type = backend.Memory
	models.Permissions().Create()
	models.Users().Create("admin", "admin")
	models.Permissions().Allow("admin", "list_users")
	models.Permissions().Allow("admin", "list_groups")
	beego.TestBeegoInit(apppath)
}

type Session struct {
	username string
}

func recordRequest(r *http.Request, session *Session) *httptest.ResponseRecorder {
	w := httptest.NewRecorder()
	if session != nil {
		s := beego.GlobalSessions.SessionStart(w, r)
		s.Set("username", session.username)
	}
	beego.BeeApp.Handlers.ServeHTTP(w, r)
	// beego.Trace("testing", "TestMain", "Code[%d]\n%s", w.Code, w.Body.String())
	return w
}

func getRequest(method string, urlStr string, session *Session) *httptest.ResponseRecorder {
	r, _ := http.NewRequest(method, urlStr, nil)
	return recordRequest(r, session)
}

func postRequest(method string, resourcePath string, data *url.Values, session *Session) *httptest.ResponseRecorder {
	u, _ := url.ParseRequestURI("/")
	u.Path = resourcePath
	urlStr := fmt.Sprintf("%v", u)

	r, _ := http.NewRequest(method, urlStr, bytes.NewBufferString(data.Encode()))
	r.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	r.Header.Add("Content-Length", strconv.Itoa(len(data.Encode())))
	return recordRequest(r, session)
}

// TestMain is a sample to run an endpoint test
func TestMain(t *testing.T) {
	var nilSession *Session
	adminSession := &Session{username: "admin"}

	Convey("Redirect to Sign In\n", t, func() {
		response := getRequest("GET", "/", nilSession)
		Convey("Redirect", func() {
			So(response.Code, ShouldEqual, 302)
			So(response.HeaderMap.Get("Location"), ShouldEqual, "/sign_in")
		})
	})

	Convey("Redirect to Profile when already Signed In\n", t, func() {
		response := getRequest("GET", "/", adminSession)
		Convey("Redirect", func() {
			So(response.Code, ShouldEqual, 200)
			So(response.Body.String(), ShouldContainSubstring, "Home")
		})
	})

	Convey("Shows Sign In\n", t, func() {
		response := getRequest("GET", "/sign_in", nilSession)
		Convey("Redirect", func() {
			So(response.Code, ShouldEqual, 200)
			So(response.Body.String(), ShouldContainSubstring, "Sign In")
		})
	})

	Convey("Successful Sign In\n", t, func() {
		data := url.Values{}
		data.Add("Username", "admin")
		data.Add("Password", "admin")
		response := postRequest("POST", "/sign_in", &data, nilSession)
		Convey("Redirect", func() {
			So(response.Code, ShouldEqual, 302)
			So(response.HeaderMap.Get("Location"), ShouldEqual, "/")
		})
	})

	Convey("Failed Sign In\n", t, func() {
		data := url.Values{}
		data.Add("Username", "adminx")
		data.Add("Password", "barx")
		response := postRequest("POST", "/sign_in", &data, nilSession)
		Convey("Redirect", func() {
			So(response.Code, ShouldEqual, 302)
			So(response.HeaderMap.Get("Location"), ShouldEqual, "/sign_in")
		})
	})

	Convey("Users without Sign In\n", t, func() {
		response := postRequest("GET", "/users", &url.Values{}, nilSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 302)
			So(response.HeaderMap.Get("Location"), ShouldEqual, "/sign_in")
		})
	})

	Convey("Users\n", t, func() {
		response := postRequest("GET", "/users", &url.Values{}, adminSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 200)
		})
	})

	Convey("Create User\n", t, func() {
		data := url.Values{}
		data.Add("Username", "mikes")
		data.Add("Password", "micke")
		response := postRequest("POST", "/users", &data, adminSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 302)
			So(response.HeaderMap.Get("Location"), ShouldStartWith, "/users/")
		})
	})

	Convey("Create User Error\n", t, func() {
		response := postRequest("POST", "/users", &url.Values{}, adminSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 200)
			So(response.Body.String(), ShouldContainSubstring, "Can not be empty")
		})
	})

	Convey("Create Group\n", t, func() {
		data := url.Values{}
		data.Add("Name", "foo")
		response := postRequest("POST", "/groups", &data, adminSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 302)
			So(response.HeaderMap.Get("Location"), ShouldStartWith, "/groups")
		})
	})

	Convey("Create User Error\n", t, func() {
		response := postRequest("POST", "/groups", &url.Values{}, adminSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 200)
			So(response.Body.String(), ShouldContainSubstring, "Can not be empty")
		})
	})
}
<commit_msg>Test that user has permission to create group.<commit_after>package test

import (
	"bytes"
	"fmt"
	_ "github.com/UserStack/ustackweb/routers"
	"net/http"
	"net/http/httptest"
	"net/url"
	"path/filepath"
	"runtime"
	"strconv"
	"testing"

	"github.com/UserStack/ustackweb/backend"
	"github.com/UserStack/ustackweb/models"
	"github.com/astaxie/beego"
	. "github.com/smartystreets/goconvey/convey"
)

func init() {
	_, file, _, _ := runtime.Caller(1)
	apppath, _ := filepath.Abs(filepath.Dir(filepath.Join(file, ".."+string(filepath.Separator))))
	backend.Type = backend.Memory
	beego.TestBeegoInit(apppath)
}

type Session struct {
	username string
}

func SetupDatabase() {
	models.Permissions().Create()
	models.Users().Create("admin", "admin")
	models.Permissions().Allow("admin", "list_users")
	models.Permissions().Allow("admin", "list_groups")
}

func recordRequest(r *http.Request, session *Session) *httptest.ResponseRecorder {
	w := httptest.NewRecorder()
	if session != nil {
		s := beego.GlobalSessions.SessionStart(w, r)
		s.Set("username", session.username)
	}
	beego.BeeApp.Handlers.ServeHTTP(w, r)
	// beego.Trace("testing", "TestMain", "Code[%d]\n%s", w.Code, w.Body.String())
	return w
}

func getRequest(method string, urlStr string, session *Session) *httptest.ResponseRecorder {
	r, _ := http.NewRequest(method, urlStr, nil)
	return recordRequest(r, session)
}

func postRequest(method string, resourcePath string, data *url.Values, session *Session) *httptest.ResponseRecorder {
	u, _ := url.ParseRequestURI("/")
	u.Path = resourcePath
	urlStr := fmt.Sprintf("%v", u)

	r, _ := http.NewRequest(method, urlStr, bytes.NewBufferString(data.Encode()))
	r.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	r.Header.Add("Content-Length", strconv.Itoa(len(data.Encode())))
	return recordRequest(r, session)
}

// TestMain is a sample to run an endpoint test
func TestMain(t *testing.T) {
	SetupDatabase()
	var nilSession *Session
	adminSession := &Session{username: "admin"}
	enduserSession := &Session{username: "enduser"}

	Convey("Redirect to Sign In\n", t, func() {
		response := getRequest("GET", "/", nilSession)
		Convey("Redirect", func() {
			So(response.Code, ShouldEqual, 302)
			So(response.HeaderMap.Get("Location"), ShouldEqual, "/sign_in")
		})
	})

	Convey("Redirect to Profile when already Signed In\n", t, func() {
		response := getRequest("GET", "/", adminSession)
		Convey("Redirect", func() {
			So(response.Code, ShouldEqual, 200)
			So(response.Body.String(), ShouldContainSubstring, "Home")
		})
	})

	Convey("Shows Sign In\n", t, func() {
		response := getRequest("GET", "/sign_in", nilSession)
		Convey("Redirect", func() {
			So(response.Code, ShouldEqual, 200)
			So(response.Body.String(), ShouldContainSubstring, "Sign In")
		})
	})

	Convey("Successful Sign In\n", t, func() {
		data := url.Values{}
		data.Add("Username", "admin")
		data.Add("Password", "admin")
		response := postRequest("POST", "/sign_in", &data, nilSession)
		Convey("Redirect", func() {
			So(response.Code, ShouldEqual, 302)
			So(response.HeaderMap.Get("Location"), ShouldEqual, "/")
		})
	})

	Convey("Failed Sign In\n", t, func() {
		data := url.Values{}
		data.Add("Username", "adminx")
		data.Add("Password", "barx")
		response := postRequest("POST", "/sign_in", &data, nilSession)
		Convey("Redirect", func() {
			So(response.Code, ShouldEqual, 302)
			So(response.HeaderMap.Get("Location"), ShouldEqual, "/sign_in")
		})
	})

	Convey("Users without Sign In\n", t, func() {
		response := postRequest("GET", "/users", &url.Values{}, nilSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 302)
			So(response.HeaderMap.Get("Location"), ShouldEqual, "/sign_in")
		})
	})

	Convey("Users\n", t, func() {
		response := postRequest("GET", "/users", &url.Values{}, adminSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 200)
		})
	})

	Convey("Create User\n", t, func() {
		data := url.Values{}
		data.Add("Username", "mikes")
		data.Add("Password", "micke")
		response := postRequest("POST", "/users", &data, adminSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 302)
			So(response.HeaderMap.Get("Location"), ShouldStartWith, "/users/")
		})
	})

	Convey("Create User Error\n", t, func() {
		response := postRequest("POST", "/users", &url.Values{}, adminSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 200)
			So(response.Body.String(), ShouldContainSubstring, "Can not be empty")
		})
	})

	Convey("Create Group\n", t, func() {
		data := url.Values{}
		data.Add("Name", "foo")
		response := postRequest("POST", "/groups", &data, adminSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 302)
			So(response.HeaderMap.Get("Location"), ShouldStartWith, "/groups")
		})
	})

	Convey("Create User Error\n", t, func() {
		response := postRequest("POST", "/groups", &url.Values{}, adminSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 200)
			So(response.Body.String(), ShouldContainSubstring, "Can not be empty")
		})
	})

	Convey("Create User Error\n", t, func() {
		response := postRequest("POST", "/groups", &url.Values{}, enduserSession)
		Convey("Render", func() {
			So(response.Code, ShouldEqual, 401)
			So(response.Body.String(), ShouldContainSubstring, "Unauthorized")
		})
	})
}
<|endoftext|>


<commit_before>// Package broker is an interface used for asynchronous messaging
package broker

// Broker is an interface used for asynchronous messaging.
type Broker interface {
	Options() Options
	Address() string
	Connect() error
	Disconnect() error
	Init(...Option) error
	Publish(string, *Message, ...PublishOption) error
	Subscribe(string, Handler, ...SubscribeOption) (Subscriber, error)
	String() string
}

// Handler is used to process messages via a subscription of a topic.
// The handler is passed a publication interface which contains the
// message and optional Ack method to acknowledge receipt of the message.
type Handler func(Publication) error

type Message struct {
	Header map[string]string
	Body   []byte
}

// Publication is given to a subscription handler for processing
type Publication interface {
	Topic() string
	Message() *Message
	Ack() error
}

// Subscriber is a convenience return type for the Subscribe method
type Subscriber interface {
	Options() SubscribeOptions
	Topic() string
	Unsubscribe() error
}

var (
	DefaultBroker Broker = newHttpBroker()
)

func NewBroker(opts ...Option) Broker {
	return newHttpBroker(opts...)
}

func Init(opts ...Option) error {
	return DefaultBroker.Init(opts...)
}

func Connect() error {
	return DefaultBroker.Connect()
}

func Disconnect() error {
	return DefaultBroker.Disconnect()
}

func Publish(topic string, msg *Message, opts ...PublishOption) error {
	return DefaultBroker.Publish(topic, msg, opts...)
}

func Subscribe(topic string, handler Handler, opts ...SubscribeOption) (Subscriber, error) {
	return DefaultBroker.Subscribe(topic, handler, opts...)
}

func String() string {
	return DefaultBroker.String()
}
<commit_msg>remove broker default funcs<commit_after>// Package broker is an interface used for asynchronous messaging
package broker

// Broker is an interface used for asynchronous messaging.
type Broker interface {
	Options() Options
	Address() string
	Connect() error
	Disconnect() error
	Init(...Option) error
	Publish(string, *Message, ...PublishOption) error
	Subscribe(string, Handler, ...SubscribeOption) (Subscriber, error)
	String() string
}

// Handler is used to process messages via a subscription of a topic.
// The handler is passed a publication interface which contains the
// message and optional Ack method to acknowledge receipt of the message.
type Handler func(Publication) error

type Message struct {
	Header map[string]string
	Body   []byte
}

// Publication is given to a subscription handler for processing
type Publication interface {
	Topic() string
	Message() *Message
	Ack() error
}

// Subscriber is a convenience return type for the Subscribe method
type Subscriber interface {
	Options() SubscribeOptions
	Topic() string
	Unsubscribe() error
}

var (
	DefaultBroker Broker = newHttpBroker()
)

func NewBroker(opts ...Option) Broker {
	return newHttpBroker(opts...)
}
<|endoftext|>


<commit_before>package broker

import (
	"crypto/tls"
	"github.com/t4i/indismqgo"
	"github.com/t4i/indismqgo/connections"
	"github.com/t4i/indismqgo/handlers"
	"github.com/t4i/indismqgo/messages"
	"github.com/t4i/indismqgo/queue"
	"github.com/t4i/indismqgo/subscribers"
	"net/http"
)

//Server ...
type Broker struct {
	name            []byte
	debug           bool
	GetCertificate  func(hello *tls.ClientHelloInfo) (*tls.Certificate, error)
	HTTPAuthHandler func(h *http.Request) (ok bool, claims map[string]string)
	WsAuthHandler   func(h *http.Request) (ok bool, claims map[string]string)
	*queue.Queues
	*messages.Messages
	*handlers.Handlers
	*connections.Connections
	*subscribers.Subscribers
	//indismqgo.ContextImpl
}

// func (s *Server) NewConnectionType(name string, defaultConnection ConnectionFunc, authHandler AuthHandlerFunc) {
// 	s.ConnectionTypes[name] = &ConnectionType{Name: name, DefaultConnection: defaultConnection, AuthHandler: authHandler}
// }
var _ indismqgo.QueueStore = (*Broker)(nil)
var _ indismqgo.MessageStore = (*Broker)(nil)
var _ indismqgo.HandlerStore = (*Broker)(nil)
var _ indismqgo.ConnectionStore = (*Broker)(nil)
var _ indismqgo.Subscribers = (*Broker)(nil)

//NewServer ...
func NewBroker(name string) *Broker {
	broker := &Broker{}
	broker.Name([]byte(name))
	broker.Queues = queue.NewQueues()
	broker.Messages = messages.NewMessages()
	broker.Handlers = handlers.NewHandlers()
	broker.Connections = connections.NewConnections()
	broker.Subscribers = subscribers.NewSubscribers()
	return broker
}
func (c *Broker) Name(set []byte) []byte {
	if set != nil {
		c.name = set
	}
	return c.name
}
func (c *Broker) Debug(set *bool) bool {
	if set != nil {
		c.debug = *set
	}
	return c.debug
}

func (c *Broker) Recieve(m *indismqgo.MsgBuffer, conn indismqgo.Connection) error {
	return indismqgo.Process(c, m, conn)
}

func (c *Broker) RecieveRaw(data []byte, conn indismqgo.Connection) error {
	return indismqgo.ProcessRaw(c, data, conn)
}
func (b *Broker) ProcessNext(q *queue.Queue, conn indismqgo.Connection) {
	// stat := q.Stats()
	// if b.Debug(nil) {
	// 	log.Println(string(b.Name(nil)), q, "queue len", stat.Queue, "reply len", stat.Reply, "dequeue len", stat.Dequeue, "prefetch", stat.Prefetch)
	// }

	// for (q.prefetch > 0 && len(q.dequeue) < q.prefetch) || len(q.queue) > 0 {

	// 	q.processLock.Lock()
	// 	defer q.processLock.Unlock()
	// 	next := q.Peek()
	// 	if q.Handler != nil {

	// 		m := q.Pop()
	// 		q.Handler(m.msg, m.conn)
	// 		return
	// 	}
	// 	//check if is a response
	// 	if reply, replyQueue := queues.GetReply(next.msg); reply != nil {
	// 		m := q.Pop()
	// 		if reply.conn != nil {
	// 			if err := reply.conn.Send(m.msg); err != nil {
	// 				q.RePut(string(m.msg.Fields.MsgId()))
	// 			} else if m.msg.Fields.Guarantee() == indismqgo.GuaranteeAT_LEAST_ONCE {
	// 				//ctx.MakeReply(reply,StatusA)
	// 				delete(replyQueue.reply, string(m.msg.Fields.MsgId()))
	// 				q.RePut(string(m.msg.Fields.MsgId()))
	// 			}

	// 		}
	// 	}

	// 	//check for connection
	// 	// if conn2 := ctx.Connections.Get(string(next.msg.Fields.To())); conn2 != nil {
	// 	// 	m := q.Pop()
	// 	// 	conn2.Send(m.msg)
	// 	// 	ctx.ProcessNext(q, conn)
	// 	// 	return
	// 	// }
	// 	// //check for subscribers
	// 	// subs := ctx.Subscribers.Get(string(next.msg.Fields.To()))
	// 	// for k, v := range subs {
	// 	// 	if v {
	// 	// 		if conn2 := ctx.Connections.Get(k); conn != nil {
	// 	// 			conn2.Send(q.Pop().msg)
	// 	// 			return
	// 	// 		}
	// 	// 	}
	// 	// }

	// }
}

func (b *Broker) OnConnection(m *indismqgo.MsgBuffer, conn indismqgo.Connection) (ok bool) {
	//b.Connections.SetConnection(string(m.Fields.From()), conn)
	return true
}

func (b *Broker) OnDisconnected(m *indismqgo.MsgBuffer, conn indismqgo.Connection) (ok bool) {
	b.Connections.DelConnection(string(m.Fields.From()))
	b.Subscribers.DelSubscriberAll(string(m.Fields.From()))
	return true
}

func (b *Broker) OnSubscribe(m *indismqgo.MsgBuffer, conn indismqgo.Connection) (ok bool) {
	b.Subscribers.AddSubscriber(string(m.Fields.From()), string(m.Fields.Path()))
	return true
}
func (b *Broker) OnUnSubscribe(m *indismqgo.MsgBuffer, conn indismqgo.Connection) {
	b.Subscribers.DelSubscriber(string(m.Fields.From()), string(m.Fields.Path()))
}
func (b *Broker) OnUnknown(m *indismqgo.MsgBuffer, conn indismqgo.Connection) (ok bool) {

	return true
}
<commit_msg>Add Context Interface check<commit_after>package broker

import (
	"crypto/tls"
	"net/http"

	"github.com/t4i/indismqgo"
	"github.com/t4i/indismqgo/connections"
	"github.com/t4i/indismqgo/handlers"
	"github.com/t4i/indismqgo/messages"
	"github.com/t4i/indismqgo/queue"
	"github.com/t4i/indismqgo/subscribers"
)

//Server ...
type Broker struct {
	name            []byte
	debug           bool
	GetCertificate  func(hello *tls.ClientHelloInfo) (*tls.Certificate, error)
	HTTPAuthHandler func(h *http.Request) (ok bool, claims map[string]string)
	WsAuthHandler   func(h *http.Request) (ok bool, claims map[string]string)
	*queue.Queues
	*messages.Messages
	*handlers.Handlers
	*connections.Connections
	*subscribers.Subscribers
	//indismqgo.ContextImpl
}

// func (s *Server) NewConnectionType(name string, defaultConnection ConnectionFunc, authHandler AuthHandlerFunc) {
// 	s.ConnectionTypes[name] = &ConnectionType{Name: name, DefaultConnection: defaultConnection, AuthHandler: authHandler}
// }
var _ indismqgo.QueueStore = (*Broker)(nil)
var _ indismqgo.MessageStore = (*Broker)(nil)
var _ indismqgo.HandlerStore = (*Broker)(nil)
var _ indismqgo.ConnectionStore = (*Broker)(nil)
var _ indismqgo.Subscribers = (*Broker)(nil)
var _ indismqgo.Context = (*Broker)(nil)

//NewServer ...
func NewBroker(name string) *Broker {
	broker := &Broker{}
	broker.Name([]byte(name))
	broker.Queues = queue.NewQueues()
	broker.Messages = messages.NewMessages()
	broker.Handlers = handlers.NewHandlers()
	broker.Connections = connections.NewConnections()
	broker.Subscribers = subscribers.NewSubscribers()
	return broker
}
func (c *Broker) Name(set []byte) []byte {
	if set != nil {
		c.name = set
	}
	return c.name
}
func (c *Broker) Debug(set *bool) bool {
	if set != nil {
		c.debug = *set
	}
	return c.debug
}

func (c *Broker) Recieve(m *indismqgo.MsgBuffer, conn indismqgo.Connection) error {
	return indismqgo.Process(c, m, conn)
}

func (c *Broker) RecieveRaw(data []byte, conn indismqgo.Connection) error {
	return indismqgo.ProcessRaw(c, data, conn)
}
func (b *Broker) ProcessNext(q *queue.Queue, conn indismqgo.Connection) {
	// stat := q.Stats()
	// if b.Debug(nil) {
	// 	log.Println(string(b.Name(nil)), q, "queue len", stat.Queue, "reply len", stat.Reply, "dequeue len", stat.Dequeue, "prefetch", stat.Prefetch)
	// }

	// for (q.prefetch > 0 && len(q.dequeue) < q.prefetch) || len(q.queue) > 0 {

	// 	q.processLock.Lock()
	// 	defer q.processLock.Unlock()
	// 	next := q.Peek()
	// 	if q.Handler != nil {

	// 		m := q.Pop()
	// 		q.Handler(m.msg, m.conn)
	// 		return
	// 	}
	// 	//check if is a response
	// 	if reply, replyQueue := queues.GetReply(next.msg); reply != nil {
	// 		m := q.Pop()
	// 		if reply.conn != nil {
	// 			if err := reply.conn.Send(m.msg); err != nil {
	// 				q.RePut(string(m.msg.Fields.MsgId()))
	// 			} else if m.msg.Fields.Guarantee() == indismqgo.GuaranteeAT_LEAST_ONCE {
	// 				//ctx.MakeReply(reply,StatusA)
	// 				delete(replyQueue.reply, string(m.msg.Fields.MsgId()))
	// 				q.RePut(string(m.msg.Fields.MsgId()))
	// 			}

	// 		}
	// 	}

	// 	//check for connection
	// 	// if conn2 := ctx.Connections.Get(string(next.msg.Fields.To())); conn2 != nil {
	// 	// 	m := q.Pop()
	// 	// 	conn2.Send(m.msg)
	// 	// 	ctx.ProcessNext(q, conn)
	// 	// 	return
	// 	// }
	// 	// //check for subscribers
	// 	// subs := ctx.Subscribers.Get(string(next.msg.Fields.To()))
	// 	// for k, v := range subs {
	// 	// 	if v {
	// 	// 		if conn2 := ctx.Connections.Get(k); conn != nil {
	// 	// 			conn2.Send(q.Pop().msg)
	// 	// 			return
	// 	// 		}
	// 	// 	}
	// 	// }

	// }
}

func (b *Broker) OnConnection(m *indismqgo.MsgBuffer, conn indismqgo.Connection) (ok bool) {
	//b.Connections.SetConnection(string(m.Fields.From()), conn)
	return true
}

func (b *Broker) OnDisconnected(m *indismqgo.MsgBuffer, conn indismqgo.Connection) (ok bool) {
	b.Connections.DelConnection(string(m.Fields.From()))
	b.Subscribers.DelSubscriberAll(string(m.Fields.From()))
	return true
}

func (b *Broker) OnSubscribe(m *indismqgo.MsgBuffer, conn indismqgo.Connection) (ok bool) {
	b.Subscribers.AddSubscriber(string(m.Fields.From()), string(m.Fields.Path()))
	return true
}
func (b *Broker) OnUnSubscribe(m *indismqgo.MsgBuffer, conn indismqgo.Connection) {
	b.Subscribers.DelSubscriber(string(m.Fields.From()), string(m.Fields.Path()))
}
func (b *Broker) OnUnknown(m *indismqgo.MsgBuffer, conn indismqgo.Connection) (ok bool) {

	return true
}
<|endoftext|>


<commit_before>package libsysinfo

import (
	"os/exec"
	"strings"
)

const ()

var (
	cacheKeys = map[string]string{
		"HOSTNAME_FULL":        "fullhostname",
		"HOSTNAME":             "hostname",
		"DOMAIN_NAME":          "domainname",
		"LSB_FULL":             "lsbfull",
		"LSB_DIST_CODE_NAME":   "lsbdistcodename",
		"LSB_DIST_DESCRIPTION": "lsbdistdescrption",
		"LSB_DIST_ID":          "lsbdistid",
		"LSB_DIST_RELEASE":     "lsbdistrelease",
	}

	globalCache = newCachedValues(len(cacheKeys))
)

func Hostname() (string, error) {
	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys["HOSTNAME"],
		Fetcher:     getFullHostname,
		Processor:   processHostname,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func Domain() (string, error) {
	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys["DOMAIN_NAME"],
		Fetcher:     getFullHostname,
		Processor:   processDomainName,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func Fqdn() (string, error) {
	fqdn := func(fullHostname string) (string, error) {
		return fullHostname, nil
	}

	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys["FQDN"],
		Fetcher:     getFullHostname,
		Processor:   fqdn,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func LsbDistCodeName() (string, error) {
	return lsbDist("LSB_DIST_CODE_NAME", "Codename")
}

func LsbDistDescription() (string, error) {
	return lsbDist("LSB_DIST_DESCRIPTION", "Description")
}

func LsbDistId() (string, error) {
	return lsbDist("LSB_DIST_ID", "Distributor ID")
}

func LsbDistRelease() (string, error) {
	return lsbDist("LSB_DIST_RELEASE", "Release")
}

func lsbDist(k string, lsbItem string) (string, error) {
	proc := func(lsb string) (string, error) {
		return processLsbItem(lsb, lsbItem)
	}

	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys[k],
		Fetcher:     getLsbRelease,
		Processor:   proc,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func processDomainName(fullHostname string) (string, error) {
	pos := strings.Index(fullHostname, ".")
	if pos == -1 {
		return "", nil
	}

	return fullHostname[pos+1:], nil
}

func processHostname(fullHostname string) (string, error) {
	pos := strings.Index(fullHostname, ".")
	if pos == -1 {
		return fullHostname, nil
	}

	return fullHostname[:pos], nil
}

func processLsbItem(lsb string, item string) (string, error) {
	var out string
	var tmp string

	for _, line := range strings.Split(lsb, "\n") {
		if len(line) <= 0 {
			continue
		}

		tmp = line[0:len(item)]
		if tmp == item {
			out = strings.TrimSpace(strings.TrimLeft(line, item+":"))
			break
		}
	}

	return strings.ToLower(out), nil
}

func getFullHostname() (string, error) {
	cacheKey := cacheKeys["HOSTNAME_FULL"]

	hf, exists := globalCache.Exists(cacheKey)
	if exists {
		return hf, nil
	}

	out, err := exec.Command("hostname", "-f").Output()
	if err != nil {
		return string(out), err
	}

	hf = string(out)

	// removing \n
	pos := strings.Index(hf, "\n")
	if pos > -1 {
		hf = hf[:pos]
	}

	globalCache.Set(cacheKey, hf)

	return hf, nil
}

func getLsbRelease() (string, error) {
	cacheKey := cacheKeys["LSB_FULL"]

	lsb, exists := globalCache.Exists(cacheKey)
	if exists {
		return lsb, nil
	}

	out, err := exec.Command("lsb_release", "-a").Output()
	if err != nil {
		return string(out), err
	}

	lsb = string(out)

	globalCache.Set(cacheKey, lsb)

	return lsb, nil
}
<commit_msg>Added build flag<commit_after>// +build linux
package libsysinfo

import (
	"os/exec"
	"strings"
)

const ()

var (
	cacheKeys = map[string]string{
		"HOSTNAME_FULL":        "fullhostname",
		"HOSTNAME":             "hostname",
		"DOMAIN_NAME":          "domainname",
		"LSB_FULL":             "lsbfull",
		"LSB_DIST_CODE_NAME":   "lsbdistcodename",
		"LSB_DIST_DESCRIPTION": "lsbdistdescrption",
		"LSB_DIST_ID":          "lsbdistid",
		"LSB_DIST_RELEASE":     "lsbdistrelease",
	}

	globalCache = newCachedValues(len(cacheKeys))
)

func Hostname() (string, error) {
	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys["HOSTNAME"],
		Fetcher:     getFullHostname,
		Processor:   processHostname,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func Domain() (string, error) {
	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys["DOMAIN_NAME"],
		Fetcher:     getFullHostname,
		Processor:   processDomainName,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func Fqdn() (string, error) {
	fqdn := func(fullHostname string) (string, error) {
		return fullHostname, nil
	}

	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys["FQDN"],
		Fetcher:     getFullHostname,
		Processor:   fqdn,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func LsbDistCodeName() (string, error) {
	return lsbDist("LSB_DIST_CODE_NAME", "Codename")
}

func LsbDistDescription() (string, error) {
	return lsbDist("LSB_DIST_DESCRIPTION", "Description")
}

func LsbDistId() (string, error) {
	return lsbDist("LSB_DIST_ID", "Distributor ID")
}

func LsbDistRelease() (string, error) {
	return lsbDist("LSB_DIST_RELEASE", "Release")
}

func lsbDist(k string, lsbItem string) (string, error) {
	proc := func(lsb string) (string, error) {
		return processLsbItem(lsb, lsbItem)
	}

	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys[k],
		Fetcher:     getLsbRelease,
		Processor:   proc,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func processDomainName(fullHostname string) (string, error) {
	pos := strings.Index(fullHostname, ".")
	if pos == -1 {
		return "", nil
	}

	return fullHostname[pos+1:], nil
}

func processHostname(fullHostname string) (string, error) {
	pos := strings.Index(fullHostname, ".")
	if pos == -1 {
		return fullHostname, nil
	}

	return fullHostname[:pos], nil
}

func processLsbItem(lsb string, item string) (string, error) {
	var out string
	var tmp string

	for _, line := range strings.Split(lsb, "\n") {
		if len(line) <= 0 {
			continue
		}

		tmp = line[0:len(item)]
		if tmp == item {
			out = strings.TrimSpace(strings.TrimLeft(line, item+":"))
			break
		}
	}

	return strings.ToLower(out), nil
}

func getFullHostname() (string, error) {
	cacheKey := cacheKeys["HOSTNAME_FULL"]

	hf, exists := globalCache.Exists(cacheKey)
	if exists {
		return hf, nil
	}

	out, err := exec.Command("hostname", "-f").Output()
	if err != nil {
		return string(out), err
	}

	hf = string(out)

	// removing \n
	pos := strings.Index(hf, "\n")
	if pos > -1 {
		hf = hf[:pos]
	}

	globalCache.Set(cacheKey, hf)

	return hf, nil
}

func getLsbRelease() (string, error) {
	cacheKey := cacheKeys["LSB_FULL"]

	lsb, exists := globalCache.Exists(cacheKey)
	if exists {
		return lsb, nil
	}

	out, err := exec.Command("lsb_release", "-a").Output()
	if err != nil {
		return string(out), err
	}

	lsb = string(out)

	globalCache.Set(cacheKey, lsb)

	return lsb, nil
}
<|endoftext|>


<commit_before>package libsysinfo

import (
	"os/exec"
	"strings"
)

const ()

var (
	cacheKeys = map[string]string{
		"HOSTNAME_FULL": "fullhostname",
		"HOSTNAME":      "hostname",
		"DOMAIN_NAME":   "domainname",
	}

	globalCache = NewCachedValues(len(cacheKeys))
)

func Hostname() (string, error) {
	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys["HOSTNAME"],
		Fetcher:     getFullHostname,
		Processor:   processHostname,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func Domain() (string, error) {
	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys["DOMAIN_NAME"],
		Fetcher:     getFullHostname,
		Processor:   processDomainName,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func Fqdn() (string, error) {
	fqdn := func(fullHostname string) (string, error) {
		return fullHostname, nil
	}

	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys["FQDN"],
		Fetcher:     getFullHostname,
		Processor:   fqdn,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func processDomainName(fullHostname string) (string, error) {
	pos := strings.Index(fullHostname, ".")
	if pos == -1 {
		return "", nil
	}

	return fullHostname[pos+1:], nil
}

func processHostname(fullHostname string) (string, error) {
	pos := strings.Index(fullHostname, ".")
	if pos == -1 {
		return fullHostname, nil
	}

	return fullHostname[:pos], nil
}

func getFullHostname() (string, error) {
	cacheKey := cacheKeys["HOSTNAME_FULL"]

	hf, exists := globalCache.Exists(cacheKey)
	if exists {
		return hf, nil
	}

	out, err := exec.Command("hostname", "-f").Output()
	if err != nil {
		return string(out), err
	}

	hf = string(out)

	// removing \n
	pos := strings.Index(hf, "\n")
	if pos > -1 {
		hf = hf[:pos]
	}

	globalCache.Set(cacheKey, hf)

	return hf, nil
}
<commit_msg>Added LsbDistCodeName()<commit_after>package libsysinfo

import (
	"os/exec"
	"strings"
)

const ()

var (
	cacheKeys = map[string]string{
		"HOSTNAME_FULL":      "fullhostname",
		"HOSTNAME":           "hostname",
		"DOMAIN_NAME":        "domainname",
		"LSB_FULL":           "lsbfull",
		"LSB_DIST_CODE_NAME": "lsbdistcodename",
	}

	globalCache = NewCachedValues(len(cacheKeys))
)

func Hostname() (string, error) {
	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys["HOSTNAME"],
		Fetcher:     getFullHostname,
		Processor:   processHostname,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func Domain() (string, error) {
	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys["DOMAIN_NAME"],
		Fetcher:     getFullHostname,
		Processor:   processDomainName,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func Fqdn() (string, error) {
	fqdn := func(fullHostname string) (string, error) {
		return fullHostname, nil
	}

	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys["FQDN"],
		Fetcher:     getFullHostname,
		Processor:   fqdn,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func LsbDistCodeName() (string, error) {
	return lsbDist("LSB_DIST_CODE_NAME", processLsbRelease)
}

func lsbDist(k string, procFunc processor) (string, error) {
	llv := &lazyLoadedValue{
		CacheKey:    cacheKeys[k],
		Fetcher:     getLsbRelease,
		Processor:   procFunc,
		CacheBucket: globalCache,
	}

	return llv.run()
}

func processDomainName(fullHostname string) (string, error) {
	pos := strings.Index(fullHostname, ".")
	if pos == -1 {
		return "", nil
	}

	return fullHostname[pos+1:], nil
}

func processHostname(fullHostname string) (string, error) {
	pos := strings.Index(fullHostname, ".")
	if pos == -1 {
		return fullHostname, nil
	}

	return fullHostname[:pos], nil
}

func processLsbRelease(lsb string) (string, error) {
	var out string
	for _, line := range strings.Split(lsb, "\n") {
		if len(line) <= 0 {
			continue
		}

		// C => Codename
		if line[0] != 'C' {
			continue
		}

		out = strings.TrimSpace(strings.TrimLeft(line, "Codename:"))

		break
	}

	return strings.ToLower(out), nil
}

func getFullHostname() (string, error) {
	cacheKey := cacheKeys["HOSTNAME_FULL"]

	hf, exists := globalCache.Exists(cacheKey)
	if exists {
		return hf, nil
	}

	out, err := exec.Command("hostname", "-f").Output()
	if err != nil {
		return string(out), err
	}

	hf = string(out)

	// removing \n
	pos := strings.Index(hf, "\n")
	if pos > -1 {
		hf = hf[:pos]
	}

	globalCache.Set(cacheKey, hf)

	return hf, nil
}

func getLsbRelease() (string, error) {
	cacheKey := cacheKeys["LSB_FULL"]

	lsb, exists := globalCache.Exists(cacheKey)
	if exists {
		return lsb, nil
	}

	out, err := exec.Command("lsb_release", "-a").Output()
	if err != nil {
		return string(out), err
	}

	lsb = string(out)

	globalCache.Set(cacheKey, lsb)

	return lsb, nil
}
<|endoftext|>


<commit_before>/*
 * nfqTrace.go - Forward/Reverse TCP traceroute using Linux NFQueue
 * Copyright (c) 2014 David Anthony Stainton
 *
 * Abstract:
 * Parasitic forward/reverse Linux-Netfilter Queue traceroute
 * for penetrating network address translation devices...
 * tracing the client's internal network.
 *
 * The MIT License (MIT)
 * Copyright (c) 2014 David Anthony Stainton
 * 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 main

import (
	"code.google.com/p/gopacket"
	"code.google.com/p/gopacket/layers"
	"code.google.com/p/gopacket/pcap"
	"encoding/binary"
	"github.com/david415/go-netfilter-queue"
	"log"
	"net"
	"sync"
	"time"
)

const (
	MAX_TTL uint8 = 255
)

// this is a composite struct type called "flowKey"
// used to track tcp/ip flows... as a hashmap key.
type flowKey [2]gopacket.Flow

// concurrent-safe hashmap of tcp/ip-flowKeys to NFQueueTraceroute`s
type FlowTracker struct {
	lock    *sync.RWMutex
	flowMap map[flowKey]*NFQueueTraceroute
}

func NewFlowTracker() *FlowTracker {
	return &FlowTracker{
		lock:    new(sync.RWMutex),
		flowMap: make(map[flowKey]*NFQueueTraceroute),
	}
}

func (f *FlowTracker) HasFlow(flow flowKey) bool {
	f.lock.RLock()
	_, ok := f.flowMap[flow]
	f.lock.RUnlock()
	return ok
}

func (f *FlowTracker) AddFlow(flow flowKey, nfqTrace *NFQueueTraceroute) {
	f.lock.Lock()
	f.flowMap[flow] = nfqTrace
	f.lock.Unlock()
}

func (f *FlowTracker) Delete(flow flowKey) {
	f.lock.Lock()
	delete(f.flowMap, flow)
	f.lock.Unlock()
}

func (f *FlowTracker) GetFlowTrace(flow flowKey) *NFQueueTraceroute {
	f.lock.RLock()
	ret := f.flowMap[flow]
	f.lock.RUnlock()
	return ret
}

type NFQueueTraceObserverOptions struct {
	// network interface to listen for ICMP responses
	iface        string
	ttlMax       uint8
	ttlRepeatMax int
	mangleFreq   int
}

type NFQueueTraceObserver struct {
	// passed in from the user in our constructor...
	options NFQueueTraceObserverOptions

	flowTracker *FlowTracker
	nfq         *netfilter.NFQueue

	// packet channel for interacting with NFQueue
	packets <-chan netfilter.NFPacket

	// this is used to stop all the traceroutes
	done chan bool

	// signal our calling party that we are finished
	// XXX get rid of this?
	finished chan bool
}

func NewNFQueueTraceObserver(options NFQueueTraceObserverOptions) *NFQueueTraceObserver {
	var err error
	o := NFQueueTraceObserver{
		options:  options,
		done:     make(chan bool),
		finished: make(chan bool),
	}

	flowTracker := NewFlowTracker()
	o.flowTracker = flowTracker
	// XXX adjust these parameters
	o.nfq, err = netfilter.NewNFQueue(0, 100, netfilter.NF_DEFAULT_PACKET_SIZE)
	if err != nil {
		panic(err)
	}
	o.packets = o.nfq.GetPackets()
	return &o
}

func (o *NFQueueTraceObserver) Start() {
	o.startReceivingReplies()
	go func() {
		for true {
			select {
			case p := <-o.packets:
				o.processPacket(p)
			case <-o.done:
				o.nfq.Close()
				close(o.done) // XXX necessary?
				// XXX todo: stop other goroutines
				break
			}
		}
	}()
}

func (o *NFQueueTraceObserver) Stop() {
	o.done <- true
}

// XXX make the locking more efficient?
func (o *NFQueueTraceObserver) processPacket(p netfilter.NFPacket) {
	ipLayer := p.Packet.Layer(layers.LayerTypeIPv4)
	tcpLayer := p.Packet.Layer(layers.LayerTypeTCP)
	if ipLayer == nil || tcpLayer == nil {
		// ignore non-tcp/ip packets
		return
	}
	ip, _ := ipLayer.(*layers.IPv4)
	tcp, _ := tcpLayer.(*layers.TCP)

	flow := flowKey{ip.NetworkFlow(), tcp.TransportFlow()}
	if o.flowTracker.HasFlow(flow) == false {
		nfqTrace := NewNFQueueTraceroute(o.options.ttlMax, o.options.ttlRepeatMax, o.options.mangleFreq)
		o.flowTracker.AddFlow(flow, nfqTrace)
	}
	nfqTrace := o.flowTracker.GetFlowTrace(flow)
	nfqTrace.processPacket(p)
}

// return a net.IP channel to report all the ICMP reponse SrcIP addresses
// that have the ICMP time exceeded flag set
func (o *NFQueueTraceObserver) startReceivingReplies() {
	snaplen := 65536
	filter := "icmp" // the idea here is to capture only ICMP packets

	var eth layers.Ethernet
	var ip layers.IPv4
	var icmp layers.ICMPv4
	var payload gopacket.Payload
	var flow flowKey

	decoded := make([]gopacket.LayerType, 0, 4)

	handle, err := pcap.OpenLive(o.options.iface, int32(snaplen), true, pcap.BlockForever)
	if err != nil {
		log.Fatal("error opening pcap handle: ", err)
	}
	if err := handle.SetBPFFilter(filter); err != nil {
		log.Fatal("error setting BPF filter: ", err)
	}

	parser := gopacket.NewDecodingLayerParser(layers.LayerTypeEthernet, &eth, &ip, &icmp, &payload)

	go func() {
		for true {
			data, _, err := handle.ReadPacketData()
			if err != nil {
				continue
			}
			err = parser.DecodeLayers(data, &decoded)
			if err != nil {
				continue
			}
			typ := uint8(icmp.TypeCode >> 8)
			if typ != layers.ICMPv4TypeTimeExceeded {
				continue
			}

			// XXX todo: check that the IP header protocol value is set to TCP
			flow = getPacketFlow(payload)

			// XXX it feels dirty to have the mutex around the hashmap
			// i'm thinking about using channels instead...
			if o.flowTracker.HasFlow(flow) == false {
				// ignore ICMP ttl expire packets that are for flows other than the ones we are currently tracking
				continue
			}

			nfqTrace := o.flowTracker.GetFlowTrace(flow)
			nfqTrace.replyReceived(ip.SrcIP)
		}
	}()
}

type NFQueueTraceroute struct {
	ttl          uint8
	ttlMax       uint8
	ttlRepeat    int
	ttlRepeatMax int
	mangleFreq   int
	count        int

	// ip.TTL -> list of ip addrs
	traceResult map[uint8][]net.IP

	stopped          bool
	responseTimedOut bool

	// XXX should it be a pointer instead?
	receivePacketChannel chan netfilter.NFPacket

	resumeTimerChannel  chan bool
	stopTimerChannel    chan bool
	restartTimerChannel chan bool
}

// conduct an nfqueue tcp traceroute;
// - send each TTL out ttlRepeatMax number of times.
// - only mangle a packet's TTL after mangleFreq number
// of packets have traversed the flow
func NewNFQueueTraceroute(ttlMax uint8, ttlRepeatMax, mangleFreq int) *NFQueueTraceroute {
	log.Print("NewNFQueueTraceroute\n")
	nfqTrace := NFQueueTraceroute{
		ttl:                 1,
		ttlMax:              ttlMax,
		ttlRepeat:           1,
		ttlRepeatMax:        ttlRepeatMax,
		mangleFreq:          mangleFreq,
		count:               1,
		traceResult:         make(map[uint8][]net.IP, 1),
		stopped:             false,
		responseTimedOut:    false,
		stopTimerChannel:    make(chan bool),
		restartTimerChannel: make(chan bool),
	}
	nfqTrace.StartResponseTimer()
	return &nfqTrace
}

func (n *NFQueueTraceroute) StartResponseTimer() {
	log.Print("StartResponseTimer\n")

	go func() {
		for {
			select {
			case <-time.After(time.Duration(200) * time.Second):
				log.Print("TimerExpired\n")

				if n.ttl >= n.ttlMax && n.ttlRepeat >= n.ttlRepeatMax {
					n.Stop()
					return
				}

				n.responseTimedOut = true
			case <-n.restartTimerChannel:
				log.Print("restartTimerChannel fired\n")
				continue
			case <-n.stopTimerChannel:
				log.Print("stopTimerChannel fired\n")
				return
			}
		}
	}()
}

func (n *NFQueueTraceroute) Stop() {
	log.Print("NFQueueTraceroute.Stop()\n")
	n.stopped = true
	n.stopTimerChannel <- true
	close(n.stopTimerChannel)
	close(n.restartTimerChannel)

}

// given a packet we decided weather or not to mangle the TTL
// for our tracerouting purposes.
func (n *NFQueueTraceroute) processPacket(p netfilter.NFPacket) {

	if n.stopped {
		p.SetVerdict(netfilter.NF_ACCEPT)
		return
	}

	if n.count%n.mangleFreq == 0 {
		log.Printf("processPacket mangle case ttl %d\n", n.ttl)

		n.ttlRepeat += 1

		if n.responseTimedOut {
			n.ttl += 1
			n.ttlRepeat = 0
			n.restartTimerChannel <- true
		} else if n.ttlRepeat >= n.ttlRepeatMax {
			log.Print("ttlRepeatMax reached case\n")
			n.ttl += 1
			n.ttlRepeat = 0
			n.restartTimerChannel <- true
		}

		p.SetModifiedVerdict(netfilter.NF_REPEAT, serializeWithTTL(p.Packet, n.ttl))
	} else {
		p.SetVerdict(netfilter.NF_ACCEPT)
	}
	n.count = n.count + 1
}

// XXX
// store the "reply" source ip address (icmp ttl expired packet with payload matching this flow)
func (n *NFQueueTraceroute) replyReceived(ip net.IP) {
	log.Printf("replyReceived: ttl %d ip %s\n", n.ttl, ip.String())

	n.traceResult[n.ttl] = append(n.traceResult[n.ttl], ip)
	if n.ttl == n.ttlMax && len(n.traceResult[n.ttl]) >= n.ttlRepeatMax {
		n.Stop() // finished!
	}
}

// This function takes a gopacket.Packet and a TTL
// and returns a byte array of the serialized packet with the specified TTL
func serializeWithTTL(p gopacket.Packet, ttl uint8) []byte {
	ipLayer := p.Layer(layers.LayerTypeIPv4)
	if ipLayer == nil {
		return nil
	}
	tcpLayer := p.Layer(layers.LayerTypeTCP)
	if tcpLayer == nil {
		return nil
	}
	ip, _ := ipLayer.(*layers.IPv4)
	ip.TTL = ttl
	tcp, _ := tcpLayer.(*layers.TCP)
	opts := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	tcp.SetNetworkLayerForChecksum(ip)
	rawPacketBuf := gopacket.NewSerializeBuffer()
	if err := gopacket.SerializeLayers(rawPacketBuf, opts, ip, tcp); err != nil {
		return nil
	}
	return rawPacketBuf.Bytes()
}

// We use this to deal with rfc792 implementations where
// the original packet is NOT sent back via ICMP payload but
// instead 64 bits of the original packet are sent.
// https://tools.ietf.org/html/rfc792
// Returns a TCP Flow.
// XXX obviously the 64 bits could be from a UDP packet or something else
// however this is *good-enough* for NFQueue TCP traceroute!
// XXX should we look at the protocol specified in the IP header
// and set it's type here? no we should probably not even get this
// far if the IP header has something other than TCP specified...
func getTCPFlowFromTCPHead(data []byte) gopacket.Flow {
	var srcPort, dstPort layers.TCPPort
	srcPort = layers.TCPPort(binary.BigEndian.Uint16(data[0:2]))
	dstPort = layers.TCPPort(binary.BigEndian.Uint16(data[2:4]))
	// XXX convert to tcp/ip flow
	tcpSrc := layers.NewTCPPortEndpoint(srcPort)
	tcpDst := layers.NewTCPPortEndpoint(dstPort)
	tcpFlow, _ := gopacket.FlowFromEndpoints(tcpSrc, tcpDst)
	// error (^ _) is only non-nil if the two endpoint types don't match
	return tcpFlow
}

// given a byte array packet return a tcp/ip flow
func getPacketFlow(packet []byte) flowKey {
	var ip layers.IPv4
	var tcp layers.TCP
	decoded := []gopacket.LayerType{}
	parser := gopacket.NewDecodingLayerParser(layers.LayerTypeIPv4, &ip, &tcp)
	err := parser.DecodeLayers(packet, &decoded)
	if err != nil {
		// XXX last 64 bits... we only use the last 32 bits
		tcpHead := packet[len(packet)-8 : len(packet)]
		tcpFlow := getTCPFlowFromTCPHead(tcpHead)
		return flowKey{ip.NetworkFlow(), tcpFlow}
	}
	return flowKey{ip.NetworkFlow(), tcp.TransportFlow()}
}

/***
use this rough POC with an iptables nfqueue rule that will select
a tcp flow direction like this:
iptables -A OUTPUT -j NFQUEUE --queue-num 0 -p tcp --dport 2666

***/
func main() {
	options := NFQueueTraceObserverOptions{
		iface:        "wlan0",
		ttlMax:       40,
		ttlRepeatMax: 3,
		mangleFreq:   6,
	}
	o := NewNFQueueTraceObserver(options)
	o.Start()
	// XXX
	<-o.finished
}
<commit_msg>Fix trace termination bug<commit_after>/*
 * nfqTrace.go - Forward/Reverse TCP traceroute using Linux NFQueue
 * Copyright (c) 2014 David Anthony Stainton
 *
 * Abstract:
 * Parasitic forward/reverse Linux-Netfilter Queue traceroute
 * for penetrating network address translation devices...
 * tracing the client's internal network.
 *
 * The MIT License (MIT)
 * Copyright (c) 2014 David Anthony Stainton
 * 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 main

import (
	"code.google.com/p/gopacket"
	"code.google.com/p/gopacket/layers"
	"code.google.com/p/gopacket/pcap"
	"encoding/binary"
	"github.com/david415/go-netfilter-queue"
	"log"
	"net"
	"sync"
	"time"
)

const (
	MAX_TTL uint8 = 255
)

// this is a composite struct type called "flowKey"
// used to track tcp/ip flows... as a hashmap key.
type flowKey [2]gopacket.Flow

// concurrent-safe hashmap of tcp/ip-flowKeys to NFQueueTraceroute`s
type FlowTracker struct {
	lock    *sync.RWMutex
	flowMap map[flowKey]*NFQueueTraceroute
}

func NewFlowTracker() *FlowTracker {
	return &FlowTracker{
		lock:    new(sync.RWMutex),
		flowMap: make(map[flowKey]*NFQueueTraceroute),
	}
}

func (f *FlowTracker) HasFlow(flow flowKey) bool {
	f.lock.RLock()
	_, ok := f.flowMap[flow]
	f.lock.RUnlock()
	return ok
}

func (f *FlowTracker) AddFlow(flow flowKey, nfqTrace *NFQueueTraceroute) {
	f.lock.Lock()
	f.flowMap[flow] = nfqTrace
	f.lock.Unlock()
}

func (f *FlowTracker) Delete(flow flowKey) {
	f.lock.Lock()
	delete(f.flowMap, flow)
	f.lock.Unlock()
}

func (f *FlowTracker) GetFlowTrace(flow flowKey) *NFQueueTraceroute {
	f.lock.RLock()
	ret := f.flowMap[flow]
	f.lock.RUnlock()
	return ret
}

type NFQueueTraceObserverOptions struct {
	// network interface to listen for ICMP responses
	iface        string
	ttlMax       uint8
	ttlRepeatMax int
	mangleFreq   int
}

type NFQueueTraceObserver struct {
	// passed in from the user in our constructor...
	options NFQueueTraceObserverOptions

	flowTracker *FlowTracker
	nfq         *netfilter.NFQueue

	// packet channel for interacting with NFQueue
	packets <-chan netfilter.NFPacket

	// this is used to stop all the traceroutes
	done chan bool

	// signal our calling party that we are finished
	// XXX get rid of this?
	finished chan bool
}

func NewNFQueueTraceObserver(options NFQueueTraceObserverOptions) *NFQueueTraceObserver {
	var err error
	o := NFQueueTraceObserver{
		options:  options,
		done:     make(chan bool),
		finished: make(chan bool),
	}

	flowTracker := NewFlowTracker()
	o.flowTracker = flowTracker
	// XXX adjust these parameters
	o.nfq, err = netfilter.NewNFQueue(0, 100, netfilter.NF_DEFAULT_PACKET_SIZE)
	if err != nil {
		panic(err)
	}
	o.packets = o.nfq.GetPackets()
	return &o
}

func (o *NFQueueTraceObserver) Start() {
	o.startReceivingReplies()
	go func() {
		for true {
			select {
			case p := <-o.packets:
				o.processPacket(p)
			case <-o.done:
				o.nfq.Close()
				close(o.done) // XXX necessary?
				// XXX todo: stop other goroutines
				break
			}
		}
	}()
}

func (o *NFQueueTraceObserver) Stop() {
	o.done <- true
}

// XXX make the locking more efficient?
func (o *NFQueueTraceObserver) processPacket(p netfilter.NFPacket) {
	ipLayer := p.Packet.Layer(layers.LayerTypeIPv4)
	tcpLayer := p.Packet.Layer(layers.LayerTypeTCP)
	if ipLayer == nil || tcpLayer == nil {
		// ignore non-tcp/ip packets
		return
	}
	ip, _ := ipLayer.(*layers.IPv4)
	tcp, _ := tcpLayer.(*layers.TCP)

	flow := flowKey{ip.NetworkFlow(), tcp.TransportFlow()}
	if o.flowTracker.HasFlow(flow) == false {
		nfqTrace := NewNFQueueTraceroute(o.options.ttlMax, o.options.ttlRepeatMax, o.options.mangleFreq)
		o.flowTracker.AddFlow(flow, nfqTrace)
	}
	nfqTrace := o.flowTracker.GetFlowTrace(flow)
	nfqTrace.processPacket(p)
}

// return a net.IP channel to report all the ICMP reponse SrcIP addresses
// that have the ICMP time exceeded flag set
func (o *NFQueueTraceObserver) startReceivingReplies() {
	snaplen := 65536
	filter := "icmp" // the idea here is to capture only ICMP packets

	var eth layers.Ethernet
	var ip layers.IPv4
	var icmp layers.ICMPv4
	var payload gopacket.Payload
	var flow flowKey

	decoded := make([]gopacket.LayerType, 0, 4)

	handle, err := pcap.OpenLive(o.options.iface, int32(snaplen), true, pcap.BlockForever)
	if err != nil {
		log.Fatal("error opening pcap handle: ", err)
	}
	if err := handle.SetBPFFilter(filter); err != nil {
		log.Fatal("error setting BPF filter: ", err)
	}

	parser := gopacket.NewDecodingLayerParser(layers.LayerTypeEthernet, &eth, &ip, &icmp, &payload)

	go func() {
		for true {
			data, _, err := handle.ReadPacketData()
			if err != nil {
				continue
			}
			err = parser.DecodeLayers(data, &decoded)
			if err != nil {
				continue
			}
			typ := uint8(icmp.TypeCode >> 8)
			if typ != layers.ICMPv4TypeTimeExceeded {
				continue
			}

			// XXX todo: check that the IP header protocol value is set to TCP
			flow = getPacketFlow(payload)

			// XXX it feels dirty to have the mutex around the hashmap
			// i'm thinking about using channels instead...
			if o.flowTracker.HasFlow(flow) == false {
				// ignore ICMP ttl expire packets that are for flows other than the ones we are currently tracking
				continue
			}

			nfqTrace := o.flowTracker.GetFlowTrace(flow)
			nfqTrace.replyReceived(ip.SrcIP)
		}
	}()
}

type NFQueueTraceroute struct {
	ttl          uint8
	ttlMax       uint8
	ttlRepeat    int
	ttlRepeatMax int
	mangleFreq   int
	count        int

	// ip.TTL -> list of ip addrs
	traceResult map[uint8][]net.IP

	stopped          bool
	responseTimedOut bool

	// XXX should it be a pointer instead?
	receivePacketChannel chan netfilter.NFPacket

	resumeTimerChannel  chan bool
	stopTimerChannel    chan bool
	restartTimerChannel chan bool
}

// conduct an nfqueue tcp traceroute;
// - send each TTL out ttlRepeatMax number of times.
// - only mangle a packet's TTL after mangleFreq number
// of packets have traversed the flow
func NewNFQueueTraceroute(ttlMax uint8, ttlRepeatMax, mangleFreq int) *NFQueueTraceroute {
	log.Print("NewNFQueueTraceroute\n")
	nfqTrace := NFQueueTraceroute{
		ttl:                 1,
		ttlMax:              ttlMax,
		ttlRepeat:           1,
		ttlRepeatMax:        ttlRepeatMax,
		mangleFreq:          mangleFreq,
		count:               1,
		traceResult:         make(map[uint8][]net.IP, 1),
		stopped:             false,
		responseTimedOut:    false,
		stopTimerChannel:    make(chan bool),
		restartTimerChannel: make(chan bool),
	}
	nfqTrace.StartResponseTimer()
	return &nfqTrace
}

func (n *NFQueueTraceroute) StartResponseTimer() {
	log.Print("StartResponseTimer\n")

	go func() {
		for {
			select {
			case <-time.After(time.Duration(200) * time.Second):
				log.Print("TimerExpired\n")

				if n.ttl >= n.ttlMax && n.ttlRepeat >= n.ttlRepeatMax {
					n.Stop()
					return
				}

				n.responseTimedOut = true
			case <-n.restartTimerChannel:
				log.Print("restartTimerChannel fired\n")
				continue
			case <-n.stopTimerChannel:
				log.Print("stopTimerChannel fired\n")
				return
			}
		}
	}()
}

func (n *NFQueueTraceroute) Stop() {
	log.Print("NFQueueTraceroute.Stop()\n")
	n.stopped = true
	n.stopTimerChannel <- true
	close(n.stopTimerChannel)
	close(n.restartTimerChannel)

}

// given a packet we decided weather or not to mangle the TTL
// for our tracerouting purposes.
func (n *NFQueueTraceroute) processPacket(p netfilter.NFPacket) {

	if n.stopped {
		p.SetVerdict(netfilter.NF_ACCEPT)
		return
	}

	if n.count%n.mangleFreq == 0 {
		log.Printf("processPacket mangle case n.ttl %d, n.ttlRepeat %d, n.ttlRepeatMax %d\n", n.ttl, n.ttlRepeat, n.ttlRepeatMax)

		n.ttlRepeat += 1

		if n.responseTimedOut {
			n.ttl += 1
			n.ttlRepeat = 0
			n.responseTimedOut = false
			n.restartTimerChannel <- true
		} else if n.ttlRepeat == n.ttlRepeatMax {
			log.Print("ttlRepeatMax reached case\n")
			n.ttl += 1
			n.ttlRepeat = 0
			n.responseTimedOut = false
			n.restartTimerChannel <- true
		}

		// terminate trace upon max ttl and ttlRepeatMax conditions
		if n.ttl > n.ttlMax && n.ttlRepeat == (n.ttlRepeatMax-1) {
			n.Stop()
			p.SetVerdict(netfilter.NF_ACCEPT)
			return
		}

		p.SetModifiedVerdict(netfilter.NF_REPEAT, serializeWithTTL(p.Packet, n.ttl))
	} else {
		p.SetVerdict(netfilter.NF_ACCEPT)
	}
	n.count = n.count + 1
}

// XXX
// store the "reply" source ip address (icmp ttl expired packet with payload matching this flow)
func (n *NFQueueTraceroute) replyReceived(ip net.IP) {
	log.Printf("replyReceived: ttl %d ip %s\n", n.ttl, ip.String())

	n.traceResult[n.ttl] = append(n.traceResult[n.ttl], ip)
	if n.ttl == n.ttlMax && len(n.traceResult[n.ttl]) >= n.ttlRepeatMax {
		n.Stop() // finished!
	}
}

// This function takes a gopacket.Packet and a TTL
// and returns a byte array of the serialized packet with the specified TTL
func serializeWithTTL(p gopacket.Packet, ttl uint8) []byte {
	ipLayer := p.Layer(layers.LayerTypeIPv4)
	if ipLayer == nil {
		return nil
	}
	tcpLayer := p.Layer(layers.LayerTypeTCP)
	if tcpLayer == nil {
		return nil
	}
	ip, _ := ipLayer.(*layers.IPv4)
	ip.TTL = ttl
	tcp, _ := tcpLayer.(*layers.TCP)
	opts := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	tcp.SetNetworkLayerForChecksum(ip)
	rawPacketBuf := gopacket.NewSerializeBuffer()
	if err := gopacket.SerializeLayers(rawPacketBuf, opts, ip, tcp); err != nil {
		return nil
	}
	return rawPacketBuf.Bytes()
}

// We use this to deal with rfc792 implementations where
// the original packet is NOT sent back via ICMP payload but
// instead 64 bits of the original packet are sent.
// https://tools.ietf.org/html/rfc792
// Returns a TCP Flow.
// XXX obviously the 64 bits could be from a UDP packet or something else
// however this is *good-enough* for NFQueue TCP traceroute!
// XXX should we look at the protocol specified in the IP header
// and set it's type here? no we should probably not even get this
// far if the IP header has something other than TCP specified...
func getTCPFlowFromTCPHead(data []byte) gopacket.Flow {
	var srcPort, dstPort layers.TCPPort
	srcPort = layers.TCPPort(binary.BigEndian.Uint16(data[0:2]))
	dstPort = layers.TCPPort(binary.BigEndian.Uint16(data[2:4]))
	// XXX convert to tcp/ip flow
	tcpSrc := layers.NewTCPPortEndpoint(srcPort)
	tcpDst := layers.NewTCPPortEndpoint(dstPort)
	tcpFlow, _ := gopacket.FlowFromEndpoints(tcpSrc, tcpDst)
	// error (^ _) is only non-nil if the two endpoint types don't match
	return tcpFlow
}

// given a byte array packet return a tcp/ip flow
func getPacketFlow(packet []byte) flowKey {
	var ip layers.IPv4
	var tcp layers.TCP
	decoded := []gopacket.LayerType{}
	parser := gopacket.NewDecodingLayerParser(layers.LayerTypeIPv4, &ip, &tcp)
	err := parser.DecodeLayers(packet, &decoded)
	if err != nil {
		// XXX last 64 bits... we only use the last 32 bits
		tcpHead := packet[len(packet)-8 : len(packet)]
		tcpFlow := getTCPFlowFromTCPHead(tcpHead)
		return flowKey{ip.NetworkFlow(), tcpFlow}
	}
	return flowKey{ip.NetworkFlow(), tcp.TransportFlow()}
}

/***
use this rough POC with an iptables nfqueue rule that will select
a tcp flow direction like this:
iptables -A OUTPUT -j NFQUEUE --queue-num 0 -p tcp --dport 2666

***/
func main() {
	options := NFQueueTraceObserverOptions{
		iface:        "wlan0",
		ttlMax:       40,
		ttlRepeatMax: 3,
		mangleFreq:   6,
	}
	o := NewNFQueueTraceObserver(options)
	o.Start()
	// XXX
	<-o.finished
}
<|endoftext|>


<commit_before>package lifted

func (n *LiteralNode) dnf() Formula { return n }

func (TrueNode) dnf() Formula { return TrueNode(1) }

func (FalseNode) dnf() Formula { return FalseNode(1) }

func (n *AndNode) dnf() Formula {
	n.Left = n.Left.dnf()
	n.Right = n.Right.dnf()

	var disj *OrNode
	var other Formula
	if l, ok := n.Left.(*OrNode); ok {
		disj = l
		other = n.Right
	} else if r, ok := n.Right.(*OrNode); ok {
		disj = r
		other = n.Left
	}

	if disj == nil {
		return n
	}

	left := Conjunct(other, disj.Left)
	right := Conjunct(other, disj.Right)
	return Disjunct(left.dnf(), right.dnf()).dnf()
}

func (n *OrNode) dnf() Formula {
	n.Left = n.Left.dnf()
	n.Right = n.Right.dnf()
	return n
}

func (n *NotNode) dnf() Formula {
	switch f := n.Formula.dnf().(type) {
	case *OrNode:
		m := AndNode{
			BinaryNode{Left: Negate(f.Left), Right: Negate(f.Right)},
		}
		return m.dnf()
	case *AndNode:
		return &OrNode{
			BinaryNode{Left: Negate(f.Left), Right: Negate(f.Right)},
		}
	}
	return Negate(n.Formula)
}

func (n *ForallNode) dnf() Formula {
	panic("ForallNode in the tree when converting to DFN")
}

func (n *ExistsNode) dnf() Formula {
	panic("ExistsNode in the tree when converting to DFN")
}

func (n *WhenNode) dnf() Formula {
	disj := Formula(FalseNode(0))
	conds := collectDisjuncts(n.Condition.dnf())
	if len(conds) == 1 {
		return n
	}
	for _, cond := range conds {
		nd := &WhenNode{
			Condition: cond,
			UnaryNode: n.UnaryNode,
		}
		disj = Disjunct(disj, Formula(nd))
	}
	return disj
}

func (n *AssignNode) dnf() Formula { return n }

func collectDisjuncts(f Formula) (fs []Formula) {
	switch n := f.(type) {
	case *OrNode:
		fs = append(fs, collectDisj(n.Left)...)
		fs = append(fs, collectDisj(n.Right)...)
	default:
		fs = append(fs, n)
	}
	return
}
<commit_msg>collectOrs instead of the collectDisjuncts.  Also, there was an error because all of these didn't get renamed in the previous commit.<commit_after>package lifted

func (n *LiteralNode) dnf() Formula { return n }

func (TrueNode) dnf() Formula { return TrueNode(1) }

func (FalseNode) dnf() Formula { return FalseNode(1) }

func (n *AndNode) dnf() Formula {
	n.Left = n.Left.dnf()
	n.Right = n.Right.dnf()

	var disj *OrNode
	var other Formula
	if l, ok := n.Left.(*OrNode); ok {
		disj = l
		other = n.Right
	} else if r, ok := n.Right.(*OrNode); ok {
		disj = r
		other = n.Left
	}

	if disj == nil {
		return n
	}

	left := Conjunct(other, disj.Left)
	right := Conjunct(other, disj.Right)
	return Disjunct(left.dnf(), right.dnf()).dnf()
}

func (n *OrNode) dnf() Formula {
	n.Left = n.Left.dnf()
	n.Right = n.Right.dnf()
	return n
}

func (n *NotNode) dnf() Formula {
	switch f := n.Formula.dnf().(type) {
	case *OrNode:
		m := AndNode{
			BinaryNode{Left: Negate(f.Left), Right: Negate(f.Right)},
		}
		return m.dnf()
	case *AndNode:
		return &OrNode{
			BinaryNode{Left: Negate(f.Left), Right: Negate(f.Right)},
		}
	}
	return Negate(n.Formula)
}

func (n *ForallNode) dnf() Formula {
	panic("ForallNode in the tree when converting to DFN")
}

func (n *ExistsNode) dnf() Formula {
	panic("ExistsNode in the tree when converting to DFN")
}

func (n *WhenNode) dnf() Formula {
	disj := Formula(FalseNode(0))
	conds := collectOrs(n.Condition.dnf())
	if len(conds) == 1 {
		return n
	}
	for _, cond := range conds {
		nd := &WhenNode{
			Condition: cond,
			UnaryNode: n.UnaryNode,
		}
		disj = Disjunct(disj, Formula(nd))
	}
	return disj
}

func (n *AssignNode) dnf() Formula { return n }

func collectOrs(f Formula) (fs []Formula) {
	switch n := f.(type) {
	case *OrNode:
		fs = append(fs, collectOrs(n.Left)...)
		fs = append(fs, collectOrs(n.Right)...)
	default:
		fs = append(fs, n)
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha256"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"os/user"
	"path"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/mitnk/goutils/encrypt"
	"github.com/orcaman/concurrent-map"
)

var VERSION = "1.7.1"
var countConnected = 0
var KEY = getKey()
var DEBUG = false

var Servers = cmap.New()

func main() {
	host := flag.String("host", "0.0.0.0", "host")
	port := flag.String("port", "12345", "port")
	_debug := flag.Bool("v", false, "verbose")
	flag.Usage = func() {
		fmt.Printf("Usage of lightsocks v%s:\n", VERSION)
		fmt.Printf("lightsocks [flags]\n")
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.Parse()
	remote, err := net.Listen("tcp", *host+":"+*port)
	if err != nil {
		fmt.Printf("net listen: %v\n", err)
		os.Exit(2)
	}
	defer remote.Close()
	DEBUG = *_debug

	info("lightsocks v%s", VERSION)
	info("listen on %s:%s", *host, *port)

	go printServersInfo()
	for {
		local, err := remote.Accept()
		if err != nil {
			info("error when accept: %v", err)
			continue
		}
		go handleLocal(local)
	}
}

func printServersInfo() {
	for {
		select {
		case <-time.After(600 * time.Second):
			ts_now := time.Now().Unix()
			keys := Servers.Keys()
			info("[REPORT] We have %d servers connected", len(keys))
			for i, key := range keys {
				if tmp, ok := Servers.Get(key); ok {
					bytes := int64(0)
					ts_span := int64(0)
					if tmp, ok := tmp.(cmap.ConcurrentMap).Get("bytes"); ok {
						bytes = tmp.(int64)
					}
					if tmp, ok := tmp.(cmap.ConcurrentMap).Get("ts"); ok {
						ts_span = ts_now - tmp.(int64)
					}

					str_bytes := ""
					if bytes > 1024*1024*1024 {
						str_bytes += fmt.Sprintf("%.2fG", float64(bytes)/(1024*1024*1024))
					} else if bytes > 1024*1024 {
						str_bytes += fmt.Sprintf("%.2fM", float64(bytes)/(1024*1024))
					} else {
						str_bytes += fmt.Sprintf("%.2fK", float64(bytes)/1024)
					}

					str_span := ""
					if ts_span > 3600 {
						str_span += fmt.Sprintf("%dh", ts_span/3600)
					}
					if ts_span > 60 {
						str_span += fmt.Sprintf("%dm", (ts_span%3600)/60)
					}
					str_span += fmt.Sprintf("%ds", ts_span%60)
					info("[REPORT] [%d][%s] %s: %s", i, str_span, key, str_bytes)
				}
			}
		}
	}
}

func handleLocal(local net.Conn) {
	countConnected += 1
	defer func() {
		local.Close()
		countConnected -= 1
		debug("closed local")
	}()

	debug("local connected: %v", local.RemoteAddr())
	buffer := make([]byte, 1)
	_, err := io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read first byte from local")
		return
	}
	buffer = make([]byte, buffer[0])
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read validation data from local")
		return
	}
	dataCheck, err := encrypt.Decrypt(buffer, KEY)
	if err != nil {
		info("invalid local: %v", err)
		return
	}
	if !reflect.DeepEqual(KEY[8:16], dataCheck) {
		info("invalid local: checker types not eq")
		return
	}
	buffer = make([]byte, 1)
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read size from local")
		return
	}

	buffer = make([]byte, buffer[0])
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read host from local")
		return
	}
	host, err := encrypt.Decrypt(buffer, KEY)
	if err != nil {
		info("ERROR: cannot decrypt host")
		return
	}

	buffer = make([]byte, 2)
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read port from local")
		return
	}
	port := binary.BigEndian.Uint16(buffer)

	url := net.JoinHostPort(string(host), strconv.Itoa(int(port)))
	if strings.Contains(url, "[[") && strings.Contains(url, "]]") {
		// not known yet why, but the url could be something like this:
		// dial tcp: missing port in address [[::ffff:220.249.243.126]]:80
		// we just fix it here.
		url = strings.Replace(url, "[[", "[", 1)
		url = strings.Replace(url, "]]", "]", 1)
	}
	server, err := net.DialTimeout("tcp", url, time.Second*60)
	if err != nil {
		info("ERROR: cannot dial to server %s: %v", url, err)
		return
	}
	info("connected to server: %s", url)
	initServers(url, 0)

	defer func() {
		server.Close()
		deleteServers(url)
		debug("closed server")
	}()

	ch_local := make(chan []byte)
	ch_server := make(chan DataInfo)
	go readDataFromLocal(ch_local, local)
	go readDataFromServer(ch_server, server, url)

	shouldStop := false
	for {
		if shouldStop {
			break
		}

		select {
		case data, ok := <-ch_local:
			if !ok {
				shouldStop = true
				break
			}
			server.Write(data)
		case di, ok := <-ch_server:
			if !ok {
				shouldStop = true
				break
			}
			buffer = encrypt.Encrypt(di.data[:di.size], KEY)
			b := make([]byte, 2)
			binary.BigEndian.PutUint16(b, uint16(len(buffer)))
			local.Write(b)
			local.Write(buffer)
		case <-time.After(3600 * time.Second):
			debug("timeout on %s", url)
			shouldStop = true
			break
		}
	}
}

func readDataFromServer(ch chan DataInfo, conn net.Conn, url string) {
	debug("enter readDataFromServer")
	defer func() {
		debug("leave readDataFromServer")
	}()
	for {
		data := make([]byte, 7000+rand.Intn(2000))
		n, err := conn.Read(data)
		if err != nil {
			break
		}
		incrServers(url, int64(n))
		debug("data from server:\n%s", data[:n])
		ch <- DataInfo{data, n}
	}
	close(ch)
}

func readDataFromLocal(ch chan []byte, conn net.Conn) {
	debug("enter readDataFromLocal")
	defer func() {
		debug("leave readDataFromLocal")
	}()
	for {
		buffer := make([]byte, 2)
		_, err := io.ReadFull(conn, buffer)
		if err != nil {
			break
		}
		size := binary.BigEndian.Uint16(buffer)
		buffer = make([]byte, size)
		_, err = io.ReadFull(conn, buffer)
		if err != nil {
			break
		}
		data, err := encrypt.Decrypt(buffer, KEY)
		if err != nil {
			fmt.Printf("ERROR: cannot decrypt data from local.")
			break
		}
		debug("data from local:\n%s", data)
		ch <- data
	}
	close(ch)
}

func getKey() []byte {
	usr, err := user.Current()
	if err != nil {
		fmt.Printf("user current: %v\n", err)
		os.Exit(2)
	}
	fileKey := path.Join(usr.HomeDir, ".lightsockskey")
	data, err := ioutil.ReadFile(fileKey)
	if err != nil {
		fmt.Printf("Failed to load key file: %v\n", err)
		os.Exit(1)
	}
	s := strings.TrimSpace(string(data))
	sum := sha256.Sum256([]byte(s))
	return sum[:]
}

func initServers(key string, bytes int64) {
	m := cmap.New()
	now := time.Now()
	m.Set("ts", now.Unix())
	m.Set("bytes", bytes)
	Servers.Set(key, m)
}

func incrServers(key string, n int64) {
	if m, ok := Servers.Get(key); ok {
		if tmp, ok := m.(cmap.ConcurrentMap).Get("bytes"); ok {
			m.(cmap.ConcurrentMap).Set("bytes", tmp.(int64)+n)
		}
	} else {
		initServers(key, n)
	}
}

func deleteServers(key string) {
	Servers.Remove(key)
}

func getGoixyConfig() []byte {
	usr, err := user.Current()
	if err != nil {
		fmt.Printf("user current: %v\n", err)
		os.Exit(2)
	}
	fileConfig := path.Join(usr.HomeDir, ".goixy/config.json")
	if _, err := os.Stat(fileConfig); os.IsNotExist(err) {
		fmt.Printf("config file is missing: %v\n", fileConfig)
		os.Exit(2)
	}

	data, err := ioutil.ReadFile(fileConfig)
	if err != nil {
		fmt.Printf("failed to load direct-servers file: %v\n", err)
		os.Exit(1)
	}
	return data
}

func info(format string, a ...interface{}) {
	ts := time.Now().Format("2006-01-02 15:04:05")
	prefix := fmt.Sprintf("[%s][%d] ", ts, countConnected)
	fmt.Printf(prefix+format+"\n", a...)
}

func debug(format string, a ...interface{}) {
	if DEBUG {
		info(format, a...)
	}
}

type DataInfo struct {
	data []byte
	size int
}
<commit_msg>clean up<commit_after>package main

import (
	"crypto/sha256"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"os/user"
	"path"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/mitnk/goutils/encrypt"
	"github.com/orcaman/concurrent-map"
)

var VERSION = "1.7.1"
var countConnected = 0
var KEY = getKey()
var DEBUG = false

var Servers = cmap.New()

func main() {
	host := flag.String("host", "0.0.0.0", "host")
	port := flag.String("port", "12345", "port")
	_debug := flag.Bool("v", false, "verbose")
	flag.Usage = func() {
		fmt.Printf("Usage of lightsocks v%s:\n", VERSION)
		fmt.Printf("lightsocks [flags]\n")
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.Parse()
	remote, err := net.Listen("tcp", *host+":"+*port)
	if err != nil {
		fmt.Printf("net listen: %v\n", err)
		os.Exit(2)
	}
	defer remote.Close()
	DEBUG = *_debug

	info("lightsocks v%s", VERSION)
	info("listen on %s:%s", *host, *port)

	go printServersInfo()
	for {
		local, err := remote.Accept()
		if err != nil {
			info("error when accept: %v", err)
			continue
		}
		go handleLocal(local)
	}
}

func printServersInfo() {
	for {
		select {
		case <-time.After(600 * time.Second):
			ts_now := time.Now().Unix()
			keys := Servers.Keys()
			info("[REPORT] We have %d servers connected", len(keys))
			for i, key := range keys {
				if tmp, ok := Servers.Get(key); ok {
					bytes := int64(0)
					ts_span := int64(0)
					if tmp, ok := tmp.(cmap.ConcurrentMap).Get("bytes"); ok {
						bytes = tmp.(int64)
					}
					if tmp, ok := tmp.(cmap.ConcurrentMap).Get("ts"); ok {
						ts_span = ts_now - tmp.(int64)
					}

					str_bytes := ""
					if bytes > 1024*1024*1024 {
						str_bytes += fmt.Sprintf("%.2fG", float64(bytes)/(1024*1024*1024))
					} else if bytes > 1024*1024 {
						str_bytes += fmt.Sprintf("%.2fM", float64(bytes)/(1024*1024))
					} else {
						str_bytes += fmt.Sprintf("%.2fK", float64(bytes)/1024)
					}

					str_span := ""
					if ts_span > 3600 {
						str_span += fmt.Sprintf("%dh", ts_span/3600)
					}
					if ts_span > 60 {
						str_span += fmt.Sprintf("%dm", (ts_span%3600)/60)
					}
					str_span += fmt.Sprintf("%ds", ts_span%60)
					info("[REPORT] [%d][%s] %s: %s", i, str_span, key, str_bytes)
				}
			}
		}
	}
}

func handleLocal(local net.Conn) {
	countConnected += 1
	defer func() {
		local.Close()
		countConnected -= 1
		debug("closed local")
	}()

	debug("local connected: %v", local.RemoteAddr())
	buffer := make([]byte, 1)
	_, err := io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read first byte from local")
		return
	}
	buffer = make([]byte, buffer[0])
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read validation data from local")
		return
	}
	dataCheck, err := encrypt.Decrypt(buffer, KEY)
	if err != nil {
		info("invalid local: %v", err)
		return
	}
	if !reflect.DeepEqual(KEY[8:16], dataCheck) {
		info("invalid local: checker types not eq")
		return
	}
	buffer = make([]byte, 1)
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read size from local")
		return
	}

	buffer = make([]byte, buffer[0])
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read host from local")
		return
	}
	host, err := encrypt.Decrypt(buffer, KEY)
	if err != nil {
		info("ERROR: cannot decrypt host")
		return
	}

	buffer = make([]byte, 2)
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read port from local")
		return
	}
	port := binary.BigEndian.Uint16(buffer)

	url := net.JoinHostPort(string(host), strconv.Itoa(int(port)))
	if strings.Contains(url, "[[") && strings.Contains(url, "]]") {
		// not known yet why, but the url could be something like this:
		// dial tcp: missing port in address [[::ffff:220.249.243.126]]:80
		// we just fix it here.
		url = strings.Replace(url, "[[", "[", 1)
		url = strings.Replace(url, "]]", "]", 1)
	}
	server, err := net.DialTimeout("tcp", url, time.Second*60)
	if err != nil {
		info("ERROR: cannot dial to server %s: %v", url, err)
		return
	}
	info("connected to server: %s", url)
	initServers(url, 0)

	defer func() {
		server.Close()
		deleteServers(url)
		debug("closed server")
	}()

	ch_local := make(chan []byte)
	ch_server := make(chan DataInfo)
	go readDataFromLocal(ch_local, local)
	go readDataFromServer(ch_server, server, url)

	shouldStop := false
	for {
		if shouldStop {
			break
		}

		select {
		case data, ok := <-ch_local:
			if !ok {
				shouldStop = true
				break
			}
			server.Write(data)
		case di, ok := <-ch_server:
			if !ok {
				shouldStop = true
				break
			}
			buffer = encrypt.Encrypt(di.data[:di.size], KEY)
			b := make([]byte, 2)
			binary.BigEndian.PutUint16(b, uint16(len(buffer)))
			local.Write(b)
			local.Write(buffer)
		case <-time.After(3600 * time.Second):
			debug("timeout on %s", url)
			shouldStop = true
			break
		}
	}
}

func readDataFromServer(ch chan DataInfo, conn net.Conn, url string) {
	debug("enter readDataFromServer")
	defer func() {
		debug("leave readDataFromServer")
	}()
	for {
		data := make([]byte, 7000+rand.Intn(2000))
		n, err := conn.Read(data)
		if err != nil {
			break
		}
		incrServers(url, int64(n))
		debug("data from server:\n%s", data[:n])
		ch <- DataInfo{data, n}
	}
	close(ch)
}

func readDataFromLocal(ch chan []byte, conn net.Conn) {
	debug("enter readDataFromLocal")
	defer func() {
		debug("leave readDataFromLocal")
	}()
	for {
		buffer := make([]byte, 2)
		_, err := io.ReadFull(conn, buffer)
		if err != nil {
			break
		}
		size := binary.BigEndian.Uint16(buffer)
		buffer = make([]byte, size)
		_, err = io.ReadFull(conn, buffer)
		if err != nil {
			break
		}
		data, err := encrypt.Decrypt(buffer, KEY)
		if err != nil {
			fmt.Printf("ERROR: cannot decrypt data from local.")
			break
		}
		debug("data from local:\n%s", data)
		ch <- data
	}
	close(ch)
}

func getKey() []byte {
	usr, err := user.Current()
	if err != nil {
		fmt.Printf("user current: %v\n", err)
		os.Exit(2)
	}
	fileKey := path.Join(usr.HomeDir, ".lightsockskey")
	data, err := ioutil.ReadFile(fileKey)
	if err != nil {
		fmt.Printf("Failed to load key file: %v\n", err)
		os.Exit(1)
	}
	s := strings.TrimSpace(string(data))
	sum := sha256.Sum256([]byte(s))
	return sum[:]
}

func initServers(key string, bytes int64) {
	m := cmap.New()
	now := time.Now()
	m.Set("ts", now.Unix())
	m.Set("bytes", bytes)
	Servers.Set(key, m)
}

func incrServers(key string, n int64) {
	if m, ok := Servers.Get(key); ok {
		if tmp, ok := m.(cmap.ConcurrentMap).Get("bytes"); ok {
			m.(cmap.ConcurrentMap).Set("bytes", tmp.(int64)+n)
		}
	} else {
		initServers(key, n)
	}
}

func deleteServers(key string) {
	Servers.Remove(key)
}

func info(format string, a ...interface{}) {
	ts := time.Now().Format("2006-01-02 15:04:05")
	prefix := fmt.Sprintf("[%s][%d] ", ts, countConnected)
	fmt.Printf(prefix+format+"\n", a...)
}

func debug(format string, a ...interface{}) {
	if DEBUG {
		info(format, a...)
	}
}

type DataInfo struct {
	data []byte
	size int
}
<|endoftext|>


<commit_before>// This is a derivative work of roger peppe's stackgraph command.
// For more details, see http://code.google.com/p/rog-go/
//

package traceback

import (
	"bufio"
	"fmt"
	"io"
	"regexp"
	"strconv"
	"strings"
)

// ParseStacks parse stacktrace of go executables.
func ParseStacks(r io.Reader) ([]*Stack, error) {
	var stacks []*Stack
	re := regexp.MustCompile(`goroutine (\d+) \[([\w ]+)\]:`)
	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		line := scanner.Text()
		stack := &Stack{}
		strs := re.FindStringSubmatch(line)
		if strs == nil || strs[0] == "" || strs[1] == "" || strs[2] == "" {
			continue
		}
		stack.ID, _ = strconv.Atoi(strs[1])
		stack.Status = StackStatus(strs[2])
		for scanner.Scan() {
			line := scanner.Text()
			if line == "" {
				// empty line signifies end of a stack
				break
			}
			if strings.HasPrefix(line, "exit status") {
				break
			}
			if strings.Contains(line, "  ") {
				// Looks like a register dump.
				// TODO better heuristic here.
				continue
			}
			var call Call
			if strings.HasSuffix(line, ")") {
				if i := strings.LastIndex(line, "("); i > 0 {
					call.Args = parseArgs(line[i+1 : len(line)-1])
					line = line[0:i]
				}
			}
			call.Func = strings.TrimPrefix(line, "created by ")
			if !scanner.Scan() {
				break
			}
			line = scanner.Text()
			if strings.HasPrefix(line, "\t") || strings.HasPrefix(line, "        ") {
				line = strings.TrimPrefix(line, "        ")
				line = strings.TrimPrefix(line, "\t")
				if i := strings.LastIndex(line, ":"); i >= 0 {
					call.Source = line[0:i]
					line = line[i:]
				}
				if i := strings.LastIndex(line, " +"); i >= 1 {
					call.Line, _ = strconv.Atoi(line[1:i])
				}
			}
			stack.Calls = append(stack.Calls, call)
		}
		if len(stack.Calls) > 0 {
			stacks = append(stacks, stack)
		}
	}
	return stacks, nil
}

func parseArgs(argList string) []uint64 {
	argList = strings.TrimSuffix(argList, ", ...")
	if argList == "" {
		return nil
	}
	parts := strings.Split(argList, ", ")
	args := make([]uint64, len(parts))
	for i, a := range parts {
		n, err := strconv.ParseUint(a, 0, 64)
		if err != nil {
			panic(fmt.Errorf("cannot parse %q (from %q)", a, argList))
			n = 0xdeadbeef
		}
		args[i] = n
	}
	return args
}
<commit_msg>fix for anonymous function<commit_after>// This is a derivative work of roger peppe's stackgraph command.
// For more details, see http://code.google.com/p/rog-go/
//

package traceback

import (
	"bufio"
	"fmt"
	"io"
	"regexp"
	"strconv"
	"strings"
)

// ParseStacks parse stacktrace of go executables.
func ParseStacks(r io.Reader) ([]*Stack, error) {
	var stacks []*Stack
	re := regexp.MustCompile(`goroutine (\d+) \[([\w ]+)\]:`)
	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		line := scanner.Text()
		stack := &Stack{}
		strs := re.FindStringSubmatch(line)
		if strs == nil || strs[0] == "" || strs[1] == "" || strs[2] == "" {
			continue
		}
		stack.ID, _ = strconv.Atoi(strs[1])
		stack.Status = StackStatus(strs[2])
		for scanner.Scan() {
			line := scanner.Text()
			if line == "" {
				// empty line signifies end of a stack
				break
			}
			if strings.HasPrefix(line, "exit status") {
				break
			}
			if strings.Contains(line, "  ") {
				// Looks like a register dump.
				// TODO better heuristic here.
				continue
			}
			var call Call
			if strings.HasSuffix(line, ")") {
				if i := strings.LastIndex(line, "("); i > 0 {
					call.Args = parseArgs(line[i+1 : len(line)-1])
					line = line[0:i]
				}
			}
			call.Func = strings.TrimPrefix(line, "created by ")
			if !scanner.Scan() {
				break
			}
			line = scanner.Text()
			if strings.HasPrefix(line, "\t") || strings.HasPrefix(line, "        ") {
				line = strings.TrimPrefix(line, "        ")
				line = strings.TrimPrefix(line, "\t")
				if i := strings.LastIndex(line, ":"); i >= 0 {
					call.Source = line[0:i]
					line = line[i:]
				}
				if i := strings.LastIndex(line, " +"); i >= 1 {
					call.Line, _ = strconv.Atoi(line[1:i])
				} else {
					call.Line, _ = strconv.Atoi(line[1:])
				}
			}
			stack.Calls = append(stack.Calls, call)
		}
		if len(stack.Calls) > 0 {
			stacks = append(stacks, stack)
		}
	}
	return stacks, nil
}

func parseArgs(argList string) []uint64 {
	argList = strings.TrimSuffix(argList, ", ...")
	if argList == "" {
		return nil
	}
	parts := strings.Split(argList, ", ")
	args := make([]uint64, len(parts))
	for i, a := range parts {
		n, err := strconv.ParseUint(a, 0, 64)
		if err != nil {
			panic(fmt.Errorf("cannot parse %q (from %q)", a, argList))
			n = 0xdeadbeef
		}
		args[i] = n
	}
	return args
}
<|endoftext|>


<commit_before>package models

import "github.com/codegangsta/martini-contrib/binding"

import "gopkg.in/mgo.v2/bson"

import "net/http"

type Event struct {
	Id          bson.ObjectId `json:"id,omitempty" bson:"_id,omitempty"`
	CreatedBy   bson.ObjectId `json:"createdBy" bson:"createdBy"`
	Title       string        `json:"title" bson:"title" binding:"required"`
	Description string        `json:"description" bson:"description" binding:"required"`
	Coordinates []float64     `json:"coordinates" bson:"coordinates" binding:"required"`
	// TODO score
}

// This method implements binding.Validator and is executed by the binding.Validate middleware
func (event Event) Validate(errors *binding.Errors, req *http.Request) {

    if len(event.Title) < 5 {
        errors.Fields["title"] = "Too short; minimum 5 characters"
    } else if len(event.Title) > 50 {
        errors.Fields["title"] = "Too long; maximum 50 characters"
    }

    if len(event.Description) < 5 {
        errors.Fields["description"] = "Too short; minimum 5 characters"
    } else if len(event.Title) > 200 {
        errors.Fields["description"] = "Too long; maximum 200 characters"
    }

    // TODO Validate Coordinates

    // TODO binding:"required" should return the json name, not the field name
}
<commit_msg>Update imports in models/event.go<commit_after>package models

import (
	"github.com/codegangsta/martini-contrib/binding"
	"gopkg.in/mgo.v2/bson"
	"net/http"
)

type Event struct {
	Id          bson.ObjectId `json:"id,omitempty" bson:"_id,omitempty"`
	CreatedBy   bson.ObjectId `json:"createdBy" bson:"createdBy"`
	Title       string        `json:"title" bson:"title" binding:"required"`
	Description string        `json:"description" bson:"description" binding:"required"`
	Coordinates []float64     `json:"coordinates" bson:"coordinates" binding:"required"`
	// TODO score
}

// This method implements binding.Validator and is executed by the binding.Validate middleware
func (event Event) Validate(errors *binding.Errors, req *http.Request) {

	if len(event.Title) < 5 {
		errors.Fields["title"] = "Too short; minimum 5 characters"
	} else if len(event.Title) > 50 {
		errors.Fields["title"] = "Too long; maximum 50 characters"
	}

	if len(event.Description) < 5 {
		errors.Fields["description"] = "Too short; minimum 5 characters"
	} else if len(event.Title) > 200 {
		errors.Fields["description"] = "Too long; maximum 200 characters"
	}

	// TODO Validate Coordinates

	// TODO binding:"required" should return the json name, not the field name
}
<|endoftext|>