<commit_before>package bitwarden

import (
	"encoding/base64"
	"net/http"
	"time"

	"github.com/cozy/cozy-stack/model/bitwarden"
	"github.com/cozy/cozy-stack/model/bitwarden/settings"
	"github.com/cozy/cozy-stack/model/permission"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/crypto"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/cozy/cozy-stack/web/middlewares"
	"github.com/gorilla/websocket"
	"github.com/labstack/echo/v4"
	"github.com/ugorji/go/codec"
)

type transport struct {
	Transport string   `json:"transport"`
	Formats   []string `json:"transferFormats"`
}

// NegotiateHub is the handler for negotiating between the server and the
// client which transport to use for bitwarden notifications. Currently,
// only websocket is supported.
func NegotiateHub(c echo.Context) error {
	if err := middlewares.AllowWholeType(c, permission.GET, consts.BitwardenCiphers); err != nil {
		return c.JSON(http.StatusUnauthorized, echo.Map{
			"error": "invalid token",
		})
	}

	transports := []transport{
		// Bitwarden jslib supports only msgpack (Binary), not JSON (Text)
		{Transport: "WebSockets", Formats: []string{"Binary"}},
	}

	connID := crypto.GenerateRandomBytes(16)
	return c.JSON(http.StatusOK, echo.Map{
		"connectionId":        base64.URLEncoding.EncodeToString(connID),
		"availableTransports": transports,
	})
}

const (
	// Time allowed to write a message to the peer
	writeWait = 10 * time.Second
	// Time allowed to read the next pong message from the peer
	pongWait = 20 * time.Second
	// Send pings to peer with this period (must be less than pongWait)
	pingPeriod = 15 * time.Second
	// Maximum message size allowed from peer
	maxMessageSize = 1024
)

var upgrader = websocket.Upgrader{
	// Don't check the origin of the connexion
	CheckOrigin:     func(r *http.Request) bool { return true },
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

// WebsocketHub is the websocket handler for the hub to send notifications in
// real-time for bitwarden stuff.
func WebsocketHub(c echo.Context) error {
	inst := middlewares.GetInstance(c)
	token := c.QueryParam("access_token")
	pdoc, err := middlewares.ParseJWT(c, inst, token)
	if err != nil || !pdoc.Permissions.AllowWholeType(permission.GET, consts.BitwardenCiphers) {
		return c.JSON(http.StatusUnauthorized, echo.Map{
			"error": "invalid token",
		})
	}
	userID := pdoc.SourceID

	settings, err := settings.Get(inst)
	if err != nil {
		return c.JSON(http.StatusInternalServerError, echo.Map{
			"error": err.Error(),
		})
	}

	ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil)
	if err != nil {
		return err
	}
	defer ws.Close()

	ws.SetReadLimit(maxMessageSize)
	if err = ws.SetReadDeadline(time.Now().Add(pongWait)); err != nil {
		return nil
	}
	ws.SetPongHandler(func(string) error {
		return ws.SetReadDeadline(time.Now().Add(pongWait))
	})

	responses := make(chan []byte)
	ds := realtime.GetHub().Subscriber(inst)
	defer ds.Close()
	go readPump(ws, ds, responses)

	handle := new(codec.MsgpackHandle)
	handle.WriteExt = true
	ticker := time.NewTicker(pingPeriod)
	defer ticker.Stop()

	for {
		select {
		case r := <-responses:
			if err := ws.SetWriteDeadline(time.Now().Add(writeWait)); err != nil {
				return err
			}
			if err := ws.WriteMessage(websocket.BinaryMessage, r); err != nil {
				logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
					Infof("Write error: %s", err)
				return nil
			}
		case e := <-ds.Channel:
			if err := ws.SetWriteDeadline(time.Now().Add(writeWait)); err != nil {
				return err
			}
			notif := buildNotification(e, userID, settings)
			if notif == nil {
				continue
			}
			serialized, err := serializeNotification(handle, *notif)
			if err != nil {
				logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
					Infof("Serialize error: %s", err)
				continue
			}
			if err := ws.WriteMessage(websocket.BinaryMessage, serialized); err != nil {
				return nil
			}
		case <-ticker.C:
			if err := ws.SetWriteDeadline(time.Now().Add(writeWait)); err != nil {
				return err
			}
			if err := ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				return nil
			}
		}
	}
}

var initialResponse = []byte{0x7b, 0x7d, 0x1e} // {}<RS>

func readPump(ws *websocket.Conn, ds *realtime.DynamicSubscriber, responses chan []byte) {
	var msg struct {
		Protocol string `json:"protocol"`
		Version  int    `json:"version"`
	}
	if err := ws.ReadJSON(&msg); err != nil {
		if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNoStatusReceived) {
			logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
				Infof("Read error: %s", err)
		}
		return
	}
	if msg.Protocol != "messagepack" || msg.Version != 1 {
		logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
			Infof("Unexpected message: %v", msg)
		return
	}
	if err := ds.Subscribe(consts.BitwardenFolders); err != nil {
		logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
			Infof("Subscribe error: %s", err)
		return
	}
	if err := ds.Subscribe(consts.BitwardenCiphers); err != nil {
		logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
			Infof("Subscribe error: %s", err)
		return
	}
	responses <- initialResponse

	// Just send back the pings from the client
	for {
		_, msg, err := ws.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNoStatusReceived) {
				logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
					Infof("Read error: %s", err)
			}
			return
		}
		responses <- msg
	}
}

type notificationResponse struct {
	ContextID string `codec:"ContextId"`
	Type      int
	Payload   map[string]interface{}
}

type notification []interface{}

// https://github.com/bitwarden/jslib/blob/master/src/enums/notificationType.ts
const (
	hubCipherUpdate = 0
	hubCipherCreate = 1
	hubLoginDelete  = 2
	hubFolderDelete = 3
	hubCiphers      = 4
	hubVault        = 5
	hubOrgKeys      = 6
	hubFolderCreate = 7
	hubFolderUpdate = 8
	hubCipherDelete = 9
	hubSettings     = 10
	hubLogOut       = 11
)

func buildNotification(e *realtime.Event, userID string, settings *settings.Settings) *notification {
	if e == nil || e.Doc == nil {
		return nil
	}

	doctype := e.Doc.DocType()
	t := -1
	var payload map[string]interface{}
	if doctype == consts.BitwardenFolders {
		payload = buildFolderPayload(e, userID)
		switch e.Verb {
		case realtime.EventCreate:
			t = hubFolderCreate
		case realtime.EventUpdate:
			t = hubFolderUpdate
		case realtime.EventDelete:
			t = hubFolderDelete
		}
	} else if doctype == consts.BitwardenCiphers {
		payload = buildCipherPayload(e, userID, settings)
		switch e.Verb {
		case realtime.EventCreate:
			t = hubCipherCreate
		case realtime.EventUpdate:
			t = hubCipherUpdate
		case realtime.EventDelete:
			t = hubCipherDelete
		}
	}
	if t < 0 {
		return nil
	}

	arg := notificationResponse{
		ContextID: "app_id",
		Type:      t,
		Payload:   payload,
	}
	msg := notification{
		1,                           // MessageType.Invocation
		[]interface{}{},             // Headers
		nil,                         // InvocationId
		"ReceiveMessage",            // Target
		[]notificationResponse{arg}, // Arguments
	}
	return &msg
}

func buildFolderPayload(e *realtime.Event, userID string) map[string]interface{} {
	var updatedAt interface{}
	var date string
	if doc, ok := e.Doc.(*couchdb.JSONDoc); ok {
		meta, _ := doc.M["cozyMetadata"].(map[string]interface{})
		date, _ = meta["updatedAt"].(string)
	} else if doc, ok := e.Doc.(*realtime.JSONDoc); ok {
		meta, _ := doc.M["cozyMetadata"].(map[string]interface{})
		date, _ = meta["updatedAt"].(string)
	} else if doc, ok := e.Doc.(*bitwarden.Folder); ok {
		if doc.Metadata != nil {
			updatedAt = doc.Metadata.UpdatedAt
		}
	}
	if date != "" {
		if t, err := time.Parse(time.RFC3339, date); err == nil {
			updatedAt = t
		}
	}
	if updatedAt == nil {
		updatedAt = time.Now()
	}
	return map[string]interface{}{
		"Id":           e.Doc.ID(),
		"UserId":       userID,
		"RevisionDate": updatedAt,
	}
}

func buildCipherPayload(e *realtime.Event, userID string, settings *settings.Settings) map[string]interface{} {
	var sharedWithCozy bool
	var updatedAt interface{}
	var date string
	if doc, ok := e.Doc.(*couchdb.JSONDoc); ok {
		sharedWithCozy, _ = doc.M["sharedWithCozy"].(bool)
		meta, _ := doc.M["cozyMetadata"].(map[string]interface{})
		date, _ = meta["updatedAt"].(string)
	} else if doc, ok := e.Doc.(*realtime.JSONDoc); ok {
		sharedWithCozy, _ = doc.M["sharedWithCozy"].(bool)
		meta, _ := doc.M["cozyMetadata"].(map[string]interface{})
		date, _ = meta["updatedAt"].(string)
	} else if doc, ok := e.Doc.(*bitwarden.Cipher); ok {
		sharedWithCozy = doc.SharedWithCozy
		if doc.Metadata != nil {
			updatedAt = doc.Metadata.UpdatedAt
		}
	}
	if date != "" {
		if t, err := time.Parse(time.RFC3339, date); err == nil {
			updatedAt = t
		}
	}
	if updatedAt == nil {
		updatedAt = time.Now()
	}
	var orgID, collIDs interface{}
	if sharedWithCozy {
		orgID = settings.OrganizationID
		collIDs = []string{settings.CollectionID}
	}
	return map[string]interface{}{
		"Id":             e.Doc.ID(),
		"UserId":         userID,
		"OrganizationId": orgID,
		"CollectionIds":  collIDs,
		"RevisionDate":   updatedAt,
	}
}

func serializeNotification(handle *codec.MsgpackHandle, notif notification) ([]byte, error) {
	// First serialize the notification to msgpack
	packed := make([]byte, 0, 256)
	encoder := codec.NewEncoderBytes(&packed, handle)
	if err := encoder.Encode(notif); err != nil {
		return nil, err
	}

	// Then, put it in a BinaryMessageFormat
	// https://github.com/aspnet/AspNetCore/blob/master/src/SignalR/clients/ts/signalr-protocol-msgpack/src/BinaryMessageFormat.ts
	size := uint(len(packed))
	lenBuf := make([]byte, 0, 8)
	for size > 0 {
		sizePart := size & 0x7f
		size >>= 7
		if size > 0 {
			sizePart |= 0x80
		}
		lenBuf = append(lenBuf, byte(sizePart))
	}
	buf := make([]byte, len(lenBuf)+len(packed))
	copy(buf[:len(lenBuf)], lenBuf)
	copy(buf[len(lenBuf):], packed)
	return buf, nil
}
<commit_msg>Fix lint<commit_after>package bitwarden

import (
	"encoding/base64"
	"net/http"
	"time"

	"github.com/cozy/cozy-stack/model/bitwarden"
	"github.com/cozy/cozy-stack/model/bitwarden/settings"
	"github.com/cozy/cozy-stack/model/permission"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/crypto"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/cozy/cozy-stack/web/middlewares"
	"github.com/gorilla/websocket"
	"github.com/labstack/echo/v4"
	"github.com/ugorji/go/codec"
)

type transport struct {
	Transport string   `json:"transport"`
	Formats   []string `json:"transferFormats"`
}

// NegotiateHub is the handler for negotiating between the server and the
// client which transport to use for bitwarden notifications. Currently,
// only websocket is supported.
func NegotiateHub(c echo.Context) error {
	if err := middlewares.AllowWholeType(c, permission.GET, consts.BitwardenCiphers); err != nil {
		return c.JSON(http.StatusUnauthorized, echo.Map{
			"error": "invalid token",
		})
	}

	transports := []transport{
		// Bitwarden jslib supports only msgpack (Binary), not JSON (Text)
		{Transport: "WebSockets", Formats: []string{"Binary"}},
	}

	connID := crypto.GenerateRandomBytes(16)
	return c.JSON(http.StatusOK, echo.Map{
		"connectionId":        base64.URLEncoding.EncodeToString(connID),
		"availableTransports": transports,
	})
}

const (
	// Time allowed to write a message to the peer
	writeWait = 10 * time.Second
	// Time allowed to read the next pong message from the peer
	pongWait = 20 * time.Second
	// Send pings to peer with this period (must be less than pongWait)
	pingPeriod = 15 * time.Second
	// Maximum message size allowed from peer
	maxMessageSize = 1024
)

var upgrader = websocket.Upgrader{
	// Don't check the origin of the connexion
	CheckOrigin:     func(r *http.Request) bool { return true },
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

// WebsocketHub is the websocket handler for the hub to send notifications in
// real-time for bitwarden stuff.
func WebsocketHub(c echo.Context) error {
	inst := middlewares.GetInstance(c)
	token := c.QueryParam("access_token")
	pdoc, err := middlewares.ParseJWT(c, inst, token)
	if err != nil || !pdoc.Permissions.AllowWholeType(permission.GET, consts.BitwardenCiphers) {
		return c.JSON(http.StatusUnauthorized, echo.Map{
			"error": "invalid token",
		})
	}
	userID := pdoc.SourceID

	settings, err := settings.Get(inst)
	if err != nil {
		return c.JSON(http.StatusInternalServerError, echo.Map{
			"error": err.Error(),
		})
	}

	ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil)
	if err != nil {
		return err
	}
	defer ws.Close()

	ws.SetReadLimit(maxMessageSize)
	if err = ws.SetReadDeadline(time.Now().Add(pongWait)); err != nil {
		return nil
	}
	ws.SetPongHandler(func(string) error {
		return ws.SetReadDeadline(time.Now().Add(pongWait))
	})

	responses := make(chan []byte)
	ds := realtime.GetHub().Subscriber(inst)
	defer ds.Close()
	go readPump(ws, ds, responses)

	handle := new(codec.MsgpackHandle)
	handle.WriteExt = true
	ticker := time.NewTicker(pingPeriod)
	defer ticker.Stop()

	for {
		select {
		case r := <-responses:
			if err := ws.SetWriteDeadline(time.Now().Add(writeWait)); err != nil {
				return err
			}
			if err := ws.WriteMessage(websocket.BinaryMessage, r); err != nil {
				logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
					Infof("Write error: %s", err)
				return nil
			}
		case e := <-ds.Channel:
			if err := ws.SetWriteDeadline(time.Now().Add(writeWait)); err != nil {
				return err
			}
			notif := buildNotification(e, userID, settings)
			if notif == nil {
				continue
			}
			serialized, err := serializeNotification(handle, *notif)
			if err != nil {
				logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
					Infof("Serialize error: %s", err)
				continue
			}
			if err := ws.WriteMessage(websocket.BinaryMessage, serialized); err != nil {
				return nil
			}
		case <-ticker.C:
			if err := ws.SetWriteDeadline(time.Now().Add(writeWait)); err != nil {
				return err
			}
			if err := ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				return nil
			}
		}
	}
}

var initialResponse = []byte{0x7b, 0x7d, 0x1e} // {}<RS>

func readPump(ws *websocket.Conn, ds *realtime.DynamicSubscriber, responses chan []byte) {
	var msg struct {
		Protocol string `json:"protocol"`
		Version  int    `json:"version"`
	}
	if err := ws.ReadJSON(&msg); err != nil {
		if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNoStatusReceived) {
			logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
				Infof("Read error: %s", err)
		}
		return
	}
	if msg.Protocol != "messagepack" || msg.Version != 1 {
		logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
			Infof("Unexpected message: %v", msg)
		return
	}
	if err := ds.Subscribe(consts.BitwardenFolders); err != nil {
		logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
			Infof("Subscribe error: %s", err)
		return
	}
	if err := ds.Subscribe(consts.BitwardenCiphers); err != nil {
		logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
			Infof("Subscribe error: %s", err)
		return
	}
	responses <- initialResponse

	// Just send back the pings from the client
	for {
		_, msg, err := ws.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNoStatusReceived) {
				logger.WithDomain(ds.DomainName()).WithField("nspace", "bitwarden").
					Infof("Read error: %s", err)
			}
			return
		}
		responses <- msg
	}
}

type notificationResponse struct {
	ContextID string `codec:"ContextId"`
	Type      int
	Payload   map[string]interface{}
}

type notification []interface{}

// https://github.com/bitwarden/jslib/blob/master/src/enums/notificationType.ts
const (
	hubCipherUpdate = 0
	hubCipherCreate = 1
	// hubLoginDelete  = 2
	hubFolderDelete = 3
	// hubCiphers      = 4
	// hubVault        = 5
	// hubOrgKeys      = 6
	hubFolderCreate = 7
	hubFolderUpdate = 8
	hubCipherDelete = 9
	// hubSettings     = 10
	// hubLogOut       = 11
)

func buildNotification(e *realtime.Event, userID string, settings *settings.Settings) *notification {
	if e == nil || e.Doc == nil {
		return nil
	}

	doctype := e.Doc.DocType()
	t := -1
	var payload map[string]interface{}
	if doctype == consts.BitwardenFolders {
		payload = buildFolderPayload(e, userID)
		switch e.Verb {
		case realtime.EventCreate:
			t = hubFolderCreate
		case realtime.EventUpdate:
			t = hubFolderUpdate
		case realtime.EventDelete:
			t = hubFolderDelete
		}
	} else if doctype == consts.BitwardenCiphers {
		payload = buildCipherPayload(e, userID, settings)
		switch e.Verb {
		case realtime.EventCreate:
			t = hubCipherCreate
		case realtime.EventUpdate:
			t = hubCipherUpdate
		case realtime.EventDelete:
			t = hubCipherDelete
		}
	}
	if t < 0 {
		return nil
	}

	arg := notificationResponse{
		ContextID: "app_id",
		Type:      t,
		Payload:   payload,
	}
	msg := notification{
		1,                           // MessageType.Invocation
		[]interface{}{},             // Headers
		nil,                         // InvocationId
		"ReceiveMessage",            // Target
		[]notificationResponse{arg}, // Arguments
	}
	return &msg
}

func buildFolderPayload(e *realtime.Event, userID string) map[string]interface{} {
	var updatedAt interface{}
	var date string
	if doc, ok := e.Doc.(*couchdb.JSONDoc); ok {
		meta, _ := doc.M["cozyMetadata"].(map[string]interface{})
		date, _ = meta["updatedAt"].(string)
	} else if doc, ok := e.Doc.(*realtime.JSONDoc); ok {
		meta, _ := doc.M["cozyMetadata"].(map[string]interface{})
		date, _ = meta["updatedAt"].(string)
	} else if doc, ok := e.Doc.(*bitwarden.Folder); ok {
		if doc.Metadata != nil {
			updatedAt = doc.Metadata.UpdatedAt
		}
	}
	if date != "" {
		if t, err := time.Parse(time.RFC3339, date); err == nil {
			updatedAt = t
		}
	}
	if updatedAt == nil {
		updatedAt = time.Now()
	}
	return map[string]interface{}{
		"Id":           e.Doc.ID(),
		"UserId":       userID,
		"RevisionDate": updatedAt,
	}
}

func buildCipherPayload(e *realtime.Event, userID string, settings *settings.Settings) map[string]interface{} {
	var sharedWithCozy bool
	var updatedAt interface{}
	var date string
	if doc, ok := e.Doc.(*couchdb.JSONDoc); ok {
		sharedWithCozy, _ = doc.M["sharedWithCozy"].(bool)
		meta, _ := doc.M["cozyMetadata"].(map[string]interface{})
		date, _ = meta["updatedAt"].(string)
	} else if doc, ok := e.Doc.(*realtime.JSONDoc); ok {
		sharedWithCozy, _ = doc.M["sharedWithCozy"].(bool)
		meta, _ := doc.M["cozyMetadata"].(map[string]interface{})
		date, _ = meta["updatedAt"].(string)
	} else if doc, ok := e.Doc.(*bitwarden.Cipher); ok {
		sharedWithCozy = doc.SharedWithCozy
		if doc.Metadata != nil {
			updatedAt = doc.Metadata.UpdatedAt
		}
	}
	if date != "" {
		if t, err := time.Parse(time.RFC3339, date); err == nil {
			updatedAt = t
		}
	}
	if updatedAt == nil {
		updatedAt = time.Now()
	}
	var orgID, collIDs interface{}
	if sharedWithCozy {
		orgID = settings.OrganizationID
		collIDs = []string{settings.CollectionID}
	}
	return map[string]interface{}{
		"Id":             e.Doc.ID(),
		"UserId":         userID,
		"OrganizationId": orgID,
		"CollectionIds":  collIDs,
		"RevisionDate":   updatedAt,
	}
}

func serializeNotification(handle *codec.MsgpackHandle, notif notification) ([]byte, error) {
	// First serialize the notification to msgpack
	packed := make([]byte, 0, 256)
	encoder := codec.NewEncoderBytes(&packed, handle)
	if err := encoder.Encode(notif); err != nil {
		return nil, err
	}

	// Then, put it in a BinaryMessageFormat
	// https://github.com/aspnet/AspNetCore/blob/master/src/SignalR/clients/ts/signalr-protocol-msgpack/src/BinaryMessageFormat.ts
	size := uint(len(packed))
	lenBuf := make([]byte, 0, 8)
	for size > 0 {
		sizePart := size & 0x7f
		size >>= 7
		if size > 0 {
			sizePart |= 0x80
		}
		lenBuf = append(lenBuf, byte(sizePart))
	}
	buf := make([]byte, len(lenBuf)+len(packed))
	copy(buf[:len(lenBuf)], lenBuf)
	copy(buf[len(lenBuf):], packed)
	return buf, nil
}
<|endoftext|>


<commit_before>package main

import (
	"net"
	"os"
	"fmt"
    "strings"
    "errors"
    "database/sql"
    _ "github.com/go-sql-driver/mysql"
)

func dieIfError(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "Fatal error: %s", err.Error())
		os.Exit(1)
	}
}

func main() {

    // connect to database to store tags
    db, err := sql.Open("mysql", "carbon_tagger:carbon_tagger_pw@tcp(graphitemachine:3306)/carbon_tagger?charset=utf8")
    dieIfError(err)
    defer db.Close()
    // Open doesn't open a connection. Validate DSN data:
    err = db.Ping()
    dieIfError(err)
    statement_insert_tag, err := db.Prepare("INSERT INTO tags (tag_key, tag_val) VALUES( ?, ? )")
    dieIfError(err)
    statement_insert_metric, err := db.Prepare("INSERT INTO metrics VALUES( ? )")
    dieIfError(err)
    statement_insert_link, err := db.Prepare("INSERT INTO metrics_tags VALUES( ?, ? )")
    dieIfError(err)

    // listen for incoming metrics
	service := ":2003"
	addr, err := net.ResolveTCPAddr("tcp4", service)
	dieIfError(err)
	listener, err := net.ListenTCP("tcp", addr)
	dieIfError(err)

    // TODO connect to outgoing carbon-relay or carbon-cache

	for {
		conn_in, err := listener.Accept()
		if err != nil {
			continue
		}
		handleClient(conn_in, db, statement_insert_tag, statement_insert_metric, statement_insert_link)
		conn_in.Close()
	}
}

func parseTagBasedMetric(metric string) (tags map[string]string, err error) {
    fmt.Printf(">incoming: %s\n", metric)
    // metric_spec value unix_timestamp
    elements := strings.Split(metric, " ")
    if len(elements) != 3  {
        return nil, errors.New("metric doesn't contain exactly 3 nodes")
    }
    nodes := strings.Split(elements[0], ".")
    tags = make(map[string]string)
    // TODO make sure incoming tags are sorted
    for _, node := range nodes {
        tag := strings.Split(node, "=")
        if len(tag) != 2 {
            return nil, errors.New("bad metric spec: each node must be a 'tag_k=tag_v' pair")
        }
        if tag[0] == "" || tag[1] == "" {
            return nil, errors.New("bad metric spec: tag_k and tag_v must be non-empty strings")
        }

        tags[tag[0]] = tag[1]
    }
    if _,ok := tags["unit"]; !ok {
        return nil, errors.New("bad metric spec: unit tag (mandatory) not specified")
    }
    if len(tags) < 2 {
        return nil, errors.New("bad metric spec: must have at least one tag_k/tag_v pair beyond unit")
    }
    return
}

func forwardMetric(metric string) {
    // forward
    //_, err2 := conn_out.Write(buf[0:n])
    //if err2 != nil {
    //    return
   // }
}

func handleClient(conn_in net.Conn, db *sql.DB, statement_insert_tag *sql.Stmt, statement_insert_metric *sql.Stmt, statement_insert_link *sql.Stmt) {
	var buf [512]byte
	for {
		bytes, err := conn_in.Read(buf[0:])
		if err != nil {
			return
		}
        str := string(buf[0:bytes])
        if(strings.ContainsAny(str, "=")) {
            str = strings.TrimSpace(str)
            tags, err:= parseTagBasedMetric(str)
            if err != nil {
                fmt.Printf("DEBUG: invalid tag based metric, ignoring (%s)\n", err)
            } else {
                fmt.Printf("DEBUG: valid tag based metric %s, storing tags and forwarding\n", strings.TrimSpace(str))
                // TODO this should go in a transaction. for now we first store all tag_k=tag_v pairs (if they are orphans, it's not so bad)
                // then add the metric, than the coupling between metric and tags. <-- all this should def. be in a transaction
                tag_ids := make([]int64, 1)
                for tag_k, tag_v := range tags {
                    fmt.Println("Key:", tag_k, "Value:", tag_v)
                    res, err := statement_insert_tag.Exec(tag_k, tag_v)
                    if err != nil {
                        fmt.Fprintf(os.Stderr, "can't store tag %s=%s\n", tag_k, tag_v)
                        return
                    }
                    // TODO on ERROR 1062 (23000), select id
                    id, err := res.LastInsertId()
                    if err != nil {
                        tag_ids = append(tag_ids, id)
                    } else {
                        fmt.Fprintf(os.Stderr, "can't store tag %s=%s\n", tag_k, tag_v)
                        return
                    }
                }
                _, err = statement_insert_metric.Exec(str)
                if err != nil {
                    fmt.Fprintf(os.Stderr, "db insert failed\n")
                    return
                }
                for _, tag_id := range tag_ids {
                    _, err = statement_insert_link.Exec(str, tag_id)
                    if err != nil {
                        fmt.Fprintf(os.Stderr, "db insert failed\n")
                        return
                    }
                }
                forwardMetric(str)
            }
        } else {
            fmt.Printf("DEBUG: not tag based, forwarding metric %s\n", strings.TrimSpace(str))
            forwardMetric(str)
        }
	}
}

<commit_msg>handle already existing tags, better error msgs<commit_after>package main

import (
	"net"
	"os"
	"fmt"
    "strings"
    "errors"
    "database/sql"
    "github.com/go-sql-driver/mysql"
)

func dieIfError(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "Fatal error: %s", err.Error())
		os.Exit(1)
	}
}

func main() {

    // connect to database to store tags
    db, err := sql.Open("mysql", "carbon_tagger:carbon_tagger_pw@tcp(graphitemachine:3306)/carbon_tagger?charset=utf8")
    dieIfError(err)
    defer db.Close()
    // Open doesn't open a connection. Validate DSN data:
    err = db.Ping()
    dieIfError(err)
    statement_insert_tag, err := db.Prepare("INSERT INTO tags (tag_key, tag_val) VALUES( ?, ? )")
    dieIfError(err)
    statement_select_tag, err := db.Prepare("SELECT tag_id FROM tags WHERE tag_key=? AND tag_val=?")
    dieIfError(err)
    statement_insert_metric, err := db.Prepare("INSERT INTO metrics VALUES( ? )")
    dieIfError(err)
    statement_insert_link, err := db.Prepare("INSERT INTO metrics_tags VALUES( ?, ? )")
    dieIfError(err)

    // listen for incoming metrics
	service := ":2003"
	addr, err := net.ResolveTCPAddr("tcp4", service)
	dieIfError(err)
	listener, err := net.ListenTCP("tcp", addr)
	dieIfError(err)

    // TODO connect to outgoing carbon-relay or carbon-cache

	for {
		conn_in, err := listener.Accept()
		if err != nil {
			continue
		}
		handleClient(conn_in, db, statement_insert_tag, statement_select_tag, statement_insert_metric, statement_insert_link)
		conn_in.Close()
	}
}

func parseTagBasedMetric(metric string) (tags map[string]string, err error) {
    fmt.Printf(">incoming: %s\n", metric)
    // metric_spec value unix_timestamp
    elements := strings.Split(metric, " ")
    if len(elements) != 3  {
        return nil, errors.New("metric doesn't contain exactly 3 nodes")
    }
    nodes := strings.Split(elements[0], ".")
    tags = make(map[string]string)
    // TODO make sure incoming tags are sorted
    for _, node := range nodes {
        tag := strings.Split(node, "=")
        if len(tag) != 2 {
            return nil, errors.New("bad metric spec: each node must be a 'tag_k=tag_v' pair")
        }
        if tag[0] == "" || tag[1] == "" {
            return nil, errors.New("bad metric spec: tag_k and tag_v must be non-empty strings")
        }

        tags[tag[0]] = tag[1]
    }
    if _,ok := tags["unit"]; !ok {
        return nil, errors.New("bad metric spec: unit tag (mandatory) not specified")
    }
    if len(tags) < 2 {
        return nil, errors.New("bad metric spec: must have at least one tag_k/tag_v pair beyond unit")
    }
    return
}

func forwardMetric(metric string) {
    // forward
    //_, err2 := conn_out.Write(buf[0:n])
    //if err2 != nil {
    //    return
   // }
}

func handleClient(conn_in net.Conn, db *sql.DB, statement_insert_tag *sql.Stmt, statement_select_tag *sql.Stmt, statement_insert_metric *sql.Stmt, statement_insert_link *sql.Stmt) {
	var buf [512]byte
	for {
		bytes, err := conn_in.Read(buf[0:])
		if err != nil {
			return
		}
        str := string(buf[0:bytes])
        if(strings.ContainsAny(str, "=")) {
            str = strings.TrimSpace(str)
            tags, err:= parseTagBasedMetric(str)
            if err != nil {
                fmt.Printf("DEBUG: invalid tag based metric, ignoring (%s)\n", err)
            } else {
                fmt.Printf("DEBUG: valid tag based metric %s, storing tags and forwarding\n", strings.TrimSpace(str))
                // TODO this should go in a transaction. for now we first store all tag_k=tag_v pairs (if they are orphans, it's not so bad)
                // then add the metric, than the coupling between metric and tags. <-- all this should def. be in a transaction
                tag_ids := make([]int64, 1)
                for tag_k, tag_v := range tags {
                    fmt.Println("Key:", tag_k, "Value:", tag_v)
                    res, err := statement_insert_tag.Exec(tag_k, tag_v)
                    if err != nil {
                        if err.(*mysql.MySQLError).Number == 1062 {  // Error 1062: Duplicate entry
                            res, err := statement_select_tag.Exec(tag_k, tag_v)
                            if err != nil {
                                fmt.Fprintf(os.Stderr, "Can't lookup the id of tag %s=%s: %s\n", tag_k, tag_v, err.Error())
                                return
                            }
                        } else {
                            fmt.Fprintf(os.Stderr, "can't store tag %s=%s: %s\n", tag_k, tag_v, err.Error())
                            return
                        }
                    } else {
                        id, err := res.LastInsertId()
                        if err != nil {
                            tag_ids = append(tag_ids, id)
                        } else {
                            fmt.Fprintf(os.Stderr, "can't get id for just inserted tag %s=%s: %s\n", tag_k, tag_v, err.Error())
                            return
                        }
                    }
                }
                _, err = statement_insert_metric.Exec(str)
                if err != nil {
                    fmt.Fprintf(os.Stderr, "can't store metric:%s\n", err.Error())
                    return
                }
                for _, tag_id := range tag_ids {
                    _, err = statement_insert_link.Exec(str, tag_id)
                    if err != nil {
                        fmt.Fprintf(os.Stderr, "can't link metric to tag:%s\n", err.Error())
                        return
                    }
                }
                forwardMetric(str)
            }
        } else {
            fmt.Printf("DEBUG: not tag based, forwarding metric %s\n", strings.TrimSpace(str))
            forwardMetric(str)
        }
	}
}

<|endoftext|>


<commit_before>package filesystem

import (
	"fmt"
	"io"
	"os"
	"testing"
	"time"

	"github.com/pkg/errors"
)

var (
	now   = time.Now()
	files = []struct {
		name   string
		stat   os.FileInfo
		data   string
		query  string
		exists bool
		inOS   bool
	}{
		{"foo", info{name: "foo", size: 4, mode: 0x1a4, modTime: now}, "1234", "foo", true, false},
		{"bar", info{name: "bar", size: 8, mode: 0x1a5, modTime: now}, "98765432", "/bar", true, false},
		{"d/alpha", info{name: "alpha", size: 8, mode: 0x1a5, modTime: now}, "98765432", "d/alpha", true, false},
		{"d/beta", info{name: "beta", size: 8, mode: 0x1a5, modTime: now}, "98765432", "./d/beta", true, false},
		{"d/gamma", info{name: "gamma", size: 8, mode: 0x1a5, modTime: now}, "98765432", "/d/gamma", true, false},
		{"", info{}, "", "fs_test.go", false, true},
		{"", info{}, "", "fs_test.stop", false, false},
	}
)

func TestFiles(t *testing.T) {
	fs := New()
	fallback := New()
	fallback.Fallback = true

	for i, fs := range []*FileSystem{fs, fallback} {
		for j, tc := range files {
			t.Run(fmt.Sprintf("case %d-%d", i, j), func(t *testing.T) {
				if tc.name != "" {
					err := fs.Add(tc.name, tc.stat.Size(), tc.stat.Mode(), tc.stat.ModTime(), tc.data)
					if err != nil {
						t.Fatalf("didn't expect error %+v", err)
					}

					err = fs.Add(tc.name, tc.stat.Size(), tc.stat.Mode(), tc.stat.ModTime(), tc.data)
					if err == nil {
						t.Fatalf("file already exist, should've gotten an error")
					} else if !os.IsExist(errors.Cause(err)) {
						t.Fatalf("expected %v, got %+v", os.ErrExist, err)
					}
				}

				f, err := fs.Open(tc.query)
				if tc.exists || fs.Fallback && tc.inOS {
					if err != nil {
						t.Fatalf("opening file: %+v", err)
					}
				} else {
					if !os.IsNotExist(errors.Cause(err)) {
						t.Fatalf("expected ErrNotExist, got %+v", err)
					}

					return
				}

				b := make([]byte, tc.stat.Size())

				_, err = f.Read(b)
				if err != nil {
					t.Fatalf("reading file: %+v", err)
				}

				if string(b) != tc.data {
					t.Fatalf("expected data %s, got %s", tc.data, string(b))
				}

				stat, err := f.Stat()
				if err != nil {
					t.Fatalf("file stat: %+v", err)
				}

				if !fs.Fallback || !tc.inOS {
					if stat.Name() != tc.stat.Name() {
						t.Fatalf("expected name %s, got %s", tc.stat.Name(), stat.Name())
					}

					if stat.Size() != tc.stat.Size() {
						t.Fatalf("expected size %s, got %s", tc.stat.Size(), stat.Size())
					}

					if stat.Mode() != tc.stat.Mode() {
						t.Fatalf("expected mode %s, got %s", tc.stat.Mode(), stat.Mode())
					}

					if stat.ModTime() != tc.stat.ModTime() {
						t.Fatalf("expected mod time %s, got %s", tc.stat.ModTime(), stat.ModTime())
					}
				}
			})
		}
	}
}

func TestDir(t *testing.T) {
	fs := New()

	for _, f := range files {
		if f.name != "" {
			fs.Add(f.name, f.stat.Size(), f.stat.Mode(), f.stat.ModTime(), f.data)
		}
	}

	cases := []struct {
		name   string
		exists bool
		n      int
		again  bool
		names  []string
		n2     int
		names2 []string
		eof    bool
	}{
		{"", true, -1, false, []string{"bar", "d", "foo"}, 0, nil, false},
		{"", true, 0, false, []string{"bar", "d", "foo"}, 0, nil, false},
		{"", true, 1, true, []string{"bar"}, 1, []string{"d"}, false},
		{"", true, 2, true, []string{"bar", "d"}, 1, []string{"foo"}, false},
		{"", true, 1, true, []string{"bar"}, 2, []string{"d", "foo"}, false},

		{".", true, 0, false, []string{"bar", "d", "foo"}, 0, nil, false},
		{"/", true, 0, false, []string{"bar", "d", "foo"}, 0, nil, false},
		{"/d", true, 0, false, []string{"alpha", "beta", "gamma"}, 0, nil, false},
		{"./d", true, 0, false, []string{"alpha", "beta", "gamma"}, 0, nil, false},
		{"d", true, 0, false, []string{"alpha", "beta", "gamma"}, 0, nil, false},
		{"d/", true, 0, false, []string{"alpha", "beta", "gamma"}, 0, nil, false},
	}

	for i, tc := range cases {
		t.Run(fmt.Sprintf("case %d", i), func(t *testing.T) {
			f, err := fs.Open(tc.name)

			if err != nil {
				if os.IsNotExist(errors.Cause(err)) != tc.exists {
					t.Fatalf("expected to exist: %v, got %+v", tc.exists, err)
				}
			}

			stats, err := f.Readdir(tc.n)
			if err != nil {
				t.Fatalf("err: %+v", err)
			}

			if len(tc.names) != len(stats) {
				t.Fatalf("expected %d entries, got %d", len(tc.names), len(stats))
			}

			for i, n := range tc.names {
				if stats[i].Name() != n {
					t.Fatalf("expected %s, got %s", n, stats[i].Name())
				}
			}

			if tc.again {
				stats, err = f.Readdir(tc.n2)
				if err != nil {
					if errors.Cause(err) != io.EOF || !tc.eof {
						t.Fatalf("err: %+v", err)
					}
				}

				for i, n := range tc.names2 {
					if stats[i].Name() != n {
						t.Fatalf("expected %s, got %s", n, stats[i].Name())
					}
				}
			}
		})
	}
}
<commit_msg>Update fs_test.go<commit_after>package filesystem

import (
	"fmt"
	"io"
	"os"
	"testing"
	"time"

	"github.com/pkg/errors"
)

var (
	now   = time.Now()
	files = []struct {
		name   string
		stat   os.FileInfo
		data   string
		query  string
		exists bool
		inOS   bool
	}{
		{"foo", info{name: "foo", size: 4, mode: 0x1a4, modTime: now}, "1234", "foo", true, false},
		{"bar", info{name: "bar", size: 8, mode: 0x1a5, modTime: now}, "98765432", "/bar", true, false},
		{"d/alpha", info{name: "alpha", size: 8, mode: 0x1a5, modTime: now}, "98765432", "d/alpha", true, false},
		{"d/beta", info{name: "beta", size: 8, mode: 0x1a5, modTime: now}, "98765432", "./d/beta", true, false},
		{"d/gamma", info{name: "gamma", size: 8, mode: 0x1a5, modTime: now}, "98765432", "/d/gamma", true, false},
		{"", info{}, "", "fs_test.go", false, true},
		{"", info{}, "", "fs_test.stop", false, false},
	}
)

func TestFiles(t *testing.T) {
	fs := New()
	fallback := New()
	fallback.Fallback = true

	for i, fs := range []*FileSystem{fs, fallback} {
		for j, tc := range files {
			t.Run(fmt.Sprintf("case %d-%d", i, j), func(t *testing.T) {
				if tc.name != "" {
					err := fs.Add(tc.name, tc.stat.Size(), tc.stat.Mode(), tc.stat.ModTime(), tc.data)
					if err != nil {
						t.Fatalf("didn't expect error %+v", err)
					}

					err = fs.Add(tc.name, tc.stat.Size(), tc.stat.Mode(), tc.stat.ModTime(), tc.data)
					if err == nil {
						t.Fatalf("file already exist, should've gotten an error")
					} else if !os.IsExist(errors.Cause(err)) {
						t.Fatalf("expected %v, got %+v", os.ErrExist, err)
					}
				}

				f, err := fs.Open(tc.query)
				if tc.exists || fs.Fallback && tc.inOS {
					if err != nil {
						t.Fatalf("opening file: %+v", err)
					}
				} else {
					if !os.IsNotExist(errors.Cause(err)) {
						t.Fatalf("expected ErrNotExist, got %+v", err)
					}

					return
				}

				b := make([]byte, tc.stat.Size())

				_, err = f.Read(b)
				if err != nil {
					t.Fatalf("reading file: %+v", err)
				}

				if string(b) != tc.data {
					t.Fatalf("expected data %s, got %s", tc.data, string(b))
				}

				stat, err := f.Stat()
				if err != nil {
					t.Fatalf("file stat: %+v", err)
				}

				if !fs.Fallback || !tc.inOS {
					if stat.Name() != tc.stat.Name() {
						t.Fatalf("expected name %s, got %s", tc.stat.Name(), stat.Name())
					}

					if stat.Size() != tc.stat.Size() {
						t.Fatalf("expected size %d, got %d", tc.stat.Size(), stat.Size())
					}

					if stat.Mode() != tc.stat.Mode() {
						t.Fatalf("expected mode %s, got %s", tc.stat.Mode(), stat.Mode())
					}

					if stat.ModTime() != tc.stat.ModTime() {
						t.Fatalf("expected mod time %s, got %s", tc.stat.ModTime(), stat.ModTime())
					}
				}
			})
		}
	}
}

func TestDir(t *testing.T) {
	fs := New()

	for _, f := range files {
		if f.name != "" {
			fs.Add(f.name, f.stat.Size(), f.stat.Mode(), f.stat.ModTime(), f.data)
		}
	}

	cases := []struct {
		name   string
		exists bool
		n      int
		again  bool
		names  []string
		n2     int
		names2 []string
		eof    bool
	}{
		{"", true, -1, false, []string{"bar", "d", "foo"}, 0, nil, false},
		{"", true, 0, false, []string{"bar", "d", "foo"}, 0, nil, false},
		{"", true, 1, true, []string{"bar"}, 1, []string{"d"}, false},
		{"", true, 2, true, []string{"bar", "d"}, 1, []string{"foo"}, false},
		{"", true, 1, true, []string{"bar"}, 2, []string{"d", "foo"}, false},

		{".", true, 0, false, []string{"bar", "d", "foo"}, 0, nil, false},
		{"/", true, 0, false, []string{"bar", "d", "foo"}, 0, nil, false},
		{"/d", true, 0, false, []string{"alpha", "beta", "gamma"}, 0, nil, false},
		{"./d", true, 0, false, []string{"alpha", "beta", "gamma"}, 0, nil, false},
		{"d", true, 0, false, []string{"alpha", "beta", "gamma"}, 0, nil, false},
		{"d/", true, 0, false, []string{"alpha", "beta", "gamma"}, 0, nil, false},
	}

	for i, tc := range cases {
		t.Run(fmt.Sprintf("case %d", i), func(t *testing.T) {
			f, err := fs.Open(tc.name)

			if err != nil {
				if os.IsNotExist(errors.Cause(err)) != tc.exists {
					t.Fatalf("expected to exist: %v, got %+v", tc.exists, err)
				}
			}

			stats, err := f.Readdir(tc.n)
			if err != nil {
				t.Fatalf("err: %+v", err)
			}

			if len(tc.names) != len(stats) {
				t.Fatalf("expected %d entries, got %d", len(tc.names), len(stats))
			}

			for i, n := range tc.names {
				if stats[i].Name() != n {
					t.Fatalf("expected %s, got %s", n, stats[i].Name())
				}
			}

			if tc.again {
				stats, err = f.Readdir(tc.n2)
				if err != nil {
					if errors.Cause(err) != io.EOF || !tc.eof {
						t.Fatalf("err: %+v", err)
					}
				}

				for i, n := range tc.names2 {
					if stats[i].Name() != n {
						t.Fatalf("expected %s, got %s", n, stats[i].Name())
					}
				}
			}
		})
	}
}
<|endoftext|>


<commit_before>package training_test

import (
	"io/ioutil"
	"log"
	"testing"
	"time"

	"github.com/boltdb/bolt"
	"github.com/nim4/DBShield/dbshield/sql"
	"github.com/nim4/DBShield/dbshield/training"
)

func TestMain(m *testing.M) {
	log.SetOutput(ioutil.Discard) // Avoid log outputs
	tmpfile, err := ioutil.TempFile("", "testdb")
	if err != nil {
		panic(err)
	}
	defer tmpfile.Close()
	path := tmpfile.Name()
	training.DBCon, err = bolt.Open(path, 0600, nil)
	if err != nil {
		panic(err)
	}
	training.DBCon.Update(func(tx *bolt.Tx) error {
		tx.CreateBucket([]byte("pattern"))
		tx.CreateBucket([]byte("abnormal"))
		tx.CreateBucket([]byte("state"))
		return nil
	})
	m.Run()
}

func TestAddToTrainingSet(t *testing.T) {
	var err error
	c := sql.QueryContext{
		Query:    []byte("select * from test;"),
		Database: []byte("test"),
		User:     []byte("test"),
		Client:   []byte("127.0.0.1"),
		Time:     time.Now(),
	}
	err = training.AddToTrainingSet(c)
	if err != nil {
		t.Error("Not Expected error", err)
	}

	tmpCon := training.DBCon
	defer func() {
		training.DBCon = tmpCon
	}()
	tmpfile, err := ioutil.TempFile("", "testdb")
	if err != nil {
		panic(err)
	}
	defer tmpfile.Close()
	path := tmpfile.Name()
	training.DBCon, err = bolt.Open(path, 0600, nil)
	if err != nil {
		panic(err)
	}
	err = training.AddToTrainingSet(c)
	if err == nil {
		t.Error("Expected error")
	}
}

func TestCheckQuery(t *testing.T) {
	c1 := sql.QueryContext{
		Query:    []byte("select * from test;"),
		Database: []byte("test"),
		User:     []byte("test"),
		Client:   []byte("127.0.0.1"),
		Time:     time.Now(),
	}
	c2 := sql.QueryContext{
		Query:    []byte("select * from user;"),
		Database: []byte("test"),
		User:     []byte("test"),
		Client:   []byte("127.0.0.1"),
		Time:     time.Now(),
	}
	training.AddToTrainingSet(c1)
	if !training.CheckQuery(c1) {
		t.Error("Expected false")
	}
	if training.CheckQuery(c2) {
		t.Error("Expected true")
	}

	tmpCon := training.DBCon
	defer func() {
		training.DBCon = tmpCon
	}()
	tmpfile, err := ioutil.TempFile("", "testdb")
	if err != nil {
		panic(err)
	}
	defer tmpfile.Close()
	path := tmpfile.Name()
	training.DBCon, err = bolt.Open(path, 0600, nil)
	if err != nil {
		panic(err)
	}
	training.DBCon.Update(func(tx *bolt.Tx) error {
		tx.CreateBucket([]byte("pattern"))
		return err
	})
	defer func() {
		if r := recover(); r == nil {
			t.Error("Expected panic")
		}
	}()
	training.CheckQuery(c1)
}

func BenchmarkAddToTrainingSet(b *testing.B) {
	b.ReportAllocs()
	b.ResetTimer()
	c := sql.QueryContext{
		Query:    []byte("select * from test;"),
		Database: []byte("test"),
		User:     []byte("test"),
		Client:   []byte("127.0.0.1"),
		Time:     time.Now(),
	}
	for i := 0; i < b.N; i++ {
		training.AddToTrainingSet(c)
	}
}
<commit_msg>Improve coverage<commit_after>package training_test

import (
	"io/ioutil"
	"log"
	"testing"
	"time"

	"github.com/boltdb/bolt"
	"github.com/nim4/DBShield/dbshield/config"
	"github.com/nim4/DBShield/dbshield/sql"
	"github.com/nim4/DBShield/dbshield/training"
)

func TestMain(m *testing.M) {
	log.SetOutput(ioutil.Discard) // Avoid log outputs
	tmpfile, err := ioutil.TempFile("", "testdb")
	if err != nil {
		panic(err)
	}
	defer tmpfile.Close()
	path := tmpfile.Name()
	training.DBCon, err = bolt.Open(path, 0600, nil)
	if err != nil {
		panic(err)
	}
	training.DBCon.Update(func(tx *bolt.Tx) error {
		tx.CreateBucket([]byte("pattern"))
		tx.CreateBucket([]byte("abnormal"))
		tx.CreateBucket([]byte("state"))
		return nil
	})
	m.Run()
}

func TestAddToTrainingSet(t *testing.T) {
	var err error
	c := sql.QueryContext{
		Query:    []byte("select * from test;"),
		Database: []byte("test"),
		User:     []byte("test"),
		Client:   []byte("127.0.0.1"),
		Time:     time.Now(),
	}
	err = training.AddToTrainingSet(c)
	if err != nil {
		t.Error("Not Expected error", err)
	}

	tmpCon := training.DBCon
	defer func() {
		training.DBCon = tmpCon
	}()
	tmpfile, err := ioutil.TempFile("", "testdb")
	if err != nil {
		panic(err)
	}
	defer tmpfile.Close()
	path := tmpfile.Name()
	training.DBCon, err = bolt.Open(path, 0600, nil)
	if err != nil {
		panic(err)
	}
	err = training.AddToTrainingSet(c)
	if err == nil {
		t.Error("Expected error")
	}
}

func TestCheckQuery(t *testing.T) {
	config.Config.CheckUser = true
	config.Config.CheckSource = true
	c1 := sql.QueryContext{
		Query:    []byte("select * from test;"),
		Database: []byte("test"),
		User:     []byte("test"),
		Client:   []byte("127.0.0.1"),
		Time:     time.Now(),
	}
	c2 := sql.QueryContext{
		Query:    []byte("select * from user;"),
		Database: []byte("test"),
		User:     []byte("test"),
		Client:   []byte("127.0.0.1"),
		Time:     time.Now(),
	}
	training.AddToTrainingSet(c1)
	if !training.CheckQuery(c1) {
		t.Error("Expected false")
	}
	if training.CheckQuery(c2) {
		t.Error("Expected true")
	}

	tmpCon := training.DBCon
	defer func() {
		training.DBCon = tmpCon
	}()
	tmpfile, err := ioutil.TempFile("", "testdb")
	if err != nil {
		panic(err)
	}
	defer tmpfile.Close()
	path := tmpfile.Name()
	training.DBCon, err = bolt.Open(path, 0600, nil)
	if err != nil {
		panic(err)
	}
	training.DBCon.Update(func(tx *bolt.Tx) error {
		tx.CreateBucket([]byte("pattern"))
		return err
	})
	defer func() {
		if r := recover(); r == nil {
			t.Error("Expected panic")
		}
	}()
	training.CheckQuery(c1)
}

func BenchmarkAddToTrainingSet(b *testing.B) {
	b.ReportAllocs()
	b.ResetTimer()
	c := sql.QueryContext{
		Query:    []byte("select * from test;"),
		Database: []byte("test"),
		User:     []byte("test"),
		Client:   []byte("127.0.0.1"),
		Time:     time.Now(),
	}
	for i := 0; i < b.N; i++ {
		training.AddToTrainingSet(c)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Christian Saide <Supernomad>
// Licensed under the MPL-2.0, for details see https://github.com/Supernomad/protond/blob/master/LICENSE

package filter

import (
	"testing"
	"time"

	"github.com/Supernomad/protond/alert"
	"github.com/Supernomad/protond/cache"
	"github.com/Supernomad/protond/common"
)

var (
	config        *common.Config
	badCfg        *common.Config
	internalCache cache.Cache
	alerts        map[string]alert.Alert
)

func init() {
	filterTimeout, _ := time.ParseDuration("10s")
	badTimeout, _ := time.ParseDuration("1ns")
	log := common.NewLogger(common.NoopLogger)
	config = &common.Config{FilterTimeout: filterTimeout, Log: log}
	badCfg = &common.Config{FilterTimeout: badTimeout, Log: log}
	internalCache, _ = cache.New(cache.MemoryCache, config, &common.PluginConfig{Name: "memory"})
	noopAlert, _ := alert.New(alert.NoopAlert, config, nil)
	alerts = map[string]alert.Alert{"Noop": noopAlert}
}

func TestNonExistentFilterPlugin(t *testing.T) {
	nonExistent, err := New("doesn't exist", nil, nil, nil, nil)
	if err == nil {
		t.Fatal("Something is very very wrong.")
	}
	if nonExistent != nil {
		t.Fatal("Something is very very wrong.")
	}
}

func TestNoop(t *testing.T) {
	noop, err := New(NoopFilter, nil, nil, nil, nil)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": 101010101,
		},
	}
	name := noop.Name()
	if name != "Noop" {
		t.Fatal("Something is very very wrong.")
	}

	test, err := noop.Run(event)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	if test != event {
		t.Fatal("Something is very very wrong.")
	}
}

func TestJavascript(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			event.message = "testing"
			event.added_field = "woot"
			event.new_array = ["this", "should", "be", "handled", 1, 2, 3]
			event.new_object = {"woot": 123, "hello": "world", "sub_array":[1,2,3,"woot"]}
		`,
	}
	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": 101010101,
		},
	}

	name := javascript.Name()
	if name != "Test Filter" {
		t.Fatal("javascript filter name is not properly handled")
	}

	test, err := javascript.Run(event)
	if err != nil {
		t.Fatalf("Error occurred: %s", err.Error())
	}

	if test.Data["message"] != "testing" {
		t.Fatalf("javascript filter failed to overwrite existing 'message' field")
	}

	if test.Data["added_field"] == nil || test.Data["added_field"] != "woot" {
		t.Fatalf("javascript filter failed to add a new field 'added_field' and/or set its value correctly")
	}

	if test.Data["new_array"] == nil || len(test.Data["new_array"].([]interface{})) != 7 {
		t.Fatalf("javascript filter failed to add a new field 'added_field' and/or set its value correctly")
	}
}

func TestJavascriptAlert(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			event.message = "testing"
			event.added_field = "woot"
			event.new_array = ["this", "should", "be", "handled", 1, 2, 3]
			event.new_object = {"woot": 123, "hello": "world", "sub_array":[1,2,3,"woot"]}
			alert.emit("Doesn't exist", event)
			alert.emit({this: "should fail"}, event)
			alert.emit("Noop", "this should fail")
			alert.emit("Noop", event)
		`,
	}
	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": 101010101,
		},
	}

	_, err = javascript.Run(event)
	if err != nil {
		t.Fatalf("Something is very very wrong. %s", err.Error())
	}
}

func TestJavascriptInternalCache(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			event.message = "testing"
			event.added_field = "woot"
			event.new_array = ["this", "should", "be", "handled", 1, 2, 3]
			event.new_object = {"woot": 123, "hello": "world", "sub_array":[1,2,3,"woot"]}
			cache.store("testing", event)
		`,
	}

	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": "woot",
		},
	}

	_, err = javascript.Run(event)
	if err != nil {
		t.Fatalf("Something is very very wrong. %s", err.Error())
	}

	test := javascript.(*Javascript).internalCache.Get("testing")
	if test == nil || len(test) != 1 || test[0].Data["message"] != "testing" {
		t.Fatal("internal cache was not properly updated")
	}
}

func TestJavascriptInternalCacheGet(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			events = cache.get("testing")
			if(events.length == 1) {
				event.stored_events = 1
			}
		`,
	}

	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": "woot",
		},
	}

	test, err := javascript.Run(event)
	if err != nil {
		t.Fatalf("Something is very very wrong. %s", err.Error())
	}

	if test == nil || test.Data["stored_events"].(float64) != 1 {
		t.Fatal("event was not properly updated based on cached events")
	}
}

func TestJavascriptInternalCacheObjectKeyGet(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			events = cache.get({this:"should fail"})
			if(events != undefined) {
				event.failed = true
			} else {
				event.failed = false
			}
		`,
	}

	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": "woot",
		},
	}

	test, err := javascript.Run(event)
	if err != nil {
		t.Fatalf("Something is very very wrong. %s", err.Error())
	}

	if test == nil || test.Data["failed"].(bool) != false {
		t.Fatal("event was not properly updated based on cached events")
	}
}

func TestJavascriptInternalCacheObjectKeyStore(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			cache.store({this:"should fail"}, event)
		`,
	}

	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": "woot",
		},
	}

	_, err = javascript.Run(event)
	if err != nil {
		t.Fatalf("Something is very very wrong. %s", err.Error())
	}

	test := javascript.(*Javascript).internalCache.Get("testing")
	if test == nil || len(test) != 1 {
		t.Fatal("internal cache was not properly updated")
	}
}

func TestJavascriptInternalCacheStringValueStore(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			cache.store("testing", "this should fail")
		`,
	}

	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": "woot",
		},
	}

	_, err = javascript.Run(event)
	if err != nil {
		t.Fatalf("Something is very very wrong. %s", err.Error())
	}

	test := javascript.(*Javascript).internalCache.Get("testing")
	if test == nil || len(test) != 1 {
		t.Fatal("internal cache was not properly updated")
	}
}

func TestJavascriptImproperTypeReturn(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			event = "testing"
		`,
	}
	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": 101010101,
		},
	}

	test, err := javascript.Run(event)
	if err == nil {
		t.Fatal("javascript filter improperly set event type but passed.")
	}
	if test == nil || test.Timestamp != event.Timestamp || test.Data["message"] != event.Data["message"] {
		t.Fatal("javascript filter improperly set event value on failure, should be the unchanged supplied event object.")
	}
}

func TestJavascriptImproperScript(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			event = "testing"
			setTimeout(function() {
				console.log("this will never work")
			}, 100)
		`,
	}
	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": 101010101,
		},
	}

	test, err := javascript.Run(event)
	if err == nil {
		t.Fatal("javascript filter improperly handled a return line in the filter.")
	}
	if test == nil || test.Timestamp != event.Timestamp || test.Data["message"] != event.Data["message"] {
		t.Fatal("javascript filter improperly set event value on failure, should be the unchanged supplied event object.")
	}
}

func TestJavascriptInterrupt(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			function square(a) {
				return a * a
			}

			event.value = 1
			while(true) {
				event.value = square(event.value)
			}
		`,
	}
	javascript, err := New(JavascriptFilter, badCfg, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": 101010101,
		},
	}

	test, err := javascript.Run(event)
	if err == nil {
		t.Fatal("javascript filter improperly handled an interrupt.")
	}
	if test == nil || test.Timestamp != event.Timestamp || test.Data["message"] != event.Data["message"] {
		t.Fatal("javascript filter improperly set event value on failure, should be the unchanged supplied event object.")
	}
}
<commit_msg>Added some more code coverage<commit_after>// Copyright (c) 2017 Christian Saide <Supernomad>
// Licensed under the MPL-2.0, for details see https://github.com/Supernomad/protond/blob/master/LICENSE

package filter

import (
	"testing"
	"time"

	"github.com/Supernomad/protond/alert"
	"github.com/Supernomad/protond/cache"
	"github.com/Supernomad/protond/common"
)

var (
	config        *common.Config
	badCfg        *common.Config
	internalCache cache.Cache
	alerts        map[string]alert.Alert
)

func init() {
	filterTimeout, _ := time.ParseDuration("10s")
	badTimeout, _ := time.ParseDuration("1ns")
	log := common.NewLogger(common.NoopLogger)
	config = &common.Config{FilterTimeout: filterTimeout, Log: log}
	badCfg = &common.Config{FilterTimeout: badTimeout, Log: log}
	internalCache, _ = cache.New(cache.MemoryCache, config, &common.PluginConfig{Name: "memory"})
	noopAlert, _ := alert.New(alert.NoopAlert, config, nil)
	alerts = map[string]alert.Alert{"Noop": noopAlert}
}

func TestNonExistentFilterPlugin(t *testing.T) {
	nonExistent, err := New("doesn't exist", nil, nil, nil, nil)
	if err == nil {
		t.Fatal("Something is very very wrong.")
	}
	if nonExistent != nil {
		t.Fatal("Something is very very wrong.")
	}
}

func TestNoop(t *testing.T) {
	noop, err := New(NoopFilter, nil, nil, nil, nil)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": 101010101,
		},
	}
	name := noop.Name()
	if name != "Noop" {
		t.Fatal("Something is very very wrong.")
	}

	test, err := noop.Run(event)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	if test != event {
		t.Fatal("Something is very very wrong.")
	}
}

func TestJavascript(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			event.message = "testing"
			event.added_field = "woot"
			event.new_array = ["this", "should", "be", "handled", 1, 2, 3]
			event.new_object = {"woot": 123, "hello": "world", "sub_array":[1,2,3,"woot"]}
		`,
	}
	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, nil)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": 101010101,
		},
	}

	name := javascript.Name()
	if name != "Test Filter" {
		t.Fatal("javascript filter name is not properly handled")
	}

	test, err := javascript.Run(event)
	if err != nil {
		t.Fatalf("Error occurred: %s", err.Error())
	}

	if test.Data["message"] != "testing" {
		t.Fatalf("javascript filter failed to overwrite existing 'message' field")
	}

	if test.Data["added_field"] == nil || test.Data["added_field"] != "woot" {
		t.Fatalf("javascript filter failed to add a new field 'added_field' and/or set its value correctly")
	}

	if test.Data["new_array"] == nil || len(test.Data["new_array"].([]interface{})) != 7 {
		t.Fatalf("javascript filter failed to add a new field 'added_field' and/or set its value correctly")
	}
}

func TestJavascriptAlert(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			event.message = "testing"
			event.added_field = "woot"
			event.new_array = ["this", "should", "be", "handled", 1, 2, 3]
			event.new_object = {"woot": 123, "hello": "world", "sub_array":[1,2,3,"woot"]}
			alert.emit("Doesn't exist", event)
			alert.emit({this: "should fail"}, event)
			alert.emit("Noop", "this should fail")
			alert.emit("Noop", event)
		`,
	}
	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": 101010101,
		},
	}

	_, err = javascript.Run(event)
	if err != nil {
		t.Fatalf("Something is very very wrong. %s", err.Error())
	}
}

func TestJavascriptInternalCache(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			event.message = "testing"
			event.added_field = "woot"
			event.new_array = ["this", "should", "be", "handled", 1, 2, 3]
			event.new_object = {"woot": 123, "hello": "world", "sub_array":[1,2,3,"woot"]}
			cache.store("testing", event)
		`,
	}

	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": "woot",
		},
	}

	_, err = javascript.Run(event)
	if err != nil {
		t.Fatalf("Something is very very wrong. %s", err.Error())
	}

	test := javascript.(*Javascript).internalCache.Get("testing")
	if test == nil || len(test) != 1 || test[0].Data["message"] != "testing" {
		t.Fatal("internal cache was not properly updated")
	}
}

func TestJavascriptInternalCacheGet(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			events = cache.get("testing")
			if(events.length == 1) {
				event.stored_events = 1
			}
		`,
	}

	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": "woot",
		},
	}

	test, err := javascript.Run(event)
	if err != nil {
		t.Fatalf("Something is very very wrong. %s", err.Error())
	}

	if test == nil || test.Data["stored_events"].(float64) != 1 {
		t.Fatal("event was not properly updated based on cached events")
	}
}

func TestJavascriptInternalCacheObjectKeyGet(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			events = cache.get({this:"should fail"})
			if(events != undefined) {
				event.failed = true
			} else {
				event.failed = false
			}
		`,
	}

	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": "woot",
		},
	}

	test, err := javascript.Run(event)
	if err != nil {
		t.Fatalf("Something is very very wrong. %s", err.Error())
	}

	if test == nil || test.Data["failed"].(bool) != false {
		t.Fatal("event was not properly updated based on cached events")
	}
}

func TestJavascriptInternalCacheObjectKeyStore(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			cache.store({this:"should fail"}, event)
		`,
	}

	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": "woot",
		},
	}

	_, err = javascript.Run(event)
	if err != nil {
		t.Fatalf("Something is very very wrong. %s", err.Error())
	}

	test := javascript.(*Javascript).internalCache.Get("testing")
	if test == nil || len(test) != 1 {
		t.Fatal("internal cache was not properly updated")
	}
}

func TestJavascriptInternalCacheStringValueStore(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			cache.store("testing", "this should fail")
		`,
	}

	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": "woot",
		},
	}

	_, err = javascript.Run(event)
	if err != nil {
		t.Fatalf("Something is very very wrong. %s", err.Error())
	}

	test := javascript.(*Javascript).internalCache.Get("testing")
	if test == nil || len(test) != 1 {
		t.Fatal("internal cache was not properly updated")
	}
}

func TestJavascriptImproperTypeReturn(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			event = "testing"
		`,
	}
	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": 101010101,
		},
	}

	test, err := javascript.Run(event)
	if err == nil {
		t.Fatal("javascript filter improperly set event type but passed.")
	}
	if test == nil || test.Timestamp != event.Timestamp || test.Data["message"] != event.Data["message"] {
		t.Fatal("javascript filter improperly set event value on failure, should be the unchanged supplied event object.")
	}
}

func TestJavascriptImproperScript(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			event = "testing"
			setTimeout(function() {
				console.log("this will never work")
			}, 100)
		`,
	}
	javascript, err := New(JavascriptFilter, config, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": 101010101,
		},
	}

	test, err := javascript.Run(event)
	if err == nil {
		t.Fatal("javascript filter improperly handled a return line in the filter.")
	}
	if test == nil || test.Timestamp != event.Timestamp || test.Data["message"] != event.Data["message"] {
		t.Fatal("javascript filter improperly set event value on failure, should be the unchanged supplied event object.")
	}
}

func TestJavascriptInterrupt(t *testing.T) {
	filterConfig := &common.FilterConfig{
		Name: "Test Filter",
		Code: `
			function square(a) {
				return a * a
			}

			event.value = 1
			while(true) {
				event.value = square(event.value)
			}
		`,
	}
	javascript, err := New(JavascriptFilter, badCfg, filterConfig, internalCache, alerts)
	if err != nil {
		t.Fatal("Something is very very wrong.")
	}

	event := &common.Event{
		Timestamp: time.Now(),
		Data: map[string]interface{}{
			"message": 101010101,
		},
	}

	test, err := javascript.Run(event)
	if err == nil {
		t.Fatal("javascript filter improperly handled an interrupt.")
	}
	if test == nil || test.Timestamp != event.Timestamp || test.Data["message"] != event.Data["message"] {
		t.Fatal("javascript filter improperly set event value on failure, should be the unchanged supplied event object.")
	}
}
<|endoftext|>


<commit_before>package croc

import (
	"net"
	"time"

	log "github.com/cihub/seelog"
	"github.com/pkg/errors"
	"github.com/schollz/peerdiscovery"
)

func init() {
	SetLogLevel("debug")
}

// Relay initiates a relay
func (c *Croc) Relay() error {
	// start relay
	go c.startRelay()

	// start server
	return c.startServer()
}

// Send will take an existing file or folder and send it through the croc relay
func (c *Croc) Send(fname string, codePhrase string) (err error) {
	log.Debugf("sending %s with compression, encryption: (%v, %v)", fname, c.UseCompression, c.UseEncryption)
	// prepare code phrase
	defer c.cleanup()
	c.cs.Lock()
	c.cs.channel.codePhrase = codePhrase
	if len(codePhrase) == 0 {
		// generate code phrase
		codePhrase = getRandomName()
	}
	if len(codePhrase) < 4 {
		err = errors.New("code phrase must be more than 4 characters")
		c.cs.Unlock()
		return
	}
	c.cs.channel.codePhrase = codePhrase
	c.cs.channel.Channel = codePhrase[:3]
	c.cs.channel.passPhrase = codePhrase[3:]
	log.Debugf("codephrase: '%s'", codePhrase)
	log.Debugf("channel: '%s'", c.cs.channel.Channel)
	log.Debugf("passPhrase: '%s'", c.cs.channel.passPhrase)
	channel := c.cs.channel.Channel
	c.cs.Unlock()

	// start peer discovery
	go func() {
		if c.NoLocal {
			return
		}
		log.Debug("listening for local croc relay...")
		go peerdiscovery.Discover(peerdiscovery.Settings{
			Limit:     1,
			TimeLimit: 600 * time.Second,
			Delay:     50 * time.Millisecond,
			Payload:   []byte(codePhrase[:3]),
		})
	}()

	if len(fname) == 0 {
		err = errors.New("must include filename")
		return
	}
	err = c.processFile(fname)
	if err != nil {
		return
	}

	// start relay for listening
	type runInfo struct {
		err           error
		bothConnected bool
	}
	runClientError := make(chan runInfo, 2)
	go func() {
		if c.NoLocal {
			return
		}
		d := Init()
		d.ServerPort = "8140"
		d.TcpPorts = []string{"27140", "27141"}
		go d.startRelay()
		go d.startServer()
		time.Sleep(100 * time.Millisecond)
		ce := Init()
		ce.WebsocketAddress = "ws://127.0.0.1:8140"
		// copy over the information
		c.cs.Lock()
		ce.cs.Lock()
		ce.cs.channel.codePhrase = codePhrase
		ce.cs.channel.Channel = codePhrase[:3]
		ce.cs.channel.passPhrase = codePhrase[3:]
		ce.cs.channel.fileMetaData = c.cs.channel.fileMetaData
		ce.crocFile = c.crocFile
		ce.crocFileEncrypted = ce.crocFileEncrypted
		ce.isLocal = true
		ce.cs.Unlock()
		c.cs.Unlock()
		defer func() {
			// delete croc files
			ce.cleanup()
		}()
		var ri runInfo
		ri.err = ce.client(0, channel)
		ri.bothConnected = ce.bothConnected
		runClientError <- ri
	}()

	// start main client
	go func() {
		if c.LocalOnly {
			return
		}
		var ri runInfo
		ri.err = c.client(0, channel)
		ri.bothConnected = c.bothConnected
		runClientError <- ri
	}()

	var ri runInfo
	ri = <-runClientError
	if ri.bothConnected || c.LocalOnly || c.NoLocal {
		return ri.err
	}
	ri = <-runClientError
	return ri.err
}

// Receive will receive something through the croc relay
func (c *Croc) Receive(codePhrase string) (err error) {
	defer c.cleanup()
	if !c.NoLocal {
		// try to discovery codephrase and server through peer network
		discovered, errDiscover := peerdiscovery.Discover(peerdiscovery.Settings{
			Limit:     1,
			TimeLimit: 1 * time.Second,
			Delay:     50 * time.Millisecond,
			Payload:   []byte(codePhrase[:3]),
		})
		if errDiscover != nil {
			log.Debug(errDiscover)
		}
		if len(discovered) > 0 {
			log.Debugf("discovered %s on %s", discovered[0].Payload, discovered[0].Address)
			_, connectTimeout := net.DialTimeout("tcp", discovered[0].Address+":27140", 1*time.Second)
			if connectTimeout == nil {
				log.Debug("connected")
				c.WebsocketAddress = "ws://" + discovered[0].Address + ":8140"
				c.isLocal = true
				log.Debug(discovered[0].Address)
				// codePhrase = string(discovered[0].Payload)
			} else {
				log.Debug("but could not connect to ports")
			}
		} else {
			log.Debug("discovered no peers")
		}
	}

	// prepare codephrase
	c.cs.Lock()
	if len(codePhrase) == 0 {
		// prompt codephrase
		codePhrase = promptCodePhrase()
	}
	if len(codePhrase) < 4 {
		err = errors.New("code phrase must be more than 4 characters")
		c.cs.Unlock()
		return
	}
	c.cs.channel.codePhrase = codePhrase
	c.cs.channel.Channel = codePhrase[:3]
	c.cs.channel.passPhrase = codePhrase[3:]
	log.Debugf("codephrase: '%s'", codePhrase)
	log.Debugf("channel: '%s'", c.cs.channel.Channel)
	log.Debugf("passPhrase: '%s'", c.cs.channel.passPhrase)
	channel := c.cs.channel.Channel
	c.cs.Unlock()

	return c.client(1, channel)
}
<commit_msg>discover faster<commit_after>package croc

import (
	"net"
	"time"

	log "github.com/cihub/seelog"
	"github.com/pkg/errors"
	"github.com/schollz/peerdiscovery"
)

func init() {
	SetLogLevel("debug")
}

// Relay initiates a relay
func (c *Croc) Relay() error {
	// start relay
	go c.startRelay()

	// start server
	return c.startServer()
}

// Send will take an existing file or folder and send it through the croc relay
func (c *Croc) Send(fname string, codePhrase string) (err error) {
	log.Debugf("sending %s with compression, encryption: (%v, %v)", fname, c.UseCompression, c.UseEncryption)
	// prepare code phrase
	defer c.cleanup()
	c.cs.Lock()
	c.cs.channel.codePhrase = codePhrase
	if len(codePhrase) == 0 {
		// generate code phrase
		codePhrase = getRandomName()
	}
	if len(codePhrase) < 4 {
		err = errors.New("code phrase must be more than 4 characters")
		c.cs.Unlock()
		return
	}
	c.cs.channel.codePhrase = codePhrase
	c.cs.channel.Channel = codePhrase[:3]
	c.cs.channel.passPhrase = codePhrase[3:]
	log.Debugf("codephrase: '%s'", codePhrase)
	log.Debugf("channel: '%s'", c.cs.channel.Channel)
	log.Debugf("passPhrase: '%s'", c.cs.channel.passPhrase)
	channel := c.cs.channel.Channel
	c.cs.Unlock()

	// start peer discovery
	go func() {
		if c.NoLocal {
			return
		}
		log.Debug("listening for local croc relay...")
		go peerdiscovery.Discover(peerdiscovery.Settings{
			Limit:     1,
			TimeLimit: 600 * time.Second,
			Delay:     50 * time.Millisecond,
			Payload:   []byte(codePhrase[:3]),
		})
	}()

	if len(fname) == 0 {
		err = errors.New("must include filename")
		return
	}
	err = c.processFile(fname)
	if err != nil {
		return
	}

	// start relay for listening
	type runInfo struct {
		err           error
		bothConnected bool
	}
	runClientError := make(chan runInfo, 2)
	go func() {
		if c.NoLocal {
			return
		}
		d := Init()
		d.ServerPort = "8140"
		d.TcpPorts = []string{"27140", "27141"}
		go d.startRelay()
		go d.startServer()
		time.Sleep(100 * time.Millisecond)
		ce := Init()
		ce.WebsocketAddress = "ws://127.0.0.1:8140"
		// copy over the information
		c.cs.Lock()
		ce.cs.Lock()
		ce.cs.channel.codePhrase = codePhrase
		ce.cs.channel.Channel = codePhrase[:3]
		ce.cs.channel.passPhrase = codePhrase[3:]
		ce.cs.channel.fileMetaData = c.cs.channel.fileMetaData
		ce.crocFile = c.crocFile
		ce.crocFileEncrypted = ce.crocFileEncrypted
		ce.isLocal = true
		ce.cs.Unlock()
		c.cs.Unlock()
		defer func() {
			// delete croc files
			ce.cleanup()
		}()
		var ri runInfo
		ri.err = ce.client(0, channel)
		ri.bothConnected = ce.bothConnected
		runClientError <- ri
	}()

	// start main client
	go func() {
		if c.LocalOnly {
			return
		}
		var ri runInfo
		ri.err = c.client(0, channel)
		ri.bothConnected = c.bothConnected
		runClientError <- ri
	}()

	var ri runInfo
	ri = <-runClientError
	if ri.bothConnected || c.LocalOnly || c.NoLocal {
		return ri.err
	}
	ri = <-runClientError
	return ri.err
}

// Receive will receive something through the croc relay
func (c *Croc) Receive(codePhrase string) (err error) {
	defer c.cleanup()
	if !c.NoLocal {
		// try to discovery codephrase and server through peer network
		discovered, errDiscover := peerdiscovery.Discover(peerdiscovery.Settings{
			Limit:     1,
			TimeLimit: 300 * time.Millisecond,
			Delay:     50 * time.Millisecond,
			Payload:   []byte("checking"),
		})
		if errDiscover != nil {
			log.Debug(errDiscover)
		}
		if len(discovered) > 0 {
			log.Debugf("discovered %s on %s", discovered[0].Payload, discovered[0].Address)
			_, connectTimeout := net.DialTimeout("tcp", discovered[0].Address+":27140", 1*time.Second)
			if connectTimeout == nil {
				log.Debug("connected")
				c.WebsocketAddress = "ws://" + discovered[0].Address + ":8140"
				c.isLocal = true
				log.Debug(discovered[0].Address)
				// codePhrase = string(discovered[0].Payload)
			} else {
				log.Debug("but could not connect to ports")
			}
		} else {
			log.Debug("discovered no peers")
		}
	}

	// prepare codephrase
	c.cs.Lock()
	if len(codePhrase) == 0 {
		// prompt codephrase
		codePhrase = promptCodePhrase()
	}
	if len(codePhrase) < 4 {
		err = errors.New("code phrase must be more than 4 characters")
		c.cs.Unlock()
		return
	}
	c.cs.channel.codePhrase = codePhrase
	c.cs.channel.Channel = codePhrase[:3]
	c.cs.channel.passPhrase = codePhrase[3:]
	log.Debugf("codephrase: '%s'", codePhrase)
	log.Debugf("channel: '%s'", c.cs.channel.Channel)
	log.Debugf("passPhrase: '%s'", c.cs.channel.passPhrase)
	channel := c.cs.channel.Channel
	c.cs.Unlock()

	return c.client(1, channel)
}
<|endoftext|>


<commit_before>package matchers_test

import (
	"errors"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/matchers"
	"regexp"
)

func Erroring() error {
	return errors.New("bam")
}

func NotErroring() error {
	return nil
}

type AnyType struct{}

func Invalid() *AnyType {
	return nil
}

var _ = Describe("Succeed", func() {
	It("should succeed if the function succeeds", func() {
		Expect(NotErroring()).Should(Succeed())
	})

	It("should succeed (in the negated) if the function errored", func() {
		Expect(Erroring()).ShouldNot(Succeed())
	})

	It("should not if passed a non-error", func() {
		success, err := (&SucceedMatcher{}).Match(Invalid())
		Expect(success).Should(BeFalse())
		Expect(err).Should(MatchError("Expected an error-type.  Got:\n    <*matchers_test.AnyType | 0x0>: nil"))
	})

	It("doesn't support non-error type", func() {
		success, err := (&SucceedMatcher{}).Match(AnyType{})
		Expect(success).Should(BeFalse())
		Expect(err).Should(MatchError("Expected an error-type.  Got:\n    <matchers_test.AnyType>: {}"))
	})

	It("doesn't support non-error pointer type", func() {
		success, err := (&SucceedMatcher{}).Match(&AnyType{})
		Expect(success).Should(BeFalse())
		Expect(err).Should(MatchError(MatchRegexp(`Expected an error-type.  Got:\n    <*matchers_test.AnyType | 0x[[:xdigit:]]+>: {}`)))
	})

	It("should not succeed with pointer types that conform to error interface", func() {
		err := &CustomErr{"ohai"}
		Expect(err).ShouldNot(Succeed())
	})

	It("should succeed with nil pointers to types that conform to error interface", func() {
		var err *CustomErr = nil
		Expect(err).Should(Succeed())
	})

	It("builds failure message", func() {
		actual := Succeed().FailureMessage(errors.New("oops"))
		actual = regexp.MustCompile(" 0x.*>").ReplaceAllString(actual, " 0x00000000>")
		Expect(actual).To(Equal("Expected success, but got an error:\n    <*errors.errorString | 0x00000000>: {s: \"oops\"}\n    oops"))
	})

	It("builds negated failure message", func() {
		actual := Succeed().NegatedFailureMessage(123)
		Expect(actual).To(Equal("Expected failure, but got no error."))
	})
})
<commit_msg>Fixed import order (#353)<commit_after>package matchers_test

import (
	"errors"
	"regexp"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/matchers"
)

func Erroring() error {
	return errors.New("bam")
}

func NotErroring() error {
	return nil
}

type AnyType struct{}

func Invalid() *AnyType {
	return nil
}

var _ = Describe("Succeed", func() {
	It("should succeed if the function succeeds", func() {
		Expect(NotErroring()).Should(Succeed())
	})

	It("should succeed (in the negated) if the function errored", func() {
		Expect(Erroring()).ShouldNot(Succeed())
	})

	It("should not if passed a non-error", func() {
		success, err := (&SucceedMatcher{}).Match(Invalid())
		Expect(success).Should(BeFalse())
		Expect(err).Should(MatchError("Expected an error-type.  Got:\n    <*matchers_test.AnyType | 0x0>: nil"))
	})

	It("doesn't support non-error type", func() {
		success, err := (&SucceedMatcher{}).Match(AnyType{})
		Expect(success).Should(BeFalse())
		Expect(err).Should(MatchError("Expected an error-type.  Got:\n    <matchers_test.AnyType>: {}"))
	})

	It("doesn't support non-error pointer type", func() {
		success, err := (&SucceedMatcher{}).Match(&AnyType{})
		Expect(success).Should(BeFalse())
		Expect(err).Should(MatchError(MatchRegexp(`Expected an error-type.  Got:\n    <*matchers_test.AnyType | 0x[[:xdigit:]]+>: {}`)))
	})

	It("should not succeed with pointer types that conform to error interface", func() {
		err := &CustomErr{"ohai"}
		Expect(err).ShouldNot(Succeed())
	})

	It("should succeed with nil pointers to types that conform to error interface", func() {
		var err *CustomErr = nil
		Expect(err).Should(Succeed())
	})

	It("builds failure message", func() {
		actual := Succeed().FailureMessage(errors.New("oops"))
		actual = regexp.MustCompile(" 0x.*>").ReplaceAllString(actual, " 0x00000000>")
		Expect(actual).To(Equal("Expected success, but got an error:\n    <*errors.errorString | 0x00000000>: {s: \"oops\"}\n    oops"))
	})

	It("builds negated failure message", func() {
		actual := Succeed().NegatedFailureMessage(123)
		Expect(actual).To(Equal("Expected failure, but got no error."))
	})
})
<|endoftext|>


<commit_before>package flatfs_test

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/jbenet/go-datastore"
	"github.com/jbenet/go-datastore/flatfs"
	"github.com/jbenet/go-datastore/query"
)

func tempdir(t testing.TB) (path string, cleanup func()) {
	path, err := ioutil.TempDir("", "test-datastore-flatfs-")
	if err != nil {
		t.Fatalf("cannot create temp directory: %v", err)
	}

	cleanup = func() {
		if err := os.RemoveAll(path); err != nil {
			t.Errorf("tempdir cleanup failed: %v", err)
		}
	}
	return path, cleanup
}

func TestBadPrefixLen(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	for i := 0; i > -3; i-- {
		_, err := flatfs.New(temp, 0)
		if g, e := err, flatfs.ErrBadPrefixLen; g != e {
			t.Errorf("expected ErrBadPrefixLen, got: %v", g)
		}
	}
}

func TestPutBadValueType(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	err = fs.Put(datastore.NewKey("quux"), 22)
	if g, e := err, datastore.ErrInvalidType; g != e {
		t.Fatalf("expected ErrInvalidType, got: %v\n", g)
	}
}

func TestPut(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	err = fs.Put(datastore.NewKey("quux"), []byte("foobar"))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}
}

func TestGet(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	const input = "foobar"
	err = fs.Put(datastore.NewKey("quux"), []byte(input))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	data, err := fs.Get(datastore.NewKey("quux"))
	if err != nil {
		t.Fatalf("Get failed: %v", err)
	}
	buf, ok := data.([]byte)
	if !ok {
		t.Fatalf("expected []byte from Get, got %T: %v", data, data)
	}
	if g, e := string(buf), input; g != e {
		t.Fatalf("Get gave wrong content: %q != %q", g, e)
	}
}

func TestPutOverwrite(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	const (
		loser  = "foobar"
		winner = "xyzzy"
	)
	err = fs.Put(datastore.NewKey("quux"), []byte(loser))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	err = fs.Put(datastore.NewKey("quux"), []byte(winner))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	data, err := fs.Get(datastore.NewKey("quux"))
	if err != nil {
		t.Fatalf("Get failed: %v", err)
	}
	if g, e := string(data.([]byte)), winner; g != e {
		t.Fatalf("Get gave wrong content: %q != %q", g, e)
	}
}

func TestGetNotFoundError(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	_, err = fs.Get(datastore.NewKey("quux"))
	if g, e := err, datastore.ErrNotFound; g != e {
		t.Fatalf("expected ErrNotFound, got: %v\n", g)
	}
}

func TestStorage(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	const prefixLen = 2
	const prefix = "7175"
	const target = prefix + string(os.PathSeparator) + "71757578.data"
	fs, err := flatfs.New(temp, prefixLen)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	err = fs.Put(datastore.NewKey("quux"), []byte("foobar"))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	seen := false
	walk := func(absPath string, fi os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		path, err := filepath.Rel(temp, absPath)
		if err != nil {
			return err
		}
		switch path {
		case ".", "..":
			// ignore
		case prefix:
			if !fi.IsDir() {
				t.Errorf("prefix directory is not a file? %v", fi.Mode())
			}
			// we know it's there if we see the file, nothing more to
			// do here
		case target:
			seen = true
			if !fi.Mode().IsRegular() {
				t.Errorf("expected a regular file, mode: %04o", fi.Mode())
			}
			if g, e := fi.Mode()&os.ModePerm&0007, os.FileMode(0000); g != e {
				t.Errorf("file should not be world accessible: %04o", fi.Mode())
			}
		default:
			t.Errorf("saw unexpected directory entry: %q %v", path, fi.Mode())
		}
		return nil
	}
	if err := filepath.Walk(temp, walk); err != nil {
		t.Fatal("walk: %v", err)
	}
	if !seen {
		t.Error("did not see the data file")
	}
}

func TestHasNotFound(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	found, err := fs.Has(datastore.NewKey("quux"))
	if err != nil {
		t.Fatalf("Has fail: %v\n", err)
	}
	if g, e := found, false; g != e {
		t.Fatalf("wrong Has: %v != %v", g, e)
	}
}

func TestHasFound(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}
	err = fs.Put(datastore.NewKey("quux"), []byte("foobar"))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	found, err := fs.Has(datastore.NewKey("quux"))
	if err != nil {
		t.Fatalf("Has fail: %v\n", err)
	}
	if g, e := found, true; g != e {
		t.Fatalf("wrong Has: %v != %v", g, e)
	}
}

func TestDeleteNotFound(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	err = fs.Delete(datastore.NewKey("quux"))
	if g, e := err, datastore.ErrNotFound; g != e {
		t.Fatalf("expected ErrNotFound, got: %v\n", g)
	}
}

func TestDeleteFound(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}
	err = fs.Put(datastore.NewKey("quux"), []byte("foobar"))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	err = fs.Delete(datastore.NewKey("quux"))
	if err != nil {
		t.Fatalf("Delete fail: %v\n", err)
	}

	// check that it's gone
	_, err = fs.Get(datastore.NewKey("quux"))
	if g, e := err, datastore.ErrNotFound; g != e {
		t.Fatalf("expected Get after Delete to give ErrNotFound, got: %v\n", g)
	}
}

func TestQuerySimple(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}
	const myKey = "quux"
	err = fs.Put(datastore.NewKey(myKey), []byte("foobar"))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	res, err := fs.Query(query.Query{KeysOnly: true})
	if err != nil {
		t.Fatalf("Query fail: %v\n", err)
	}
	entries, err := res.Rest()
	if err != nil {
		t.Fatalf("Query Results.Rest fail: %v\n", err)
	}
	seen := false
	for _, e := range entries {
		switch e.Key {
		case datastore.NewKey(myKey).String():
			seen = true
		default:
			t.Errorf("saw unexpected key: %q", e.Key)
		}
	}
	if !seen {
		t.Errorf("did not see wanted key %q in %+v", myKey, entries)
	}
}
<commit_msg>flatfs: Don't test file modes on Windows<commit_after>package flatfs_test

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"testing"

	"github.com/jbenet/go-datastore"
	"github.com/jbenet/go-datastore/flatfs"
	"github.com/jbenet/go-datastore/query"
)

func tempdir(t testing.TB) (path string, cleanup func()) {
	path, err := ioutil.TempDir("", "test-datastore-flatfs-")
	if err != nil {
		t.Fatalf("cannot create temp directory: %v", err)
	}

	cleanup = func() {
		if err := os.RemoveAll(path); err != nil {
			t.Errorf("tempdir cleanup failed: %v", err)
		}
	}
	return path, cleanup
}

func TestBadPrefixLen(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	for i := 0; i > -3; i-- {
		_, err := flatfs.New(temp, 0)
		if g, e := err, flatfs.ErrBadPrefixLen; g != e {
			t.Errorf("expected ErrBadPrefixLen, got: %v", g)
		}
	}
}

func TestPutBadValueType(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	err = fs.Put(datastore.NewKey("quux"), 22)
	if g, e := err, datastore.ErrInvalidType; g != e {
		t.Fatalf("expected ErrInvalidType, got: %v\n", g)
	}
}

func TestPut(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	err = fs.Put(datastore.NewKey("quux"), []byte("foobar"))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}
}

func TestGet(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	const input = "foobar"
	err = fs.Put(datastore.NewKey("quux"), []byte(input))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	data, err := fs.Get(datastore.NewKey("quux"))
	if err != nil {
		t.Fatalf("Get failed: %v", err)
	}
	buf, ok := data.([]byte)
	if !ok {
		t.Fatalf("expected []byte from Get, got %T: %v", data, data)
	}
	if g, e := string(buf), input; g != e {
		t.Fatalf("Get gave wrong content: %q != %q", g, e)
	}
}

func TestPutOverwrite(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	const (
		loser  = "foobar"
		winner = "xyzzy"
	)
	err = fs.Put(datastore.NewKey("quux"), []byte(loser))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	err = fs.Put(datastore.NewKey("quux"), []byte(winner))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	data, err := fs.Get(datastore.NewKey("quux"))
	if err != nil {
		t.Fatalf("Get failed: %v", err)
	}
	if g, e := string(data.([]byte)), winner; g != e {
		t.Fatalf("Get gave wrong content: %q != %q", g, e)
	}
}

func TestGetNotFoundError(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	_, err = fs.Get(datastore.NewKey("quux"))
	if g, e := err, datastore.ErrNotFound; g != e {
		t.Fatalf("expected ErrNotFound, got: %v\n", g)
	}
}

func TestStorage(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	const prefixLen = 2
	const prefix = "7175"
	const target = prefix + string(os.PathSeparator) + "71757578.data"
	fs, err := flatfs.New(temp, prefixLen)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	err = fs.Put(datastore.NewKey("quux"), []byte("foobar"))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	seen := false
	walk := func(absPath string, fi os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		path, err := filepath.Rel(temp, absPath)
		if err != nil {
			return err
		}
		switch path {
		case ".", "..":
			// ignore
		case prefix:
			if !fi.IsDir() {
				t.Errorf("prefix directory is not a file? %v", fi.Mode())
			}
			// we know it's there if we see the file, nothing more to
			// do here
		case target:
			seen = true
			if !fi.Mode().IsRegular() {
				t.Errorf("expected a regular file, mode: %04o", fi.Mode())
			}
			if runtime.GOOS != "windows" {
				if g, e := fi.Mode()&os.ModePerm&0007, os.FileMode(0000); g != e {
					t.Errorf("file should not be world accessible: %04o", fi.Mode())
				}
			}
		default:
			t.Errorf("saw unexpected directory entry: %q %v", path, fi.Mode())
		}
		return nil
	}
	if err := filepath.Walk(temp, walk); err != nil {
		t.Fatal("walk: %v", err)
	}
	if !seen {
		t.Error("did not see the data file")
	}
}

func TestHasNotFound(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	found, err := fs.Has(datastore.NewKey("quux"))
	if err != nil {
		t.Fatalf("Has fail: %v\n", err)
	}
	if g, e := found, false; g != e {
		t.Fatalf("wrong Has: %v != %v", g, e)
	}
}

func TestHasFound(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}
	err = fs.Put(datastore.NewKey("quux"), []byte("foobar"))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	found, err := fs.Has(datastore.NewKey("quux"))
	if err != nil {
		t.Fatalf("Has fail: %v\n", err)
	}
	if g, e := found, true; g != e {
		t.Fatalf("wrong Has: %v != %v", g, e)
	}
}

func TestDeleteNotFound(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}

	err = fs.Delete(datastore.NewKey("quux"))
	if g, e := err, datastore.ErrNotFound; g != e {
		t.Fatalf("expected ErrNotFound, got: %v\n", g)
	}
}

func TestDeleteFound(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}
	err = fs.Put(datastore.NewKey("quux"), []byte("foobar"))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	err = fs.Delete(datastore.NewKey("quux"))
	if err != nil {
		t.Fatalf("Delete fail: %v\n", err)
	}

	// check that it's gone
	_, err = fs.Get(datastore.NewKey("quux"))
	if g, e := err, datastore.ErrNotFound; g != e {
		t.Fatalf("expected Get after Delete to give ErrNotFound, got: %v\n", g)
	}
}

func TestQuerySimple(t *testing.T) {
	temp, cleanup := tempdir(t)
	defer cleanup()

	fs, err := flatfs.New(temp, 2)
	if err != nil {
		t.Fatalf("New fail: %v\n", err)
	}
	const myKey = "quux"
	err = fs.Put(datastore.NewKey(myKey), []byte("foobar"))
	if err != nil {
		t.Fatalf("Put fail: %v\n", err)
	}

	res, err := fs.Query(query.Query{KeysOnly: true})
	if err != nil {
		t.Fatalf("Query fail: %v\n", err)
	}
	entries, err := res.Rest()
	if err != nil {
		t.Fatalf("Query Results.Rest fail: %v\n", err)
	}
	seen := false
	for _, e := range entries {
		switch e.Key {
		case datastore.NewKey(myKey).String():
			seen = true
		default:
			t.Errorf("saw unexpected key: %q", e.Key)
		}
	}
	if !seen {
		t.Errorf("did not see wanted key %q in %+v", myKey, entries)
	}
}
<|endoftext|>