<commit_before>package main

import (
    "bytes"
    "strconv"
    "strings"
    "time"

    "labix.org/v2/mgo/bson"
)

type RedisMessage struct {
    Ts            time.Time
    NumberOfBulks int64
    Bulks         []string

    Stream_id    uint32
    Tuple        *IpPortTuple
    CmdlineTuple *CmdlineTuple
    Direction    uint8

    IsRequest bool
    Message   string
}

type RedisStream struct {
    tcpStream *TcpStream

    data []byte

    parseOffset   int
    parseState    int
    bytesReceived int

    message *RedisMessage
}

type RedisTransaction struct {
    Type         string
    tuple        TcpTuple
    Src          Endpoint
    Dst          Endpoint
    ResponseTime int32
    Ts           int64
    JsTs         time.Time
    ts           time.Time
    cmdline      *CmdlineTuple

    Redis bson.M

    Request_raw  string
    Response_raw string

    timer *time.Timer
}

// Keep sorted for future command addition
var RedisCommands = map[string]struct{}{
    "APPEND":           struct{}{},
    "AUTH":             struct{}{},
    "BGREWRITEAOF":     struct{}{},
    "BGSAVE":           struct{}{},
    "BITCOUNT":         struct{}{},
    "BITOP":            struct{}{},
    "BITPOS":           struct{}{},
    "BLPOP":            struct{}{},
    "BRPOP":            struct{}{},
    "BRPOPLPUSH":       struct{}{},
    "CLIENT GETNAME":   struct{}{},
    "CLIENT KILL":      struct{}{},
    "CLIENT LIST":      struct{}{},
    "CLIENT PAUSE":     struct{}{},
    "CLIENT SETNAME":   struct{}{},
    "CONFIG GET":       struct{}{},
    "CONFIG RESETSTAT": struct{}{},
    "CONFIG REWRITE":   struct{}{},
    "CONFIG SET":       struct{}{},
    "DBSIZE":           struct{}{},
    "DEBUG OBJECT":     struct{}{},
    "DEBUG SEGFAULT":   struct{}{},
    "DECR":             struct{}{},
    "DECRBY":           struct{}{},
    "DEL":              struct{}{},
    "DISCARD":          struct{}{},
    "DUMP":             struct{}{},
    "ECHO":             struct{}{},
    "EVAL":             struct{}{},
    "EVALSHA":          struct{}{},
    "EXEC":             struct{}{},
    "EXISTS":           struct{}{},
    "EXPIRE":           struct{}{},
    "EXPIREAT":         struct{}{},
    "FLUSHALL":         struct{}{},
    "FLUSHDB":          struct{}{},
    "GET":              struct{}{},
    "GETBIT":           struct{}{},
    "GETRANGE":         struct{}{},
    "GETSET":           struct{}{},
    "HDEL":             struct{}{},
    "HEXISTS":          struct{}{},
    "HGET":             struct{}{},
    "HGETALL":          struct{}{},
    "HINCRBY":          struct{}{},
    "HINCRBYFLOAT":     struct{}{},
    "HKEYS":            struct{}{},
    "HLEN":             struct{}{},
    "HMGET":            struct{}{},
    "HMSET":            struct{}{},
    "HSCAN":            struct{}{},
    "HSET":             struct{}{},
    "HSETINX":          struct{}{},
    "HVALS":            struct{}{},
    "INCR":             struct{}{},
    "INCRBY":           struct{}{},
    "INCRBYFLOAT":      struct{}{},
    "INFO":             struct{}{},
    "KEYS":             struct{}{},
    "LASTSAVE":         struct{}{},
    "LINDEX":           struct{}{},
    "LINSERT":          struct{}{},
    "LLEN":             struct{}{},
    "LPOP":             struct{}{},
    "LPUSH":            struct{}{},
    "LPUSHX":           struct{}{},
    "LRANGE":           struct{}{},
    "LREM":             struct{}{},
    "LSET":             struct{}{},
    "LTRIM":            struct{}{},
    "MGET":             struct{}{},
    "MIGRATE":          struct{}{},
    "MONITOR":          struct{}{},
    "MOVE":             struct{}{},
    "MSET":             struct{}{},
    "MSETNX":           struct{}{},
    "MULTI":            struct{}{},
    "OBJECT":           struct{}{},
    "PERSIST":          struct{}{},
    "PEXPIRE":          struct{}{},
    "PEXPIREAT":        struct{}{},
    "PFADD":            struct{}{},
    "PFCOUNT":          struct{}{},
    "PFMERGE":          struct{}{},
    "PING":             struct{}{},
    "PSETEX":           struct{}{},
    "PSUBSCRIBE":       struct{}{},
    "PTTL":             struct{}{},
    "PUBLISH":          struct{}{},
    "PUBSUB":           struct{}{},
    "PUNSUBSCRIBE":     struct{}{},
    "QUIT":             struct{}{},
    "RANDOMKEY":        struct{}{},
    "RENAME":           struct{}{},
    "RENAMENX":         struct{}{},
    "RESTORE":          struct{}{},
    "RPOP":             struct{}{},
    "RPOPLPUSH":        struct{}{},
    "RPUSH":            struct{}{},
    "RPUSHX":           struct{}{},
    "SADD":             struct{}{},
    "SAVE":             struct{}{},
    "SCAN":             struct{}{},
    "SCARD":            struct{}{},
    "SCRIPT EXISTS":    struct{}{},
    "SCRIPT FLUSH":     struct{}{},
    "SCRIPT KILL":      struct{}{},
    "SCRIPT LOAD":      struct{}{},
    "SDIFF":            struct{}{},
    "SDIFFSTORE":       struct{}{},
    "SELECT":           struct{}{},
    "SET":              struct{}{},
    "SETBIT":           struct{}{},
    "SETEX":            struct{}{},
    "SETNX":            struct{}{},
    "SETRANGE":         struct{}{},
    "SHUTDOWN":         struct{}{},
    "SINTER":           struct{}{},
    "SINTERSTORE":      struct{}{},
    "SISMEMBER":        struct{}{},
    "SLAVEOF":          struct{}{},
    "SLOWLOG":          struct{}{},
    "SMEMBERS":         struct{}{},
    "SMOVE":            struct{}{},
    "SORT":             struct{}{},
    "SPOP":             struct{}{},
    "SRANDMEMBER":      struct{}{},
    "SREM":             struct{}{},
    "SSCAN":            struct{}{},
    "STRLEN":           struct{}{},
    "SUBSCRIBE":        struct{}{},
    "SUNION":           struct{}{},
    "SUNIONSTORE":      struct{}{},
    "SYNC":             struct{}{},
    "TIME":             struct{}{},
    "TTL":              struct{}{},
    "TYPE":             struct{}{},
    "UNSUBSCRIBE":      struct{}{},
    "UNWATCH":          struct{}{},
    "WATCH":            struct{}{},
    "ZADD":             struct{}{},
    "ZCARD":            struct{}{},
    "ZCOUNT":           struct{}{},
    "ZINCRBY":          struct{}{},
    "ZINTERSTORE":      struct{}{},
    "ZRANGE":           struct{}{},
    "ZRANGEBYSCORE":    struct{}{},
    "ZRANK":            struct{}{},
    "ZREM":             struct{}{},
    "ZREMRANGEBYLEX":   struct{}{},
    "ZREMRANGEBYRANK":  struct{}{},
    "ZREMRANGEBYSCORE": struct{}{},
    "ZREVRANGE":        struct{}{},
    "ZREVRANGEBYSCORE": struct{}{},
    "ZREVRANK":         struct{}{},
    "ZSCAN":            struct{}{},
    "ZSCORE":           struct{}{},
    "ZUNIONSTORE":      struct{}{},
}

var redisTransactionsMap = make(map[TcpTuple]*RedisTransaction, TransactionsHashSize)

func (stream *RedisStream) PrepareForNewMessage() {
    stream.message.NumberOfBulks = 0
    stream.message.Bulks = []string{}
    stream.message.IsRequest = false
}

func redisMessageParser(s *RedisStream) (bool, bool) {

    var err error
    var value string
    m := s.message

    for s.parseOffset < len(s.data) {

        if s.data[s.parseOffset] == '*' {
            //Multi Bulk Message

            line, off := readLine(s.data, s.parseOffset)

            if len(line) == 3 && line[1] == '-' && line[2] == '1' {
                //NULL Multi Bulk
                s.parseOffset = off
                value = "nil"
            } else {

                m.NumberOfBulks, err = strconv.ParseInt(line[1:], 10, 64)

                if err != nil {
                    ERR("Failed to read number of bulk messages: %s", err)
                    return false, false
                }
                s.parseOffset = off
                m.Bulks = []string{}

                continue
            }

        } else if s.data[s.parseOffset] == '$' {
            // Bulk Reply

            line, off := readLine(s.data, s.parseOffset)

            if len(line) == 3 && line[1] == '-' && line[2] == '1' {
                // NULL Bulk Reply
                value = "nil"
                s.parseOffset = off
            } else {
                length, err := strconv.ParseInt(line[1:], 10, 64)
                if err != nil {
                    ERR("Failed to read bulk message: %s", err)
                    return false, false
                }

                s.parseOffset = off

                line, off = readLine(s.data, s.parseOffset)

                if int64(len(line)) != length {
                    ERR("Wrong length of data: %d instead of %d", len(line), length)
                    return false, false
                }
                value = line
                s.parseOffset = off
            }

        } else if s.data[s.parseOffset] == ':' {
            // Integer reply
            line, off := readLine(s.data, s.parseOffset)
            n, err := strconv.ParseInt(line[1:], 10, 64)

            if err != nil {
                ERR("Failed to read integer reply: %s", err)
                return false, false
            }
            value = string(n)
            s.parseOffset = off

        } else if s.data[s.parseOffset] == '+' {
            // Status Reply
            line, off := readLine(s.data, s.parseOffset)

            value = line[1:]
            s.parseOffset = off
        } else if s.data[s.parseOffset] == '-' {
            // Error Reply
            line, off := readLine(s.data, s.parseOffset)

            value = line[1:]
            s.parseOffset = off
        } else {
            DEBUG("redis", "Unexpected message starting with %s", s.data[s.parseOffset:])
            return false, false
        }

        // add value
        if m.NumberOfBulks > 0 {
            m.NumberOfBulks = m.NumberOfBulks - 1
            m.Bulks = append(m.Bulks, value)

            if len(m.Bulks) == 1 {
                // check if it's a command
                if isRedisCommand(value) {
                    m.IsRequest = true
                }
            }

            if m.NumberOfBulks == 0 {
                // the last bulk received
                m.Message = strings.Join(m.Bulks, " ")
                return true, true
            }
        } else {
            m.Message = value
            return true, true
        }

    }   //end for

    return true, false
}

func readLine(data []byte, offset int) (string, int) {
    q := bytes.Index(data[offset:], []byte("\r\n"))
    return string(data[offset : offset+q]), offset + q + 2
}

func ParseRedis(pkt *Packet, tcp *TcpStream, dir uint8) {
    defer RECOVER("ParseRedis exception")

    if tcp.redisData[dir] == nil {
        tcp.redisData[dir] = &RedisStream{
            tcpStream: tcp,
            data:      pkt.payload,
            message:   &RedisMessage{Ts: pkt.ts},
        }
    } else {
        // concatenate bytes
        tcp.redisData[dir].data = append(tcp.redisData[dir].data, pkt.payload...)
    }

    stream := tcp.redisData[dir]
    if stream.message == nil {
        stream.message = &RedisMessage{Ts: pkt.ts}
    }

    ok, complete := redisMessageParser(tcp.redisData[dir])
    if !ok {
        // drop this tcp stream. Will retry parsing with the next
        // segment in it
        tcp.redisData[dir] = nil
        return
    }

    if !ok {
        // drop this tcp stream. Will retry parsing with the next
        // segment in it
        tcp.redisData[dir] = nil
        return
    }

    if complete {

        if stream.message.IsRequest {
            DEBUG("redis", "REDIS request message: %s", stream.message.Message)
        } else {
            DEBUG("redis", "REDIS response message: %s", stream.message.Message)
        }

        // all ok, go to next level
        handleRedis(stream.message, tcp, dir)

        // and reset message
        stream.PrepareForNewMessage()
    }

}

func isRedisCommand(key string) bool {
    _, exists := RedisCommands[key]
    return exists
}

func handleRedis(m *RedisMessage, tcp *TcpStream,
    dir uint8) {

    m.Stream_id = tcp.id
    m.Tuple = tcp.tuple
    m.Direction = dir
    m.CmdlineTuple = procWatcher.FindProcessesTuple(tcp.tuple)

    if m.IsRequest {
        receivedRedisRequest(m)
    } else {
        receivedRedisResponse(m)
    }
}

func receivedRedisRequest(msg *RedisMessage) {
    // Add it to the HT
    tuple := TcpTuple{
        Src_ip:    msg.Tuple.Src_ip,
        Dst_ip:    msg.Tuple.Dst_ip,
        Src_port:  msg.Tuple.Src_port,
        Dst_port:  msg.Tuple.Dst_port,
        stream_id: msg.Stream_id,
    }

    trans := redisTransactionsMap[tuple]
    if trans != nil {
        if len(trans.Redis) != 0 {
            WARN("Two requests without a Response. Dropping old request")
        }
    } else {
        trans = &RedisTransaction{Type: "redis", tuple: tuple}
        redisTransactionsMap[tuple] = trans
    }

    trans.Redis = bson.M{
        "request": msg.Message,
    }
    trans.Request_raw = msg.Message

    trans.cmdline = msg.CmdlineTuple
    trans.ts = msg.Ts
    trans.Ts = int64(trans.ts.UnixNano() / 1000) // transactions have microseconds resolution
    trans.JsTs = msg.Ts
    trans.Src = Endpoint{
        Ip:   Ipv4_Ntoa(tuple.Src_ip),
        Port: tuple.Src_port,
        Proc: string(msg.CmdlineTuple.Src),
    }
    trans.Dst = Endpoint{
        Ip:   Ipv4_Ntoa(tuple.Dst_ip),
        Port: tuple.Dst_port,
        Proc: string(msg.CmdlineTuple.Dst),
    }

    if trans.timer != nil {
        trans.timer.Stop()
    }
    trans.timer = time.AfterFunc(TransactionTimeout, func() { trans.Expire() })

}

func (trans *RedisTransaction) Expire() {

    // remove from map
    delete(redisTransactionsMap, trans.tuple)
}

func receivedRedisResponse(msg *RedisMessage) {

    tuple := TcpTuple{
        Src_ip:    msg.Tuple.Src_ip,
        Dst_ip:    msg.Tuple.Dst_ip,
        Src_port:  msg.Tuple.Src_port,
        Dst_port:  msg.Tuple.Dst_port,
        stream_id: msg.Stream_id,
    }
    trans := redisTransactionsMap[tuple]
    if trans == nil {
        WARN("Response from unknown transaction. Ignoring.")
        return
    }
    // check if the request was received
    if len(trans.Redis) == 0 {
        WARN("Response from unknown transaction. Ignoring.")
        return

    }

    trans.Redis["response"] = msg.Message

    trans.Response_raw = msg.Message

    trans.ResponseTime = int32(msg.Ts.Sub(trans.ts).Nanoseconds() / 1e6) // resp_time in milliseconds

    err := Publisher.PublishRedisTransaction(trans)
    if err != nil {
        WARN("Publish failure: %s", err)
    }

    DEBUG("redis", "Redis transaction completed: %s", trans.Redis)

    // remove from map
    delete(redisTransactionsMap, trans.tuple)
    if trans.timer != nil {
        trans.timer.Stop()
    }

}
<commit_msg>Remove duplicated block<commit_after>package main

import (
    "bytes"
    "strconv"
    "strings"
    "time"

    "labix.org/v2/mgo/bson"
)

type RedisMessage struct {
    Ts            time.Time
    NumberOfBulks int64
    Bulks         []string

    Stream_id    uint32
    Tuple        *IpPortTuple
    CmdlineTuple *CmdlineTuple
    Direction    uint8

    IsRequest bool
    Message   string
}

type RedisStream struct {
    tcpStream *TcpStream

    data []byte

    parseOffset   int
    parseState    int
    bytesReceived int

    message *RedisMessage
}

type RedisTransaction struct {
    Type         string
    tuple        TcpTuple
    Src          Endpoint
    Dst          Endpoint
    ResponseTime int32
    Ts           int64
    JsTs         time.Time
    ts           time.Time
    cmdline      *CmdlineTuple

    Redis bson.M

    Request_raw  string
    Response_raw string

    timer *time.Timer
}

// Keep sorted for future command addition
var RedisCommands = map[string]struct{}{
    "APPEND":           struct{}{},
    "AUTH":             struct{}{},
    "BGREWRITEAOF":     struct{}{},
    "BGSAVE":           struct{}{},
    "BITCOUNT":         struct{}{},
    "BITOP":            struct{}{},
    "BITPOS":           struct{}{},
    "BLPOP":            struct{}{},
    "BRPOP":            struct{}{},
    "BRPOPLPUSH":       struct{}{},
    "CLIENT GETNAME":   struct{}{},
    "CLIENT KILL":      struct{}{},
    "CLIENT LIST":      struct{}{},
    "CLIENT PAUSE":     struct{}{},
    "CLIENT SETNAME":   struct{}{},
    "CONFIG GET":       struct{}{},
    "CONFIG RESETSTAT": struct{}{},
    "CONFIG REWRITE":   struct{}{},
    "CONFIG SET":       struct{}{},
    "DBSIZE":           struct{}{},
    "DEBUG OBJECT":     struct{}{},
    "DEBUG SEGFAULT":   struct{}{},
    "DECR":             struct{}{},
    "DECRBY":           struct{}{},
    "DEL":              struct{}{},
    "DISCARD":          struct{}{},
    "DUMP":             struct{}{},
    "ECHO":             struct{}{},
    "EVAL":             struct{}{},
    "EVALSHA":          struct{}{},
    "EXEC":             struct{}{},
    "EXISTS":           struct{}{},
    "EXPIRE":           struct{}{},
    "EXPIREAT":         struct{}{},
    "FLUSHALL":         struct{}{},
    "FLUSHDB":          struct{}{},
    "GET":              struct{}{},
    "GETBIT":           struct{}{},
    "GETRANGE":         struct{}{},
    "GETSET":           struct{}{},
    "HDEL":             struct{}{},
    "HEXISTS":          struct{}{},
    "HGET":             struct{}{},
    "HGETALL":          struct{}{},
    "HINCRBY":          struct{}{},
    "HINCRBYFLOAT":     struct{}{},
    "HKEYS":            struct{}{},
    "HLEN":             struct{}{},
    "HMGET":            struct{}{},
    "HMSET":            struct{}{},
    "HSCAN":            struct{}{},
    "HSET":             struct{}{},
    "HSETINX":          struct{}{},
    "HVALS":            struct{}{},
    "INCR":             struct{}{},
    "INCRBY":           struct{}{},
    "INCRBYFLOAT":      struct{}{},
    "INFO":             struct{}{},
    "KEYS":             struct{}{},
    "LASTSAVE":         struct{}{},
    "LINDEX":           struct{}{},
    "LINSERT":          struct{}{},
    "LLEN":             struct{}{},
    "LPOP":             struct{}{},
    "LPUSH":            struct{}{},
    "LPUSHX":           struct{}{},
    "LRANGE":           struct{}{},
    "LREM":             struct{}{},
    "LSET":             struct{}{},
    "LTRIM":            struct{}{},
    "MGET":             struct{}{},
    "MIGRATE":          struct{}{},
    "MONITOR":          struct{}{},
    "MOVE":             struct{}{},
    "MSET":             struct{}{},
    "MSETNX":           struct{}{},
    "MULTI":            struct{}{},
    "OBJECT":           struct{}{},
    "PERSIST":          struct{}{},
    "PEXPIRE":          struct{}{},
    "PEXPIREAT":        struct{}{},
    "PFADD":            struct{}{},
    "PFCOUNT":          struct{}{},
    "PFMERGE":          struct{}{},
    "PING":             struct{}{},
    "PSETEX":           struct{}{},
    "PSUBSCRIBE":       struct{}{},
    "PTTL":             struct{}{},
    "PUBLISH":          struct{}{},
    "PUBSUB":           struct{}{},
    "PUNSUBSCRIBE":     struct{}{},
    "QUIT":             struct{}{},
    "RANDOMKEY":        struct{}{},
    "RENAME":           struct{}{},
    "RENAMENX":         struct{}{},
    "RESTORE":          struct{}{},
    "RPOP":             struct{}{},
    "RPOPLPUSH":        struct{}{},
    "RPUSH":            struct{}{},
    "RPUSHX":           struct{}{},
    "SADD":             struct{}{},
    "SAVE":             struct{}{},
    "SCAN":             struct{}{},
    "SCARD":            struct{}{},
    "SCRIPT EXISTS":    struct{}{},
    "SCRIPT FLUSH":     struct{}{},
    "SCRIPT KILL":      struct{}{},
    "SCRIPT LOAD":      struct{}{},
    "SDIFF":            struct{}{},
    "SDIFFSTORE":       struct{}{},
    "SELECT":           struct{}{},
    "SET":              struct{}{},
    "SETBIT":           struct{}{},
    "SETEX":            struct{}{},
    "SETNX":            struct{}{},
    "SETRANGE":         struct{}{},
    "SHUTDOWN":         struct{}{},
    "SINTER":           struct{}{},
    "SINTERSTORE":      struct{}{},
    "SISMEMBER":        struct{}{},
    "SLAVEOF":          struct{}{},
    "SLOWLOG":          struct{}{},
    "SMEMBERS":         struct{}{},
    "SMOVE":            struct{}{},
    "SORT":             struct{}{},
    "SPOP":             struct{}{},
    "SRANDMEMBER":      struct{}{},
    "SREM":             struct{}{},
    "SSCAN":            struct{}{},
    "STRLEN":           struct{}{},
    "SUBSCRIBE":        struct{}{},
    "SUNION":           struct{}{},
    "SUNIONSTORE":      struct{}{},
    "SYNC":             struct{}{},
    "TIME":             struct{}{},
    "TTL":              struct{}{},
    "TYPE":             struct{}{},
    "UNSUBSCRIBE":      struct{}{},
    "UNWATCH":          struct{}{},
    "WATCH":            struct{}{},
    "ZADD":             struct{}{},
    "ZCARD":            struct{}{},
    "ZCOUNT":           struct{}{},
    "ZINCRBY":          struct{}{},
    "ZINTERSTORE":      struct{}{},
    "ZRANGE":           struct{}{},
    "ZRANGEBYSCORE":    struct{}{},
    "ZRANK":            struct{}{},
    "ZREM":             struct{}{},
    "ZREMRANGEBYLEX":   struct{}{},
    "ZREMRANGEBYRANK":  struct{}{},
    "ZREMRANGEBYSCORE": struct{}{},
    "ZREVRANGE":        struct{}{},
    "ZREVRANGEBYSCORE": struct{}{},
    "ZREVRANK":         struct{}{},
    "ZSCAN":            struct{}{},
    "ZSCORE":           struct{}{},
    "ZUNIONSTORE":      struct{}{},
}

var redisTransactionsMap = make(map[TcpTuple]*RedisTransaction, TransactionsHashSize)

func (stream *RedisStream) PrepareForNewMessage() {
    stream.message.NumberOfBulks = 0
    stream.message.Bulks = []string{}
    stream.message.IsRequest = false
}

func redisMessageParser(s *RedisStream) (bool, bool) {

    var err error
    var value string
    m := s.message

    for s.parseOffset < len(s.data) {

        if s.data[s.parseOffset] == '*' {
            //Multi Bulk Message

            line, off := readLine(s.data, s.parseOffset)

            if len(line) == 3 && line[1] == '-' && line[2] == '1' {
                //NULL Multi Bulk
                s.parseOffset = off
                value = "nil"
            } else {

                m.NumberOfBulks, err = strconv.ParseInt(line[1:], 10, 64)

                if err != nil {
                    ERR("Failed to read number of bulk messages: %s", err)
                    return false, false
                }
                s.parseOffset = off
                m.Bulks = []string{}

                continue
            }

        } else if s.data[s.parseOffset] == '$' {
            // Bulk Reply

            line, off := readLine(s.data, s.parseOffset)

            if len(line) == 3 && line[1] == '-' && line[2] == '1' {
                // NULL Bulk Reply
                value = "nil"
                s.parseOffset = off
            } else {
                length, err := strconv.ParseInt(line[1:], 10, 64)
                if err != nil {
                    ERR("Failed to read bulk message: %s", err)
                    return false, false
                }

                s.parseOffset = off

                line, off = readLine(s.data, s.parseOffset)

                if int64(len(line)) != length {
                    ERR("Wrong length of data: %d instead of %d", len(line), length)
                    return false, false
                }
                value = line
                s.parseOffset = off
            }

        } else if s.data[s.parseOffset] == ':' {
            // Integer reply
            line, off := readLine(s.data, s.parseOffset)
            n, err := strconv.ParseInt(line[1:], 10, 64)

            if err != nil {
                ERR("Failed to read integer reply: %s", err)
                return false, false
            }
            value = string(n)
            s.parseOffset = off

        } else if s.data[s.parseOffset] == '+' {
            // Status Reply
            line, off := readLine(s.data, s.parseOffset)

            value = line[1:]
            s.parseOffset = off
        } else if s.data[s.parseOffset] == '-' {
            // Error Reply
            line, off := readLine(s.data, s.parseOffset)

            value = line[1:]
            s.parseOffset = off
        } else {
            DEBUG("redis", "Unexpected message starting with %s", s.data[s.parseOffset:])
            return false, false
        }

        // add value
        if m.NumberOfBulks > 0 {
            m.NumberOfBulks = m.NumberOfBulks - 1
            m.Bulks = append(m.Bulks, value)

            if len(m.Bulks) == 1 {
                // check if it's a command
                if isRedisCommand(value) {
                    m.IsRequest = true
                }
            }

            if m.NumberOfBulks == 0 {
                // the last bulk received
                m.Message = strings.Join(m.Bulks, " ")
                return true, true
            }
        } else {
            m.Message = value
            return true, true
        }

    }   //end for

    return true, false
}

func readLine(data []byte, offset int) (string, int) {
    q := bytes.Index(data[offset:], []byte("\r\n"))
    return string(data[offset : offset+q]), offset + q + 2
}

func ParseRedis(pkt *Packet, tcp *TcpStream, dir uint8) {
    defer RECOVER("ParseRedis exception")

    if tcp.redisData[dir] == nil {
        tcp.redisData[dir] = &RedisStream{
            tcpStream: tcp,
            data:      pkt.payload,
            message:   &RedisMessage{Ts: pkt.ts},
        }
    } else {
        // concatenate bytes
        tcp.redisData[dir].data = append(tcp.redisData[dir].data, pkt.payload...)
    }

    stream := tcp.redisData[dir]
    if stream.message == nil {
        stream.message = &RedisMessage{Ts: pkt.ts}
    }

    ok, complete := redisMessageParser(tcp.redisData[dir])

    if !ok {
        // drop this tcp stream. Will retry parsing with the next
        // segment in it
        tcp.redisData[dir] = nil
        return
    }

    if complete {

        if stream.message.IsRequest {
            DEBUG("redis", "REDIS request message: %s", stream.message.Message)
        } else {
            DEBUG("redis", "REDIS response message: %s", stream.message.Message)
        }

        // all ok, go to next level
        handleRedis(stream.message, tcp, dir)

        // and reset message
        stream.PrepareForNewMessage()
    }

}

func isRedisCommand(key string) bool {
    _, exists := RedisCommands[key]
    return exists
}

func handleRedis(m *RedisMessage, tcp *TcpStream,
    dir uint8) {

    m.Stream_id = tcp.id
    m.Tuple = tcp.tuple
    m.Direction = dir
    m.CmdlineTuple = procWatcher.FindProcessesTuple(tcp.tuple)

    if m.IsRequest {
        receivedRedisRequest(m)
    } else {
        receivedRedisResponse(m)
    }
}

func receivedRedisRequest(msg *RedisMessage) {
    // Add it to the HT
    tuple := TcpTuple{
        Src_ip:    msg.Tuple.Src_ip,
        Dst_ip:    msg.Tuple.Dst_ip,
        Src_port:  msg.Tuple.Src_port,
        Dst_port:  msg.Tuple.Dst_port,
        stream_id: msg.Stream_id,
    }

    trans := redisTransactionsMap[tuple]
    if trans != nil {
        if len(trans.Redis) != 0 {
            WARN("Two requests without a Response. Dropping old request")
        }
    } else {
        trans = &RedisTransaction{Type: "redis", tuple: tuple}
        redisTransactionsMap[tuple] = trans
    }

    trans.Redis = bson.M{
        "request": msg.Message,
    }
    trans.Request_raw = msg.Message

    trans.cmdline = msg.CmdlineTuple
    trans.ts = msg.Ts
    trans.Ts = int64(trans.ts.UnixNano() / 1000) // transactions have microseconds resolution
    trans.JsTs = msg.Ts
    trans.Src = Endpoint{
        Ip:   Ipv4_Ntoa(tuple.Src_ip),
        Port: tuple.Src_port,
        Proc: string(msg.CmdlineTuple.Src),
    }
    trans.Dst = Endpoint{
        Ip:   Ipv4_Ntoa(tuple.Dst_ip),
        Port: tuple.Dst_port,
        Proc: string(msg.CmdlineTuple.Dst),
    }

    if trans.timer != nil {
        trans.timer.Stop()
    }
    trans.timer = time.AfterFunc(TransactionTimeout, func() { trans.Expire() })

}

func (trans *RedisTransaction) Expire() {

    // remove from map
    delete(redisTransactionsMap, trans.tuple)
}

func receivedRedisResponse(msg *RedisMessage) {

    tuple := TcpTuple{
        Src_ip:    msg.Tuple.Src_ip,
        Dst_ip:    msg.Tuple.Dst_ip,
        Src_port:  msg.Tuple.Src_port,
        Dst_port:  msg.Tuple.Dst_port,
        stream_id: msg.Stream_id,
    }
    trans := redisTransactionsMap[tuple]
    if trans == nil {
        WARN("Response from unknown transaction. Ignoring.")
        return
    }
    // check if the request was received
    if len(trans.Redis) == 0 {
        WARN("Response from unknown transaction. Ignoring.")
        return

    }

    trans.Redis["response"] = msg.Message

    trans.Response_raw = msg.Message

    trans.ResponseTime = int32(msg.Ts.Sub(trans.ts).Nanoseconds() / 1e6) // resp_time in milliseconds

    err := Publisher.PublishRedisTransaction(trans)
    if err != nil {
        WARN("Publish failure: %s", err)
    }

    DEBUG("redis", "Redis transaction completed: %s", trans.Redis)

    // remove from map
    delete(redisTransactionsMap, trans.tuple)
    if trans.timer != nil {
        trans.timer.Stop()
    }

}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// kube_checks contains functions which poll Kubernetes objects until
// they get into the state desired by the caller or time out.

package test

import (
	"context"
	"fmt"
	"time"

	"go.opencensus.io/trace"
	corev1 "k8s.io/api/core/v1"
	apiv1beta1 "k8s.io/api/extensions/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
)

const (
	interval = 1 * time.Second
	timeout  = 5 * time.Minute
)

// WaitForDeploymentState polls the status of the Deployment called name
// from client every interval until inState returns `true` indicating it
// is done, returns an error or timeout. desc will be used to name the metric
// that is emitted to track how long it took for name to get into the state checked by inState.
func WaitForDeploymentState(client *KubeClient, name string, inState func(d *apiv1beta1.Deployment) (bool, error), desc string, namespace string) error {
	d := client.Kube.ExtensionsV1beta1().Deployments(namespace)
	metricName := fmt.Sprintf("WaitForDeploymentState/%s/%s", name, desc)
	_, span := trace.StartSpan(context.Background(), metricName)
	defer span.End()

	return wait.PollImmediate(interval, timeout, func() (bool, error) {
		d, err := d.Get(name, metav1.GetOptions{})
		if err != nil {
			return true, err
		}
		return inState(d)
	})
}

// WaitForPodListState polls the status of the PodList
// from client every interval until inState returns `true` indicating it
// is done, returns an error or timeout. desc will be used to name the metric
// that is emitted to track how long it took to get into the state checked by inState.
func WaitForPodListState(client *KubeClient, inState func(p *corev1.PodList) (bool, error), desc string, namespace string) error {
	p := client.Kube.CoreV1().Pods(namespace)
	metricName := fmt.Sprintf("WaitForPodListState/%s", desc)
	_, span := trace.StartSpan(context.Background(), metricName)
	defer span.End()

	return wait.PollImmediate(interval, timeout, func() (bool, error) {
		p, err := p.List(metav1.ListOptions{})
		if err != nil {
			return true, err
		}
		return inState(p)
	})
}
<commit_msg>Move duration as an arg for WaitForDeploymentState (#113)<commit_after>/*
Copyright 2018 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// kube_checks contains functions which poll Kubernetes objects until
// they get into the state desired by the caller or time out.

package test

import (
	"context"
	"fmt"
	"time"

	"go.opencensus.io/trace"
	corev1 "k8s.io/api/core/v1"
	apiv1beta1 "k8s.io/api/extensions/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
)

const (
	interval   = 1 * time.Second
	podTimeout = 5 * time.Minute
)

// WaitForDeploymentState polls the status of the Deployment called name
// from client every interval until inState returns `true` indicating it
// is done, returns an error or timeout. desc will be used to name the metric
// that is emitted to track how long it took for name to get into the state checked by inState.
func WaitForDeploymentState(client *KubeClient, name string, inState func(d *apiv1beta1.Deployment) (bool, error), desc string, namespace string, timeout time.Duration) error {
	d := client.Kube.ExtensionsV1beta1().Deployments(namespace)
	metricName := fmt.Sprintf("WaitForDeploymentState/%s/%s", name, desc)
	_, span := trace.StartSpan(context.Background(), metricName)
	defer span.End()

	return wait.PollImmediate(interval, timeout, func() (bool, error) {
		d, err := d.Get(name, metav1.GetOptions{})
		if err != nil {
			return true, err
		}
		return inState(d)
	})
}

// WaitForPodListState polls the status of the PodList
// from client every interval until inState returns `true` indicating it
// is done, returns an error or timeout. desc will be used to name the metric
// that is emitted to track how long it took to get into the state checked by inState.
func WaitForPodListState(client *KubeClient, inState func(p *corev1.PodList) (bool, error), desc string, namespace string) error {
	p := client.Kube.CoreV1().Pods(namespace)
	metricName := fmt.Sprintf("WaitForPodListState/%s", desc)
	_, span := trace.StartSpan(context.Background(), metricName)
	defer span.End()

	return wait.PollImmediate(interval, podTimeout, func() (bool, error) {
		p, err := p.List(metav1.ListOptions{})
		if err != nil {
			return true, err
		}
		return inState(p)
	})
}
<|endoftext|>


<commit_before>package main

type Activating struct {
	partialHandler
	*ability
	receiver *unit
}

// NewActivating returns a activating handler
func NewActivating(performer, receiver *unit, ability *ability) *Activating {
	return &Activating{
		partialHandler: partialHandler{
			unit:           performer,
			expirationTime: performer.after(ability.activationDuration),
		},
		ability:  ability,
		receiver: receiver,
	}
}

// OnAttach checks requirements
func (a *Activating) OnAttach() {
	a.AddEventHandler(a, EventDead)
	a.AddEventHandler(a, EventDisableInterrupt)
	a.AddEventHandler(a, EventGameTick)
	a.AddEventHandler(a, EventResourceDecreased)
	for ha := range a.handlers {
		switch ha.(type) {
		case *Activating:
			a.detachHandler(a)
			return
		}
	}
	if err := a.checkRequirements(a.unit, a.receiver); err != nil {
		a.detachHandler(a)
		return
	}
	if a.isExpired() {
		a.perform()
		return
	}
	a.publish(message{
	// TODO pack message
	})
}

// OnDetach removes the EventHandlers
func (a *Activating) OnDetach() {
	a.RemoveEventHandler(a, EventDead)
	a.RemoveEventHandler(a, EventDisableInterrupt)
	a.RemoveEventHandler(a, EventGameTick)
	a.RemoveEventHandler(a, EventResourceDecreased)
}

// HandleEvent handles the event
func (a *Activating) HandleEvent(e Event) {
	switch e {
	case EventDead:
		a.detachHandler(a)
	case EventDisableInterrupt:
		a.perform()
	case EventGameTick:
		a.perform()
	case EventResourceDecreased:
		a.perform()
	}
}

// perform performs the ability
func (a *Activating) perform() {
	if err := a.checkRequirements(a.unit, a.receiver); err != nil {
		a.publish(message{
		// TODO pack message
		})
		a.detachHandler(a)
	}
	if !a.isExpired() {
		return
	}
	a.publish(message{
	// TODO pack message
	})
	a.ability.perform(a.unit, a.receiver)
	a.detachHandler(a)
}
<commit_msg>Add comments<commit_after>package main

type Activating struct {
	partialHandler
	*ability
	receiver *unit
}

// NewActivating returns a activating handler
func NewActivating(performer, receiver *unit, ability *ability) *Activating {
	return &Activating{
		partialHandler: partialHandler{
			unit:           performer,
			expirationTime: performer.after(ability.activationDuration),
		},
		ability:  ability,
		receiver: receiver,
	}
}

// OnAttach checks requirements
func (a *Activating) OnAttach() {
	a.AddEventHandler(a, EventDead)
	a.AddEventHandler(a, EventDisableInterrupt)
	a.AddEventHandler(a, EventGameTick)
	a.AddEventHandler(a, EventResourceDecreased)
	for ha := range a.handlers {
		switch ha.(type) {
		case *Activating:
			a.detachHandler(a)
			return
		}
	}
	if err := a.checkRequirements(a.unit, a.receiver); err != nil {
		a.detachHandler(a)
		return
	}
	if a.isExpired() {
		a.perform()
		return
	}
	a.publish(message{
	// TODO pack message
	})
}

// OnDetach removes the EventHandlers
func (a *Activating) OnDetach() {
	a.RemoveEventHandler(a, EventDead)
	a.RemoveEventHandler(a, EventDisableInterrupt)
	a.RemoveEventHandler(a, EventGameTick)
	a.RemoveEventHandler(a, EventResourceDecreased)
}

// HandleEvent handles the event
func (a *Activating) HandleEvent(e Event) {
	switch e {
	case EventDead:
		a.detachHandler(a)
	case EventDisableInterrupt:
		a.perform()
	case EventGameTick:
		a.perform()
	case EventResourceDecreased:
		a.perform()
	}
}

// perform performs the ability
func (a *Activating) perform() {
	if err := a.checkRequirements(a.unit, a.receiver); err != nil {
		a.publish(message{
		// TODO pack message
		})
		a.detachHandler(a)
	}
	if !a.isExpired() {
		return
	}
	a.publish(message{
	// TODO pack message
	})
	// TODO consume health
	// TODO consume mana
	a.ability.perform(a.unit, a.receiver)
	a.detachHandler(a)
}
<|endoftext|>


<commit_before>package discord

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"time"

	"sync"

	"github.com/influxdata/kapacitor/alert"
	khttp "github.com/influxdata/kapacitor/http"
	"github.com/influxdata/kapacitor/keyvalue"
	"github.com/influxdata/kapacitor/tlsconfig"
	"github.com/pkg/errors"
)

type Diagnostic interface {
	WithContext(ctx ...keyvalue.T) Diagnostic

	InsecureSkipVerify()

	Error(msg string, err error)
}

type Workspace struct {
	mu     sync.RWMutex
	config Config
	client *http.Client
}

func NewWorkspace(c Config) (*Workspace, error) {
	tlsConfig, err := tlsconfig.Create(c.SSLCA, c.SSLCert, c.SSLKey, c.InsecureSkipVerify)
	if err != nil {
		return nil, err
	}

	cl := &http.Client{
		Transport: khttp.NewDefaultTransportWithTLS(tlsConfig),
	}

	return &Workspace{
		config: c,
		client: cl,
	}, nil
}

func (w *Workspace) Config() Config {
	w.mu.RLock()
	defer w.mu.RUnlock()
	return w.config
}

func (w *Workspace) Client() *http.Client {
	w.mu.RLock()
	defer w.mu.RUnlock()
	return w.client
}

func (w *Workspace) Update(c Config) error {
	tlsConfig, err := tlsconfig.Create(c.SSLCA, c.SSLCert, c.SSLKey, c.InsecureSkipVerify)
	if err != nil {
		return err
	}

	cl := &http.Client{
		Transport: khttp.NewDefaultTransportWithTLS(tlsConfig),
	}

	w.client = cl
	w.config = c

	return nil
}

type Service struct {
	mu         sync.RWMutex
	workspaces map[string]*Workspace
	diag       Diagnostic
}

func NewService(confs []Config, d Diagnostic) (*Service, error) {
	s := &Service{
		diag:       d,
		workspaces: make(map[string]*Workspace),
	}

	if len(confs) == 1 {
		confs[0].Default = true
	}
	for _, c := range confs {

		if c.InsecureSkipVerify {
			s.diag.InsecureSkipVerify()
		}

		w, err := NewWorkspace(c)
		if err != nil {
			return nil, err
		}
		s.workspaces[c.Workspace] = w

		// We'll stash the default workspace with the empty string as a key.
		// Either there's a single config with no workspace name, or else
		// we have multiple configs that all have names.
		if c.Default && c.Workspace != "" {
			s.workspaces[""] = s.workspaces[c.Workspace]
		}

	}

	return s, nil
}

func (s *Service) Open() error {
	return nil
}

func (s *Service) Close() error {
	return nil
}

func (s *Service) workspace(wid string) (*Workspace, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if len(s.workspaces) == 0 {
		return &Workspace{}, errors.New("no discord configuration found")
	}
	v, ok := s.workspaces[wid]
	if !ok {
		return &Workspace{}, errors.New("workspace id not found")
	}
	return v, nil
}

func (s *Service) config(wid string) (Config, error) {
	w, err := s.workspace(wid)
	if err != nil {
		return Config{}, err
	}

	return w.Config(), nil
}

func (s *Service) defaultConfig() Config {
	conf, _ := s.config("")
	return conf
}

func (s *Service) client(wid string) (*http.Client, error) {
	w, err := s.workspace(wid)
	if err != nil {
		return &http.Client{}, err
	}
	return w.Client(), nil
}

func (s *Service) Update(newConfigs []interface{}) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	for _, v := range newConfigs {
		if conf, ok := v.(Config); ok {
			if conf.InsecureSkipVerify {
				s.diag.InsecureSkipVerify()
			}
			_, ok := s.workspaces[conf.Workspace]
			if !ok {
				w, err := NewWorkspace(conf)
				s.workspaces[conf.Workspace] = w
				if err != nil {
					return err
				}
			} else {
				s.workspaces[conf.Workspace].Update(conf)
			}

			// We'll stash the default workspace with the empty string as a key.
			// Either there's a single config with no workspace name, or else
			// we have multiple configs that all have names.
			if conf.Default && conf.Workspace != "" {
				s.workspaces[""] = s.workspaces[conf.Workspace]
			}
		} else {
			return fmt.Errorf("expected config object to be of type %T, got %T", v, conf)
		}
	}
	return nil
}

func (s *Service) Global() bool {
	return s.defaultConfig().Global
}

func (s *Service) StateChangesOnly() bool {
	return s.defaultConfig().Global
}

// Discord rich embed info
type embed struct {
	Color       int    `json:"color"`
	Title       string `json:"title"`
	Description string `json:"description"`
	Timestamp   string `json:"timestamp"`
}

type testOptions struct {
	Workspace  string      `json:"workspace"`
	Message    string      `json:"message"`
	Level      alert.Level `json:"level"`
	Username   string      `json:"username"`
	AvatarURL  string      `json:"avatar-url"`
	EmbedTitle string      `json:"embed-title"`
	Timestamp  bool        `json:"timestamp"`
	Time       time.Time   `json:"time-val"`
}

func (s *Service) TestOptions() interface{} {
	c, _ := s.config("")
	t, _ := time.Parse(time.RFC3339, "1970-01-01T00:00:01Z")
	return &testOptions{
		Workspace:  c.Workspace,
		Message:    "test discord message",
		Level:      alert.Info,
		AvatarURL:  "https://influxdata.github.io/branding/img/downloads/influxdata-logo--symbol--pool-alpha.png",
		Username:   "Kapacitor",
		EmbedTitle: "Kapacitor Alert",
		Timestamp:  true,
		Time:       t,
	}
}

func (s *Service) Test(options interface{}) error {
	o, ok := options.(*testOptions)
	if !ok {
		return fmt.Errorf("unexpected options type %T", options)
	}
	return s.Alert(o.Workspace, o.Message, o.Username, o.AvatarURL, o.EmbedTitle, o.Timestamp, o.Time, o.Level)
}

// Alert sends a message to the specified room.
func (s *Service) Alert(workspace, message, username, avatarURL, embedTitle string, timestamp bool, time time.Time, level alert.Level) error {
	url, post, err := s.preparePost(workspace, message, username, avatarURL, embedTitle, timestamp, time, level)
	if err != nil {
		return err
	}

	client, err := s.client(workspace)
	if err != nil {
		return err
	}

	resp, err := client.Post(url, "application/json", post)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusNoContent {
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return err
		}
		type response struct {
			Error string `json:"error"`
		}
		r := &response{Error: fmt.Sprintf("failed to understand Discord response. code: %d content: %s", resp.StatusCode, string(body))}
		b := bytes.NewReader(body)
		dec := json.NewDecoder(b)
		dec.Decode(r)
		return errors.New(r.Error)
	}
	return nil
}

type HandlerConfig struct {
	// Discord workspace ID to use when posting to webhook
	// If empty uses the default config
	Workspace string `mapstructure:"workspace"`
	// Whether to display timestamp in the footer of the embed
	// If empty uses the default config
	Timestamp bool `mapstructure:"timestamp"`
	// Username of webhook
	// If empty uses the default config
	Username string `mapstructure:"username"`
	// URL of webhook's avatar
	// If empty uses the default config
	AvatarURL string `mapstructure:"avatar-url"`
	// Embed title
	// If empty uses the default config
	EmbedTitle string `mapstructure:"embed-title"`
}

func (s *Service) preparePost(workspace, message, username, avatarURL, embedTitle string, timestamp bool, timeVal time.Time, level alert.Level) (string, io.Reader, error) {
	c, err := s.config(workspace)
	if err != nil {
		return "", nil, err
	}
	if !c.Enabled {
		return "", nil, errors.New("service is not enabled")
	}
	var color int
	switch level {
	case alert.Critical:
		color = 0xF95F53 // #F95F53
	case alert.Warning:
		color = 0xF48D38 // #F48D38
	default:
		color = 0x7A65F2 // #7A65F2
	}
	var timeStr string
	if timestamp {
		timeStr = timeVal.Format(time.RFC3339)
	}
	a := embed{
		Description: message,
		Color:       color,
		Title:       embedTitle,
		Timestamp:   timeStr,
	}
	postData := make(map[string]interface{})
	if username == "" {
		username = c.Username
	}
	postData["username"] = username
	if avatarURL == "" {
		avatarURL = c.AvatarURL
	}
	postData["avatar_url"] = avatarURL
	postData["embeds"] = []embed{a}

	var post bytes.Buffer
	enc := json.NewEncoder(&post)
	err = enc.Encode(postData)
	if err != nil {
		return "", nil, err
	}

	return c.URL, &post, nil
}

type handler struct {
	s    *Service
	c    HandlerConfig
	diag Diagnostic
}

func (s *Service) Handler(c HandlerConfig, ctx ...keyvalue.T) alert.Handler {
	return &handler{
		s:    s,
		c:    c,
		diag: s.diag.WithContext(ctx...),
	}
}

func (h *handler) Handle(event alert.Event) {
	if err := h.s.Alert(
		h.c.Workspace,
		event.State.Message,
		h.c.Username,
		h.c.AvatarURL,
		h.c.EmbedTitle,
		h.c.Timestamp,
		event.State.Time,
		event.State.Level,
	); err != nil {
		h.diag.Error("failed to send event", err)
	}
}
<commit_msg>Allow embedTitle to be parsed as a template<commit_after>package discord

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	text "text/template"
	"time"

	"sync"

	"github.com/influxdata/kapacitor/alert"
	khttp "github.com/influxdata/kapacitor/http"
	"github.com/influxdata/kapacitor/keyvalue"
	"github.com/influxdata/kapacitor/tlsconfig"
	"github.com/pkg/errors"
)

type Diagnostic interface {
	WithContext(ctx ...keyvalue.T) Diagnostic
	TemplateError(err error, kv keyvalue.T)
	InsecureSkipVerify()

	Error(msg string, err error)
}

type Workspace struct {
	mu     sync.RWMutex
	config Config
	client *http.Client
}

func NewWorkspace(c Config) (*Workspace, error) {
	tlsConfig, err := tlsconfig.Create(c.SSLCA, c.SSLCert, c.SSLKey, c.InsecureSkipVerify)
	if err != nil {
		return nil, err
	}

	cl := &http.Client{
		Transport: khttp.NewDefaultTransportWithTLS(tlsConfig),
	}

	return &Workspace{
		config: c,
		client: cl,
	}, nil
}

func (w *Workspace) Config() Config {
	w.mu.RLock()
	defer w.mu.RUnlock()
	return w.config
}

func (w *Workspace) Client() *http.Client {
	w.mu.RLock()
	defer w.mu.RUnlock()
	return w.client
}

func (w *Workspace) Update(c Config) error {
	tlsConfig, err := tlsconfig.Create(c.SSLCA, c.SSLCert, c.SSLKey, c.InsecureSkipVerify)
	if err != nil {
		return err
	}

	cl := &http.Client{
		Transport: khttp.NewDefaultTransportWithTLS(tlsConfig),
	}

	w.client = cl
	w.config = c

	return nil
}

type Service struct {
	mu         sync.RWMutex
	workspaces map[string]*Workspace
	diag       Diagnostic
}

func NewService(confs []Config, d Diagnostic) (*Service, error) {
	s := &Service{
		diag:       d,
		workspaces: make(map[string]*Workspace),
	}

	if len(confs) == 1 {
		confs[0].Default = true
	}
	for _, c := range confs {

		if c.InsecureSkipVerify {
			s.diag.InsecureSkipVerify()
		}

		w, err := NewWorkspace(c)
		if err != nil {
			return nil, err
		}
		s.workspaces[c.Workspace] = w

		// We'll stash the default workspace with the empty string as a key.
		// Either there's a single config with no workspace name, or else
		// we have multiple configs that all have names.
		if c.Default && c.Workspace != "" {
			s.workspaces[""] = s.workspaces[c.Workspace]
		}

	}

	return s, nil
}

func (s *Service) Open() error {
	return nil
}

func (s *Service) Close() error {
	return nil
}

func (s *Service) workspace(wid string) (*Workspace, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if len(s.workspaces) == 0 {
		return &Workspace{}, errors.New("no discord configuration found")
	}
	v, ok := s.workspaces[wid]
	if !ok {
		return &Workspace{}, errors.New("workspace id not found")
	}
	return v, nil
}

func (s *Service) config(wid string) (Config, error) {
	w, err := s.workspace(wid)
	if err != nil {
		return Config{}, err
	}

	return w.Config(), nil
}

func (s *Service) defaultConfig() Config {
	conf, _ := s.config("")
	return conf
}

func (s *Service) client(wid string) (*http.Client, error) {
	w, err := s.workspace(wid)
	if err != nil {
		return &http.Client{}, err
	}
	return w.Client(), nil
}

func (s *Service) Update(newConfigs []interface{}) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	for _, v := range newConfigs {
		if conf, ok := v.(Config); ok {
			if conf.InsecureSkipVerify {
				s.diag.InsecureSkipVerify()
			}
			_, ok := s.workspaces[conf.Workspace]
			if !ok {
				w, err := NewWorkspace(conf)
				s.workspaces[conf.Workspace] = w
				if err != nil {
					return err
				}
			} else {
				s.workspaces[conf.Workspace].Update(conf)
			}

			// We'll stash the default workspace with the empty string as a key.
			// Either there's a single config with no workspace name, or else
			// we have multiple configs that all have names.
			if conf.Default && conf.Workspace != "" {
				s.workspaces[""] = s.workspaces[conf.Workspace]
			}
		} else {
			return fmt.Errorf("expected config object to be of type %T, got %T", v, conf)
		}
	}
	return nil
}

func (s *Service) Global() bool {
	return s.defaultConfig().Global
}

func (s *Service) StateChangesOnly() bool {
	return s.defaultConfig().Global
}

// Discord rich embed info
type embed struct {
	Color       int    `json:"color"`
	Title       string `json:"title"`
	Description string `json:"description"`
	Timestamp   string `json:"timestamp"`
}

type testOptions struct {
	Workspace  string      `json:"workspace"`
	Message    string      `json:"message"`
	Level      alert.Level `json:"level"`
	Username   string      `json:"username"`
	AvatarURL  string      `json:"avatar-url"`
	EmbedTitle string      `json:"embed-title"`
	Timestamp  bool        `json:"timestamp"`
	Time       time.Time   `json:"time-val"`
}

func (s *Service) TestOptions() interface{} {
	c, _ := s.config("")
	t, _ := time.Parse(time.RFC3339, "1970-01-01T00:00:01Z")
	return &testOptions{
		Workspace:  c.Workspace,
		Message:    "test discord message",
		Level:      alert.Info,
		AvatarURL:  "https://influxdata.github.io/branding/img/downloads/influxdata-logo--symbol--pool-alpha.png",
		Username:   "Kapacitor",
		EmbedTitle: "Kapacitor Alert",
		Timestamp:  true,
		Time:       t,
	}
}

func (s *Service) Test(options interface{}) error {
	o, ok := options.(*testOptions)
	if !ok {
		return fmt.Errorf("unexpected options type %T", options)
	}
	return s.Alert(o.Workspace, o.Message, o.Username, o.AvatarURL, o.EmbedTitle, o.Timestamp, o.Time, o.Level)
}

// Alert sends a message to the specified room.
func (s *Service) Alert(workspace, message, username, avatarURL, embedTitle string, timestamp bool, time time.Time, level alert.Level) error {
	url, post, err := s.preparePost(workspace, message, username, avatarURL, embedTitle, timestamp, time, level)
	if err != nil {
		return err
	}

	client, err := s.client(workspace)
	if err != nil {
		return err
	}

	resp, err := client.Post(url, "application/json", post)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusNoContent {
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return err
		}
		type response struct {
			Error string `json:"error"`
		}
		r := &response{Error: fmt.Sprintf("failed to understand Discord response. code: %d content: %s", resp.StatusCode, string(body))}
		b := bytes.NewReader(body)
		dec := json.NewDecoder(b)
		dec.Decode(r)
		return errors.New(r.Error)
	}
	return nil
}

type HandlerConfig struct {
	// Discord workspace ID to use when posting to webhook
	// If empty uses the default config
	Workspace string `mapstructure:"workspace"`
	// Whether to display timestamp in the footer of the embed
	// If empty uses the default config
	Timestamp bool `mapstructure:"timestamp"`
	// Username of webhook
	// If empty uses the default config
	Username string `mapstructure:"username"`
	// URL of webhook's avatar
	// If empty uses the default config
	AvatarURL string `mapstructure:"avatar-url"`
	// Embed title
	// If empty uses the default config
	EmbedTitle string `mapstructure:"embed-title"`
}

func (s *Service) preparePost(workspace, message, username, avatarURL, embedTitle string, timestamp bool, timeVal time.Time, level alert.Level) (string, io.Reader, error) {
	c, err := s.config(workspace)
	if err != nil {
		return "", nil, err
	}
	if !c.Enabled {
		return "", nil, errors.New("service is not enabled")
	}
	var color int
	switch level {
	case alert.Critical:
		color = 0xF95F53 // #F95F53
	case alert.Warning:
		color = 0xF48D38 // #F48D38
	default:
		color = 0x7A65F2 // #7A65F2
	}
	var timeStr string
	if timestamp {
		timeStr = timeVal.Format(time.RFC3339)
	}
	a := embed{
		Description: message,
		Color:       color,
		Title:       embedTitle,
		Timestamp:   timeStr,
	}
	postData := make(map[string]interface{})
	if username == "" {
		username = c.Username
	}
	postData["username"] = username
	if avatarURL == "" {
		avatarURL = c.AvatarURL
	}
	postData["avatar_url"] = avatarURL
	postData["embeds"] = []embed{a}

	var post bytes.Buffer
	enc := json.NewEncoder(&post)
	err = enc.Encode(postData)
	if err != nil {
		return "", nil, err
	}

	return c.URL, &post, nil
}

type handler struct {
	s    *Service
	c    HandlerConfig
	diag Diagnostic

	embedTitleTmpl *text.Template
}

func (s *Service) Handler(c HandlerConfig, ctx ...keyvalue.T) (alert.Handler, error) {
	ettmpl, err := text.New("embedTitle").Parse(c.EmbedTitle)
	if err != nil {
		return nil, err
	}
	return &handler{
		s:              s,
		c:              c,
		diag:           s.diag.WithContext(ctx...),
		embedTitleTmpl: ettmpl,
	}, nil
}

func (h *handler) Handle(event alert.Event) {
	td := event.TemplateData()
	var buf bytes.Buffer

	if err := h.embedTitleTmpl.Execute(&buf, td); err != nil {
		h.diag.TemplateError(err, keyvalue.KV("embedTitle", h.c.EmbedTitle))
		return
	}
	if err := h.s.Alert(
		h.c.Workspace,
		event.State.Message,
		h.c.Username,
		h.c.AvatarURL,
		buf.String(), // Parsed embedtitle template
		h.c.Timestamp,
		event.State.Time,
		event.State.Level,
	); err != nil {
		h.diag.Error("failed to send event to Discord", err)
	}
}
<|endoftext|>


<commit_before>// Package round is a command line spinner. Start one with Start.
package round

import (
	"io"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	term "golang.org/x/crypto/ssh/terminal"
)

var (
	// Stdout is a spin-safe version of os.Stdout
	Stdout io.Writer

	// Stderr is a spin-safe version of os.Stderr
	Stderr io.Writer

	// Spin Control
	out  io.Writer
	spin *spinMe
	mu   = spinit()

	// Terminal escape sequences.
	hide      = []byte{27, '[', '?', '2', '5', 'l'}
	show      = []byte{27, '[', '?', '2', '5', 'h'}
	save      = []byte{27, '[', 's'}
	clear     = []byte{27, '[', 'u', 27, '[', 'K'}
	saveHide  = append(save, hide...)
	clearShow = append(clear, show...)
)

// spinit is the init for the spin it.
func spinit() *sync.Mutex {
	o := term.IsTerminal(int(os.Stdout.Fd()))
	e := term.IsTerminal(int(os.Stderr.Fd()))
	switch {
	case o && e:
		Stdout = &blockingWriter{os.Stdout}
		Stderr = &blockingWriter{os.Stderr}
		out = os.Stdout
	case o:
		Stdout = &blockingWriter{os.Stdout}
		Stderr = os.Stderr
		out = os.Stdout
	case e:
		Stdout = os.Stdout
		Stderr = &blockingWriter{os.Stderr}
		out = os.Stderr
	default:
		Stdout = os.Stdout
		Stderr = os.Stderr
	}

	sigs := make(chan os.Signal)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigs
		Stop()
	}()

	return &sync.Mutex{}
}

// Start starts a spinner.
func Start(s Style) {
	if out == nil {
		return
	}
	if spin != nil {
		Stop()
	}
	spin = &spinMe{s.Frames[0], make(chan bool)}
	out.Write(saveHide)
	go spin.writeRound(s.Frames, time.NewTicker(s.Rate))
}

// Stop will stop and remove the spinner.
func Stop() {
	if out == nil || spin == nil {
		return
	}
	spin.stop <- true
	out.Write(clearShow)
}

type blockingWriter struct {
	out io.Writer
}

// Write writes to the underlying Writer, moving the spinner to the end of what's written.
func (w *blockingWriter) Write(p []byte) (int, error) {
	mu.Lock()
	w.out.Write(clear)
	n, err := w.out.Write(p)
	w.out.Write(append(save, spin.now...))
	mu.Unlock()
	return n, err
}

// spinMe goes right round.
type spinMe struct {
	now  string
	stop chan bool
}

// writeRound spins the spinner right round. Like a record, baby.
func (u *spinMe) writeRound(baby []string, rightRound *time.Ticker) {
	var f int
	for {
		select {
		case <-rightRound.C:
			f = (f + 1) % len(baby)
			mu.Lock()
			u.now = baby[f]
			out.Write(append(clear, u.now...))
			mu.Unlock()
		case <-u.stop:
			rightRound.Stop()
			break
		}
	}
}
<commit_msg>move thing<commit_after>// Package round is a command line spinner. Start one with Start.
package round

import (
	"io"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	term "golang.org/x/crypto/ssh/terminal"
)

var (
	// Stdout is a spin-safe version of os.Stdout
	Stdout io.Writer

	// Stderr is a spin-safe version of os.Stderr
	Stderr io.Writer

	// Spin Control
	out  io.Writer
	spin *spinMe
	mu   = spinit()

	// Terminal escape sequences.
	hide      = []byte{27, '[', '?', '2', '5', 'l'}
	show      = []byte{27, '[', '?', '2', '5', 'h'}
	save      = []byte{27, '[', 's'}
	clear     = []byte{27, '[', 'u', 27, '[', 'K'}
	saveHide  = append(save, hide...)
	clearShow = append(clear, show...)
)

// Start starts a spinner.
func Start(s Style) {
	if out == nil {
		return
	}
	if spin != nil {
		Stop()
	}
	spin = &spinMe{s.Frames[0], make(chan bool)}
	out.Write(saveHide)
	go spin.writeRound(s.Frames, time.NewTicker(s.Rate))
}

// Stop will stop and remove the spinner.
func Stop() {
	if out == nil || spin == nil {
		return
	}
	spin.stop <- true
	out.Write(clearShow)
}

type blockingWriter struct {
	out io.Writer
}

// Write writes to the underlying Writer, moving the spinner to the end of what's written.
func (w *blockingWriter) Write(p []byte) (int, error) {
	mu.Lock()
	w.out.Write(clear)
	n, err := w.out.Write(p)
	w.out.Write(append(save, spin.now...))
	mu.Unlock()
	return n, err
}

// spinMe goes right round.
type spinMe struct {
	now  string
	stop chan bool
}

// writeRound spins the spinner right round. Like a record, baby.
func (u *spinMe) writeRound(baby []string, rightRound *time.Ticker) {
	var f int
	for {
		select {
		case <-rightRound.C:
			f = (f + 1) % len(baby)
			mu.Lock()
			u.now = baby[f]
			out.Write(append(clear, u.now...))
			mu.Unlock()
		case <-u.stop:
			rightRound.Stop()
			break
		}
	}
}

// spinit is the init for the spin it.
func spinit() *sync.Mutex {
	o := term.IsTerminal(int(os.Stdout.Fd()))
	e := term.IsTerminal(int(os.Stderr.Fd()))
	switch {
	case o && e:
		Stdout = &blockingWriter{os.Stdout}
		Stderr = &blockingWriter{os.Stderr}
		out = os.Stdout
	case o:
		Stdout = &blockingWriter{os.Stdout}
		Stderr = os.Stderr
		out = os.Stdout
	case e:
		Stdout = os.Stdout
		Stderr = &blockingWriter{os.Stderr}
		out = os.Stderr
	default:
		Stdout = os.Stdout
		Stderr = os.Stderr
	}

	sigs := make(chan os.Signal)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigs
		Stop()
	}()

	return &sync.Mutex{}
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 Miquel Sabaté Solà <mikisabate@gmail.com>
// This file is licensed under the MIT license.
// See the LICENSE file.

package lib

import (
	"errors"
	"testing"

	"github.com/stretchr/testify/assert"
)

// Helper used to test custom errors. This is used throughout the test suite.
func testError(t *testing.T, err error, msg, see string) {
	assert.NotNil(t, err)
	e, ce := err.(*Error)
	assert.True(t, ce)
	assert.Equal(t, e.message, msg)
	assert.Equal(t, e.see, see)
}

func TestError(t *testing.T) {
	assert.Nil(t, fromError(nil))
	assert.Equal(t, fromError(errors.New("a")).Error(), "td: a.\n")
	msg := See("hello", "help").String()
	assert.Equal(t, msg, "td: hello. See: 'td help'.\n")
}
<commit_msg>lib: updated the tests.<commit_after>// Copyright (C) 2014 Miquel Sabaté Solà <mikisabate@gmail.com>
// This file is licensed under the MIT license.
// See the LICENSE file.

package lib

import (
	"errors"
	"testing"

	"github.com/stretchr/testify/assert"
)

// Helper used to test custom errors. This is used throughout the test suite.
func testError(t *testing.T, err error, msg, see string) {
	assert.NotNil(t, err)
	e, ce := err.(*Error)
	assert.True(t, ce)
	assert.Equal(t, e.message, msg)
	assert.Equal(t, e.see, see)
}

func TestError(t *testing.T) {
	assert.Nil(t, fromError(nil))
	assert.Equal(t, fromError(errors.New("a")).Error(),
		"\x1b[1;49;31merror:\x1b[0;m a.\n")
	msg := See("hello", "help").String()
	assert.Equal(t, msg,
		"\x1b[1;49;31merror:\x1b[0;m hello. \x1b[1;49;39mSee:\x1b[0;m 'td help'.\n")
}
<|endoftext|>


<commit_before>// author  			sigu-399
// author-github 	https://github.com/sigu-399
// author-mail		sigu.399@gmail.com
//
// repository-name	gojsonschema
// repository-desc 	An implementation of JSON Schema, based on IETF's draft v4 - Go language.
//
// description		Extends JsonSchemaDocument and jsonSchema, implements the validation phase.
//
// created      	28-02-2013

package jsonschema

import (
	"fmt"
	"reflect"
)

type ValidationResult struct {
	valid         bool
	errorMessages []string
}

func (v *ValidationResult) IsValid() bool {
	return v.valid
}

func (v *ValidationResult) addErrorMessage(message string) {
	v.errorMessages = append(v.errorMessages, message)
	v.valid = false
}

func (v *JsonSchemaDocument) Validate(document interface{}) ValidationResult {

	result := ValidationResult{valid: true}
	v.rootSchema.validateRecursive(v.rootSchema, document, &result)
	return result
}

func (v *jsonSchema) Validate(document interface{}) ValidationResult {

	result := ValidationResult{valid: true}
	v.validateRecursive(v, document, &result)
	return result
}

func (v *jsonSchema) validateRecursive(currentSchema *jsonSchema, currentNode interface{}, result *ValidationResult) {

	fmt.Printf("Validation of schema %s\n", currentSchema.property)

	schProperty := currentSchema.property
	schTypes := currentSchema.types

	if currentNode == nil {
		if !schTypes.HasType(TYPE_NULL) {
			result.addErrorMessage(fmt.Sprintf("%s must be of type %s", schProperty, schTypes.String()))
			return
		}
	} else {

		rValue := reflect.ValueOf(currentNode)
		rKind := rValue.Kind()

		var nextNode interface{}
		var ok bool

		fmt.Printf("Type %s\n", rKind.String())

		switch rKind {

		case reflect.Slice:

			if !schTypes.HasType(TYPE_ARRAY) {
				result.addErrorMessage(fmt.Sprintf("%s must be of type %s", schProperty, schTypes.String()))
				return
			}

			castCurrentNode := currentNode.([]interface{})

			v.validateArray(currentSchema, castCurrentNode, result)
			v.validateCommon(currentSchema, castCurrentNode, result)

			for _, nextNode = range castCurrentNode {
				v.validateRecursive(currentSchema.itemsChild, nextNode, result)
			}

		case reflect.Map:

			if !schTypes.HasType(TYPE_OBJECT) {
				result.addErrorMessage(fmt.Sprintf("%s must be of type %s", schProperty, schTypes.String()))
				return
			}

			castCurrentNode := currentNode.(map[string]interface{})

			currentSchema.validateSchema(currentSchema, castCurrentNode, result)

			v.validateObject(currentSchema, castCurrentNode, result)
			v.validateCommon(currentSchema, castCurrentNode, result)

			for _, pSchema := range currentSchema.propertiesChildren {
				nextNode, ok = castCurrentNode[pSchema.property]
				if ok {
					v.validateRecursive(pSchema, nextNode, result)
				}
			}

		case reflect.Bool:

			if !schTypes.HasType(TYPE_BOOLEAN) {
				result.addErrorMessage(fmt.Sprintf("%s must be of type %s", schProperty, schTypes.String()))
				return
			}

			value := currentNode.(bool)
			v.validateCommon(currentSchema, value, result)

		case reflect.String:

			if !schTypes.HasType(TYPE_STRING) {
				result.addErrorMessage(fmt.Sprintf("%s must be of type %s", schProperty, schTypes.String()))
				return
			}

			value := currentNode.(string)

			v.validateString(currentSchema, value, result)
			v.validateCommon(currentSchema, value, result)

		case reflect.Float64:

			value := currentNode.(float64)
			isInteger := isFloat64AnInteger(value)

			formatIsCorrect := schTypes.HasType(TYPE_NUMBER) || (isInteger && schTypes.HasType(TYPE_INTEGER))

			if !formatIsCorrect {
				result.addErrorMessage(fmt.Sprintf("%s must be of type %s", schProperty, schTypes.String()))
				return
			}

			v.validateNumber(currentSchema, value, result)
			v.validateCommon(currentSchema, value, result)
		}
	}
}

func (v *jsonSchema) validateSchema(currentSchema *jsonSchema, currentNode interface{}, result *ValidationResult) {

	if len(currentSchema.anyOf) > 0 {
		validatedAnyOf := false

		for _, anyOfSchema := range currentSchema.anyOf {
			if !validatedAnyOf {
				validationResult := anyOfSchema.Validate(currentNode)
				validatedAnyOf = validationResult.IsValid()
			}
		}
		if !validatedAnyOf {
			result.addErrorMessage(fmt.Sprintf("%s failed to validate any of the schema", currentSchema.property))
		}
	}

	if len(currentSchema.oneOf) > 0 {
		nbValidated := 0

		for _, oneOfSchema := range currentSchema.oneOf {
			validationResult := oneOfSchema.Validate(currentNode)
			if validationResult.IsValid() {
				nbValidated++
			}
		}

		if nbValidated != 1 {
			result.addErrorMessage(fmt.Sprintf("%s failed to validate one of the schema", currentSchema.property))
		}
	}

	if len(currentSchema.allOf) > 0 {
		nbValidated := 0

		for _, allOfSchema := range currentSchema.allOf {
			validationResult := allOfSchema.Validate(currentNode)
			if validationResult.IsValid() {
				nbValidated++
			}
		}

		if nbValidated != len(currentSchema.allOf) {
			result.addErrorMessage(fmt.Sprintf("%s failed to validate all of the schema", currentSchema.property))
		}
	}

	if currentSchema.not != nil {
		validationResult := currentSchema.not.Validate(currentNode)
		if validationResult.IsValid() {
			result.addErrorMessage(fmt.Sprintf("%s is not allowed to validate the schema", currentSchema.property))
		}
	}

}

func (v *jsonSchema) validateCommon(currentSchema *jsonSchema, value interface{}, result *ValidationResult) {

	if len(currentSchema.enum) > 0 {
		has, err := currentSchema.HasEnum(value)
		if err != nil {
			result.addErrorMessage(err.Error())
		}
		if !has {
			result.addErrorMessage(fmt.Sprintf("%s must match one of the enum values", currentSchema.property))
		}
	}
}

func (v *jsonSchema) validateArray(currentSchema *jsonSchema, value []interface{}, result *ValidationResult) {

	if currentSchema.minItems != nil {
		if len(value) < *currentSchema.minItems {
			result.addErrorMessage(fmt.Sprintf("%s must have at least %d items", currentSchema.property, *currentSchema.minItems))
		}
	}

	if currentSchema.maxItems != nil {
		if len(value) > *currentSchema.maxItems {
			result.addErrorMessage(fmt.Sprintf("%s must have at the most %d items", currentSchema.property, *currentSchema.maxItems))
		}
	}

	if currentSchema.uniqueItems {
		var stringifiedItems []string
		for _, v := range value {
			vString, err := marshalToString(v)
			if err != nil {
				result.addErrorMessage(fmt.Sprintf("%s could not be marshalled", currentSchema.property))
			}
			if isStringInSlice(stringifiedItems, *vString) {
				result.addErrorMessage(fmt.Sprintf("%s items must be unique", currentSchema.property))
			}
			stringifiedItems = append(stringifiedItems, *vString)
		}
	}

}

func (v *jsonSchema) validateObject(currentSchema *jsonSchema, value map[string]interface{}, result *ValidationResult) {

	if currentSchema.minProperties != nil {
		if len(value) < *currentSchema.minProperties {
			result.addErrorMessage(fmt.Sprintf("%s must have at least %d properties", currentSchema.property, *currentSchema.minProperties))
		}
	}

	if currentSchema.maxProperties != nil {
		if len(value) > *currentSchema.maxProperties {
			result.addErrorMessage(fmt.Sprintf("%s must have at the most %d properties", currentSchema.property, *currentSchema.maxProperties))
		}
	}
	for _, requiredProperty := range currentSchema.required {
		_, ok := value[requiredProperty]
		if !ok {
			result.addErrorMessage(fmt.Sprintf("%s property is required", requiredProperty))
		}
	}

}

func (v *jsonSchema) validateString(currentSchema *jsonSchema, value string, result *ValidationResult) {

	if currentSchema.minLength != nil {
		if len(value) < *currentSchema.minLength {
			result.addErrorMessage(fmt.Sprintf("%s's length must be greater or equal to %d", currentSchema.property, *currentSchema.minLength))
		}
	}

	if currentSchema.maxLength != nil {
		if len(value) > *currentSchema.maxLength {
			result.addErrorMessage(fmt.Sprintf("%s's length must be lower or equal to %d", currentSchema.property, *currentSchema.maxLength))
		}
	}

	if currentSchema.pattern != nil {
		if !currentSchema.pattern.MatchString(value) {
			result.addErrorMessage(fmt.Sprintf("%s has an invalid format", currentSchema.property))
		}

	}
}

func (v *jsonSchema) validateNumber(currentSchema *jsonSchema, value float64, result *ValidationResult) {

	if currentSchema.multipleOf != nil {
		if !isFloat64AnInteger(value / *currentSchema.multipleOf) {
			result.addErrorMessage(fmt.Sprintf("%f is not a multiple of %f", value, *currentSchema.multipleOf))
		}
	}

	if currentSchema.maximum != nil {
		if currentSchema.exclusiveMaximum {
			if value >= *currentSchema.maximum {
				result.addErrorMessage(fmt.Sprintf("%f must be lower than or equal to %f", value, *currentSchema.maximum))
			}
		} else {
			if value > *currentSchema.maximum {
				result.addErrorMessage(fmt.Sprintf("%f must be lower than %f", value, *currentSchema.maximum))
			}
		}
	}

	if currentSchema.minimum != nil {

		if currentSchema.exclusiveMinimum {
			if value <= *currentSchema.minimum {
				result.addErrorMessage(fmt.Sprintf("%f must be greater than or equal to %f", value, *currentSchema.minimum))
			}
		} else {
			if value < *currentSchema.minimum {
				result.addErrorMessage(fmt.Sprintf("%f must be greater than %f", value, *currentSchema.minimum))
			}
		}
	}

}
<commit_msg>added validation error on properties not handled by the schema<commit_after>// author  			sigu-399
// author-github 	https://github.com/sigu-399
// author-mail		sigu.399@gmail.com
//
// repository-name	gojsonschema
// repository-desc 	An implementation of JSON Schema, based on IETF's draft v4 - Go language.
//
// description		Extends JsonSchemaDocument and jsonSchema, implements the validation phase.
//
// created      	28-02-2013

package jsonschema

import (
	"fmt"
	"reflect"
)

type ValidationResult struct {
	valid         bool
	errorMessages []string
}

func (v *ValidationResult) IsValid() bool {
	return v.valid
}

func (v *ValidationResult) addErrorMessage(message string) {
	v.errorMessages = append(v.errorMessages, message)
	v.valid = false
}

func (v *JsonSchemaDocument) Validate(document interface{}) ValidationResult {

	result := ValidationResult{valid: true}
	v.rootSchema.validateRecursive(v.rootSchema, document, &result)
	return result
}

func (v *jsonSchema) Validate(document interface{}) ValidationResult {

	result := ValidationResult{valid: true}
	v.validateRecursive(v, document, &result)
	return result
}

func (v *jsonSchema) validateRecursive(currentSchema *jsonSchema, currentNode interface{}, result *ValidationResult) {

	fmt.Printf("Validation of schema %s\n", currentSchema.property)

	schProperty := currentSchema.property
	schTypes := currentSchema.types

	if currentNode == nil {
		if !schTypes.HasType(TYPE_NULL) {
			result.addErrorMessage(fmt.Sprintf("%s must be of type %s", schProperty, schTypes.String()))
			return
		}
	} else {

		rValue := reflect.ValueOf(currentNode)
		rKind := rValue.Kind()

		var nextNode interface{}
		var ok bool

		fmt.Printf("Type %s\n", rKind.String())

		switch rKind {

		case reflect.Slice:

			if !schTypes.HasType(TYPE_ARRAY) {
				result.addErrorMessage(fmt.Sprintf("%s must be of type %s", schProperty, schTypes.String()))
				return
			}

			castCurrentNode := currentNode.([]interface{})

			v.validateArray(currentSchema, castCurrentNode, result)
			v.validateCommon(currentSchema, castCurrentNode, result)

			for _, nextNode = range castCurrentNode {
				v.validateRecursive(currentSchema.itemsChild, nextNode, result)
			}

		case reflect.Map:

			if !schTypes.HasType(TYPE_OBJECT) {
				result.addErrorMessage(fmt.Sprintf("%s must be of type %s", schProperty, schTypes.String()))
				return
			}

			castCurrentNode := currentNode.(map[string]interface{})

			currentSchema.validateSchema(currentSchema, castCurrentNode, result)

			v.validateObject(currentSchema, castCurrentNode, result)
			v.validateCommon(currentSchema, castCurrentNode, result)

			for _, pSchema := range currentSchema.propertiesChildren {
				nextNode, ok = castCurrentNode[pSchema.property]
				if ok {
					v.validateRecursive(pSchema, nextNode, result)
				}
			}

		case reflect.Bool:

			if !schTypes.HasType(TYPE_BOOLEAN) {
				result.addErrorMessage(fmt.Sprintf("%s must be of type %s", schProperty, schTypes.String()))
				return
			}

			value := currentNode.(bool)
			v.validateCommon(currentSchema, value, result)

		case reflect.String:

			if !schTypes.HasType(TYPE_STRING) {
				result.addErrorMessage(fmt.Sprintf("%s must be of type %s", schProperty, schTypes.String()))
				return
			}

			value := currentNode.(string)

			v.validateString(currentSchema, value, result)
			v.validateCommon(currentSchema, value, result)

		case reflect.Float64:

			value := currentNode.(float64)
			isInteger := isFloat64AnInteger(value)

			formatIsCorrect := schTypes.HasType(TYPE_NUMBER) || (isInteger && schTypes.HasType(TYPE_INTEGER))

			if !formatIsCorrect {
				result.addErrorMessage(fmt.Sprintf("%s must be of type %s", schProperty, schTypes.String()))
				return
			}

			v.validateNumber(currentSchema, value, result)
			v.validateCommon(currentSchema, value, result)
		}
	}
}

func (v *jsonSchema) validateSchema(currentSchema *jsonSchema, currentNode interface{}, result *ValidationResult) {

	if len(currentSchema.anyOf) > 0 {
		validatedAnyOf := false

		for _, anyOfSchema := range currentSchema.anyOf {
			if !validatedAnyOf {
				validationResult := anyOfSchema.Validate(currentNode)
				validatedAnyOf = validationResult.IsValid()
			}
		}
		if !validatedAnyOf {
			result.addErrorMessage(fmt.Sprintf("%s failed to validate any of the schema", currentSchema.property))
		}
	}

	if len(currentSchema.oneOf) > 0 {
		nbValidated := 0

		for _, oneOfSchema := range currentSchema.oneOf {
			validationResult := oneOfSchema.Validate(currentNode)
			if validationResult.IsValid() {
				nbValidated++
			}
		}

		if nbValidated != 1 {
			result.addErrorMessage(fmt.Sprintf("%s failed to validate one of the schema", currentSchema.property))
		}
	}

	if len(currentSchema.allOf) > 0 {
		nbValidated := 0

		for _, allOfSchema := range currentSchema.allOf {
			validationResult := allOfSchema.Validate(currentNode)
			if validationResult.IsValid() {
				nbValidated++
			}
		}

		if nbValidated != len(currentSchema.allOf) {
			result.addErrorMessage(fmt.Sprintf("%s failed to validate all of the schema", currentSchema.property))
		}
	}

	if currentSchema.not != nil {
		validationResult := currentSchema.not.Validate(currentNode)
		if validationResult.IsValid() {
			result.addErrorMessage(fmt.Sprintf("%s is not allowed to validate the schema", currentSchema.property))
		}
	}

}

func (v *jsonSchema) validateCommon(currentSchema *jsonSchema, value interface{}, result *ValidationResult) {

	if len(currentSchema.enum) > 0 {
		has, err := currentSchema.HasEnum(value)
		if err != nil {
			result.addErrorMessage(err.Error())
		}
		if !has {
			result.addErrorMessage(fmt.Sprintf("%s must match one of the enum values", currentSchema.property))
		}
	}
}

func (v *jsonSchema) validateArray(currentSchema *jsonSchema, value []interface{}, result *ValidationResult) {

	if currentSchema.minItems != nil {
		if len(value) < *currentSchema.minItems {
			result.addErrorMessage(fmt.Sprintf("%s must have at least %d items", currentSchema.property, *currentSchema.minItems))
		}
	}

	if currentSchema.maxItems != nil {
		if len(value) > *currentSchema.maxItems {
			result.addErrorMessage(fmt.Sprintf("%s must have at the most %d items", currentSchema.property, *currentSchema.maxItems))
		}
	}

	if currentSchema.uniqueItems {
		var stringifiedItems []string
		for _, v := range value {
			vString, err := marshalToString(v)
			if err != nil {
				result.addErrorMessage(fmt.Sprintf("%s could not be marshalled", currentSchema.property))
			}
			if isStringInSlice(stringifiedItems, *vString) {
				result.addErrorMessage(fmt.Sprintf("%s items must be unique", currentSchema.property))
			}
			stringifiedItems = append(stringifiedItems, *vString)
		}
	}

}

func (v *jsonSchema) validateObject(currentSchema *jsonSchema, value map[string]interface{}, result *ValidationResult) {

	if currentSchema.minProperties != nil {
		if len(value) < *currentSchema.minProperties {
			result.addErrorMessage(fmt.Sprintf("%s must have at least %d properties", currentSchema.property, *currentSchema.minProperties))
		}
	}

	if currentSchema.maxProperties != nil {
		if len(value) > *currentSchema.maxProperties {
			result.addErrorMessage(fmt.Sprintf("%s must have at the most %d properties", currentSchema.property, *currentSchema.maxProperties))
		}
	}

	for _, requiredProperty := range currentSchema.required {
		_, ok := value[requiredProperty]
		if !ok {
			result.addErrorMessage(fmt.Sprintf("%s property is required", requiredProperty))
		}
	}

	for k, _ := range value {
		if !currentSchema.HasProperty(k) {
			result.addErrorMessage(fmt.Sprintf("%s property is not recognized", k))
		}
	}

}

func (v *jsonSchema) validateString(currentSchema *jsonSchema, value string, result *ValidationResult) {

	if currentSchema.minLength != nil {
		if len(value) < *currentSchema.minLength {
			result.addErrorMessage(fmt.Sprintf("%s's length must be greater or equal to %d", currentSchema.property, *currentSchema.minLength))
		}
	}

	if currentSchema.maxLength != nil {
		if len(value) > *currentSchema.maxLength {
			result.addErrorMessage(fmt.Sprintf("%s's length must be lower or equal to %d", currentSchema.property, *currentSchema.maxLength))
		}
	}

	if currentSchema.pattern != nil {
		if !currentSchema.pattern.MatchString(value) {
			result.addErrorMessage(fmt.Sprintf("%s has an invalid format", currentSchema.property))
		}

	}
}

func (v *jsonSchema) validateNumber(currentSchema *jsonSchema, value float64, result *ValidationResult) {

	if currentSchema.multipleOf != nil {
		if !isFloat64AnInteger(value / *currentSchema.multipleOf) {
			result.addErrorMessage(fmt.Sprintf("%f is not a multiple of %f", value, *currentSchema.multipleOf))
		}
	}

	if currentSchema.maximum != nil {
		if currentSchema.exclusiveMaximum {
			if value >= *currentSchema.maximum {
				result.addErrorMessage(fmt.Sprintf("%f must be lower than or equal to %f", value, *currentSchema.maximum))
			}
		} else {
			if value > *currentSchema.maximum {
				result.addErrorMessage(fmt.Sprintf("%f must be lower than %f", value, *currentSchema.maximum))
			}
		}
	}

	if currentSchema.minimum != nil {

		if currentSchema.exclusiveMinimum {
			if value <= *currentSchema.minimum {
				result.addErrorMessage(fmt.Sprintf("%f must be greater than or equal to %f", value, *currentSchema.minimum))
			}
		} else {
			if value < *currentSchema.minimum {
				result.addErrorMessage(fmt.Sprintf("%f must be greater than %f", value, *currentSchema.minimum))
			}
		}
	}

}
<|endoftext|>


<commit_before>package fuzzaldrin

import (
	"math"
	"os"
	"strings"
)

func basenameScore(str, query string, _score float64) float64 {
	index := len(str) - 1
	for os.IsPathSeparator(str[index]) {
		index--
	}
	slashCount := 0
	lastCharacter := index
	var base string
	for index >= 0 {
		if os.IsPathSeparator(str[index]) {
			slashCount++
			base = str[index+1 : lastCharacter+1]
		} else if index == 0 {
			if lastCharacter < len(str)-1 {
				base = str[0 : lastCharacter+1]
			} else {
				base = str
			}
		}
		index--
	}

	if base == str {
		_score *= 2
	} else if base != "" {
		_score += score(base, query)
	}

	segmentCount := slashCount + 1
	depth := math.Max(1.0, float64(10-segmentCount))
	_score *= depth * 0.01
	return _score
}

func score(str, query string) float64 {
	if str == query {
		return 1
	}

	if queryIsLastPathSegment(str, query) {
		return 1
	}

	totalScore := 0.0
	queryLength := len(query)
	strLength := len(str)

	indexInQuery := 0
	indexInStr := 0

	for indexInQuery < queryLength {
		ch := query[indexInQuery]
		indexInQuery++
		lcIndex := strings.Index(str, strings.ToLower(string(ch)))
		ucIndex := strings.Index(str, strings.ToUpper(string(ch)))
		minIndex := math.Min(float64(lcIndex), float64(ucIndex))
		if minIndex == -1 {
			minIndex = math.Max(float64(lcIndex), float64(ucIndex))
		}
		indexInStr = int(minIndex)
		if indexInStr == -1 {
			return 0
		}

		chScore := 0.1

		if str[indexInStr] == ch {
			chScore += 0.1
		}

		if indexInStr == 0 || os.IsPathSeparator(str[indexInStr-1]) {
			// Start of string bonus
			chScore += 0.8
		} else if strings.Contains("-_ ", string(str[indexInStr-1])) {
			// Start of word bonus
			chScore += 0.7
		}

		// Trim string to after current abbreviation match
		str = str[indexInStr+1 : len(str)]

		totalScore += chScore
	}

	queryScore := totalScore / float64(queryLength)
	return ((queryScore * (float64(queryLength) / float64(strLength))) + queryScore) / 2.0
}

func queryIsLastPathSegment(str, query string) bool {
	if len(str) <= len(query) {
		return false
	}
	if os.IsPathSeparator(str[len(str)-len(query)-1]) {
		return strings.Index(str, query) == (len(str) - len(query))
	}
	return false
}
<commit_msg>Add comment<commit_after>package fuzzaldrin

import (
	"math"
	"os"
	"strings"
)

func basenameScore(str, query string, _score float64) float64 {
	index := len(str) - 1
	for os.IsPathSeparator(str[index]) {
		index--
	}
	slashCount := 0
	lastCharacter := index
	var base string
	for index >= 0 {
		if os.IsPathSeparator(str[index]) {
			slashCount++
			base = str[index+1 : lastCharacter+1]
		} else if index == 0 {
			if lastCharacter < len(str)-1 {
				base = str[0 : lastCharacter+1]
			} else {
				base = str
			}
		}
		index--
	}

	// Basename matches count for more.
	if base == str {
		_score *= 2
	} else if base != "" {
		_score += score(base, query)
	}

	// Shallow files are scored higher
	segmentCount := slashCount + 1
	depth := math.Max(1.0, float64(10-segmentCount))
	_score *= depth * 0.01
	return _score
}

func score(str, query string) float64 {
	if str == query {
		return 1
	}

	if queryIsLastPathSegment(str, query) {
		return 1
	}

	totalScore := 0.0
	queryLength := len(query)
	strLength := len(str)

	indexInQuery := 0
	indexInStr := 0

	for indexInQuery < queryLength {
		ch := query[indexInQuery]
		indexInQuery++
		lcIndex := strings.Index(str, strings.ToLower(string(ch)))
		ucIndex := strings.Index(str, strings.ToUpper(string(ch)))
		minIndex := math.Min(float64(lcIndex), float64(ucIndex))
		if minIndex == -1 {
			minIndex = math.Max(float64(lcIndex), float64(ucIndex))
		}
		indexInStr = int(minIndex)
		if indexInStr == -1 {
			return 0
		}

		chScore := 0.1

		if str[indexInStr] == ch {
			chScore += 0.1
		}

		if indexInStr == 0 || os.IsPathSeparator(str[indexInStr-1]) {
			// Start of string bonus
			chScore += 0.8
		} else if strings.Contains("-_ ", string(str[indexInStr-1])) {
			// Start of word bonus
			chScore += 0.7
		}

		// Trim string to after current abbreviation match
		str = str[indexInStr+1 : len(str)]

		totalScore += chScore
	}

	queryScore := totalScore / float64(queryLength)
	return ((queryScore * (float64(queryLength) / float64(strLength))) + queryScore) / 2.0
}

func queryIsLastPathSegment(str, query string) bool {
	if len(str) <= len(query) {
		return false
	}
	if os.IsPathSeparator(str[len(str)-len(query)-1]) {
		return strings.Index(str, query) == (len(str) - len(query))
	}
	return false
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com

package main

import (
	"fmt"
	elastigo "github.com/mattbaird/elastigo/lib"
	"time"
)

type queryRequest struct {
	startTime time.Time
	endTime   time.Time
}

func queryUsingPlugin(p plugin, req queryRequest) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("queryUsingPlugin() -> %v", e)
		}
	}()

	template := `{
		"size": 10000,
		"query": {
			"bool": {
				"must": [
				%v
				]
			}
		},
		"filter": {
			"range": {
				"utctimestamp": {
					"from": "%v",
					"to": "%v"
				}
			}
		}
	}`

	// Add plugins search terms to the query
	mult := false
	temp := ""
	if len(p.searchTerms) == 0 {
		panic("plugin has no search terms defined")
	}
	for _, x := range p.searchTerms {
		if mult {
			temp += ","
		}
		termtemplate := `{
			"term": {
				"%v": "%v"
			}
		}`
		termbuf := fmt.Sprintf(termtemplate, x.key, x.value)
		temp += termbuf
		mult = true
	}
	for _, x := range p.searchQS {
		if mult {
			temp += ","
		}
		qstemplate := `{
			"query_string": {
				"query": "%v"
			}
		}`
		qsbuf := fmt.Sprintf(qstemplate, x)
		temp += qsbuf
		mult = true
	}
	querybuf := fmt.Sprintf(template, temp, req.startTime.Format(time.RFC3339), req.endTime.Format(time.RFC3339))
	conn := elastigo.NewConn()
	defer conn.Close()
	conn.Domain = cfg.ES.EventESHost
	res, err := conn.Search(cfg.ES.EventIndex, "", nil, querybuf)
	if err != nil {
		panic(err)
	}
	logf("plugin %v returned %v hits", p.name, res.Hits.Len())

	if res.Hits.Len() == 0 {
		return nil
	}

	pluginInput, err := pluginRequestDataFromES(res)
	if err != nil {
		panic(err)
	}
	err = p.runPlugin(pluginInput)
	if err != nil {
		panic(err)
	}

	return nil
}

func handleQueryRequest(q queryRequest) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("handleQueryRequest() -> %v", e)
		}
	}()
	logf("handling new query request")

	// Execute a query for each registered plugin
	for _, x := range pluginList {
		err := queryUsingPlugin(x, q)
		if err != nil {
			panic(err)
		}
	}
	return nil
}

func queryHandler(exitCh chan bool, notifyCh chan bool) {
	defer func() {
		if e := recover(); e != nil {
			logf("queryHandler() -> %v", e)
		}
		logf("query handler exiting")
		notifyCh <- true
	}()
	logf("query handler started")

	for {
		select {
		case qr := <-queryRequestCh:
			err := handleQueryRequest(qr)
			if err != nil {
				panic(err)
			}
		case <-exitCh:
			return
		}
	}
}
<commit_msg>Update query format to support elasticsearch 5<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com

package main

import (
	"fmt"
	elastigo "github.com/mattbaird/elastigo/lib"
	"time"
)

type queryRequest struct {
	startTime time.Time
	endTime   time.Time
}

func queryUsingPlugin(p plugin, req queryRequest) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("queryUsingPlugin() -> %v", e)
		}
	}()

	template := `{
		"size": 10000,
		"query": {
			"bool": {
				"must": [
				%v
				],
				"filter": {
					"range": {
						"utctimestamp": {
							"gte": "%v",
							"lte": "%v"
						}
					}
				}
			}
		}
	}`

	// Add plugins search terms to the query
	mult := false
	temp := ""
	if len(p.searchTerms) == 0 {
		panic("plugin has no search terms defined")
	}
	for _, x := range p.searchTerms {
		if mult {
			temp += ","
		}
		termtemplate := `{
			"term": {
				"%v": "%v"
			}
		}`
		termbuf := fmt.Sprintf(termtemplate, x.key, x.value)
		temp += termbuf
		mult = true
	}
	for _, x := range p.searchQS {
		if mult {
			temp += ","
		}
		qstemplate := `{
			"query_string": {
				"query": "%v"
			}
		}`
		qsbuf := fmt.Sprintf(qstemplate, x)
		temp += qsbuf
		mult = true
	}
	querybuf := fmt.Sprintf(template, temp, req.startTime.Format(time.RFC3339), req.endTime.Format(time.RFC3339))
	conn := elastigo.NewConn()
	defer conn.Close()
	conn.Domain = cfg.ES.EventESHost
	res, err := conn.Search(cfg.ES.EventIndex, "", nil, querybuf)
	if err != nil {
		panic(err)
	}
	logf("plugin %v returned %v hits", p.name, res.Hits.Len())

	if res.Hits.Len() == 0 {
		return nil
	}

	pluginInput, err := pluginRequestDataFromES(res)
	if err != nil {
		panic(err)
	}
	err = p.runPlugin(pluginInput)
	if err != nil {
		panic(err)
	}

	return nil
}

func handleQueryRequest(q queryRequest) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("handleQueryRequest() -> %v", e)
		}
	}()
	logf("handling new query request")

	// Execute a query for each registered plugin
	for _, x := range pluginList {
		err := queryUsingPlugin(x, q)
		if err != nil {
			panic(err)
		}
	}
	return nil
}

func queryHandler(exitCh chan bool, notifyCh chan bool) {
	defer func() {
		if e := recover(); e != nil {
			logf("queryHandler() -> %v", e)
		}
		logf("query handler exiting")
		notifyCh <- true
	}()
	logf("query handler started")

	for {
		select {
		case qr := <-queryRequestCh:
			err := handleQueryRequest(qr)
			if err != nil {
				panic(err)
			}
		case <-exitCh:
			return
		}
	}
}
<|endoftext|>


<commit_before>package pagecache

import (
	"bytes"
	"net/http"
	"strings"
	"sync"

	"time"

	"github.com/herb-go/herb/cache"
)

//Debug if pagecache module is in debug mode.
var Debug bool

var pagePool = sync.Pool{
	New: func() interface{} {
		return []byte{}
	},
}

//PageCacheKeyHeader http respnose header wihch will echo page cache key when debug is on.
var PageCacheKeyHeader = "herbgo-debug-pagecache"

//New create page cache with given cache.
func New(c cache.Cacheable) *PageCache {
	p := PageCache{
		Cache:           c,
		KeyPrefix:       defaultKeyPrefix,
		StatusValidator: defualtStatusValidator,
	}
	return &p
}

type cachedPage struct {
	Status   int
	Header   map[string][]string
	Response []byte
}

//PageCache page cache struct
type PageCache struct {
	Cache           cache.Cacheable
	KeyPrefix       string
	StatusValidator func(status int) bool
}

var defaultKeyPrefix = ""

func defualtStatusValidator(status int) bool {
	return status < 500
}

//ValidateStatus validate http status code if page should be cached.
//Return validated result.
//Page cache StatusValidator will be used.
func (p *PageCache) ValidateStatus(status int) bool {
	if p.StatusValidator != nil {
		return p.StatusValidator(status)
	}
	return defualtStatusValidator(status)
}
func (p *PageCache) serve(key string, ttl time.Duration, w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	page := cachedPage{
		Response: pagePool.Get().([]byte),
	}
	defer pagePool.Put(page.Response)
	err := p.Cache.Load(key, &page, ttl, func(key string) (interface{}, error) {
		cw := cacheResponseWriter{
			writer: *(bytes.NewBuffer([]byte{})),
			header: http.Header{},
			status: 0,
		}
		next(&cw, r)
		page.Header = map[string][]string(cw.Header())
		page.Response = cw.writer.Bytes()
		page.Status = cw.status
		if p.ValidateStatus(cw.status) {
			return page, nil
		}
		return nil, cache.ErrNotCacheable
	})
	if err != nil {
		if err != cache.ErrEntryTooLarge && err != cache.ErrNotCacheable {
			panic(err)
		}
	}
	h := w.Header()
	for k, v := range page.Header {
		h[k] = v
	}
	if Debug {
		k, err := p.Cache.FinalKey(key)
		if err == nil {
			h.Set(PageCacheKeyHeader, strings.Replace(k, cache.KeyPrefix, " ", -1))
		}
	}
	if page.Status != 0 {
		w.WriteHeader(page.Status)
	}
	_, err = w.Write(page.Response)
	if err != nil {
		panic(err)
	}
}

//Middleware create page cache middleware with give key generator and ttl.
//Return middleware created.
func (p *PageCache) Middleware(keyGenerator func(r *http.Request) string, ttl time.Duration) func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	return func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		key := keyGenerator(r)
		if key == "" {
			next(w, r)
			return
		}
		p.serve(p.KeyPrefix+key, ttl, w, r, next)
	}
}

//FieldMiddleware create page cache middleare with given cache field,ttl,and status validator.
//Return middleware created.
func FieldMiddleware(FieldGenerator func(r *http.Request) *cache.Field, ttl time.Duration, statusValidator func(status int) bool) func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	return func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		f := FieldGenerator(r)
		if f == nil {
			next(w, r)
			return
		}
		p := &PageCache{
			Cache: f.Cache,
		}
		p.serve(f.FieldName, ttl, w, r, next)
	}
}

type cacheResponseWriter struct {
	writer bytes.Buffer
	header http.Header
	status int
}

func (w *cacheResponseWriter) Write(data []byte) (int, error) {
	return w.writer.Write(data)
}
func (w *cacheResponseWriter) Header() http.Header {
	return w.header
}
func (w *cacheResponseWriter) WriteHeader(status int) {
	w.status = status
}
<commit_msg>update<commit_after>package pagecache

import (
	"bytes"
	"net/http"
	"strings"

	"time"

	"github.com/herb-go/herb/cache"
)

//Debug if pagecache module is in debug mode.
var Debug bool

// var pagePool = sync.Pool{
// 	New: func() interface{} {
// 		return []byte{}
// 	},
// }

//PageCacheKeyHeader http respnose header wihch will echo page cache key when debug is on.
var PageCacheKeyHeader = "herbgo-debug-pagecache"

//New create page cache with given cache.
func New(c cache.Cacheable) *PageCache {
	p := PageCache{
		Cache:           c,
		KeyPrefix:       defaultKeyPrefix,
		StatusValidator: defualtStatusValidator,
	}
	return &p
}

type cachedPage struct {
	Status   int
	Header   map[string][]string
	Response []byte
}

//PageCache page cache struct
type PageCache struct {
	Cache           cache.Cacheable
	KeyPrefix       string
	StatusValidator func(status int) bool
}

var defaultKeyPrefix = ""

func defualtStatusValidator(status int) bool {
	return status < 500
}

//ValidateStatus validate http status code if page should be cached.
//Return validated result.
//Page cache StatusValidator will be used.
func (p *PageCache) ValidateStatus(status int) bool {
	if p.StatusValidator != nil {
		return p.StatusValidator(status)
	}
	return defualtStatusValidator(status)
}
func (p *PageCache) serve(key string, ttl time.Duration, w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	// page := cachedPage{
	// 	Response: pagePool.Get().([]byte),
	// }
	// defer pagePool.Put(page.Response)
	page := cachedPage{}
	err := p.Cache.Load(key, &page, ttl, func(key string) (interface{}, error) {
		cw := cacheResponseWriter{
			writer: *(bytes.NewBuffer([]byte{})),
			header: http.Header{},
			status: 0,
		}
		next(&cw, r)
		page.Header = map[string][]string(cw.Header())
		page.Response = cw.writer.Bytes()
		page.Status = cw.status
		if p.ValidateStatus(cw.status) {
			return page, nil
		}
		return nil, cache.ErrNotCacheable
	})
	if err != nil {
		if err != cache.ErrEntryTooLarge && err != cache.ErrNotCacheable {
			panic(err)
		}
	}
	h := w.Header()
	for k, v := range page.Header {
		h[k] = v
	}
	if Debug {
		k, err := p.Cache.FinalKey(key)
		if err == nil {
			h.Set(PageCacheKeyHeader, strings.Replace(k, cache.KeyPrefix, " ", -1))
		}
	}
	if page.Status != 0 {
		w.WriteHeader(page.Status)
	}
	_, err = w.Write(page.Response)
	if err != nil {
		panic(err)
	}
}

//Middleware create page cache middleware with give key generator and ttl.
//Return middleware created.
func (p *PageCache) Middleware(keyGenerator func(r *http.Request) string, ttl time.Duration) func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	return func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		key := keyGenerator(r)
		if key == "" {
			next(w, r)
			return
		}
		p.serve(p.KeyPrefix+key, ttl, w, r, next)
	}
}

//FieldMiddleware create page cache middleare with given cache field,ttl,and status validator.
//Return middleware created.
func FieldMiddleware(FieldGenerator func(r *http.Request) *cache.Field, ttl time.Duration, statusValidator func(status int) bool) func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	return func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		f := FieldGenerator(r)
		if f == nil {
			next(w, r)
			return
		}
		p := &PageCache{
			Cache: f.Cache,
		}
		p.serve(f.FieldName, ttl, w, r, next)
	}
}

type cacheResponseWriter struct {
	writer bytes.Buffer
	header http.Header
	status int
}

func (w *cacheResponseWriter) Write(data []byte) (int, error) {
	return w.writer.Write(data)
}
func (w *cacheResponseWriter) Header() http.Header {
	return w.header
}
func (w *cacheResponseWriter) WriteHeader(status int) {
	w.status = status
}
<|endoftext|>