<commit_before>package jsonselect

import (
    "log"
    "errors"
    "regexp"
    "strconv"
    "strings"
    "github.com/latestrevision/go-simplejson"
)

type Parser struct {
    data *simplejson.Json
    nodes []*Node
}

func CreateParser(body string) (*Parser, error) {
    json, err := simplejson.NewJson([]byte(body))
    if err != nil {
        return nil, err
    }
    parser := Parser{json, nil}
    parser.mapDocument()
    return &parser, err
}

func (p *Parser) Parse(selector string) ([]*Node, error) {
    tokens, err := Lex(selector, SCANNER)
    if err != nil {
        return nil, err
    }

    results, err := p.selectorProduction(tokens)
    if err != nil {
        return nil, err
    }

    return results, nil
}

func (p *Parser) selectorProduction(tokens []*token) ([]*Node, error) {
    var results []*Node
    var validators []func(*Node)bool
    var matched bool
    var value interface{}
    var validator func(*Node)bool

    _, matched, _ = p.peek(tokens, S_TYPE)
    if matched {
        value, tokens, _ = p.match(tokens, S_TYPE)
        validators = append(
            validators,
            p.typeProduction(value),
        )
    }
    _, matched, _ = p.peek(tokens, S_IDENTIFIER)
    if matched {
        value, tokens, _ = p.match(tokens, S_IDENTIFIER)
        log.Print("Found ", S_IDENTIFIER, ": ", value)
        validators = append(
            validators,
            p.keyProduction(value),
        )
    }
    _, matched, _ = p.peek(tokens, S_PCLASS)
    if matched {
        value, tokens, _ = p.match(tokens, S_PCLASS)
        log.Print("Found ", S_PCLASS, ": ", value)
        validators = append(
            validators,
            p.pclassProduction(value),
        )
    }
    _, matched, _ = p.peek(tokens, S_NTH_FUNC)
    if matched {
        value, tokens, _ = p.match(tokens, S_NTH_FUNC)
        log.Print("Found ", S_NTH_FUNC, ": ", value)
        validator, tokens = p.nthChildProduction(value, tokens)
        validators = append(validators, validator)
    }
    _, matched, _ = p.peek(tokens, S_PCLASS_FUNC)
    if matched {
        value, tokens, _ = p.match(tokens, S_PCLASS_FUNC)
        log.Print("Found ", S_PCLASS_FUNC, ": ", value)
        validator, tokens = p.pclassFuncProduction(value, tokens)
        validators = append(validators, validator)
    }

    log.Print(len(validators), " validators found.")

    if len(validators) < 1 {
        return nil, errors.New("No selector recognized")
    }

    results, err := p.matchNodes(validators)
    if err != nil {
        return nil, err
    }

    _, matched, _ = p.peek(tokens, S_OPER)
    if matched {
        value, tokens, _ = p.match(tokens, S_OPER)
        rvals, err := p.selectorProduction(tokens)
        if err != nil {
            return nil, err
        }
        switch value {
            case ",":
                for _, val := range rvals {
                    // TODO: This is quite slow
                    // it seems like it's probably quite easy to expand
                    // the list just once
                    results = append(results, val)
                }
            case ">":
                results = parents(results, rvals)
            case "~":
                results = siblings(results, rvals)
            case " ":
                results = ancestors(results, rvals)
            default:
                return nil, errors.New("Unrecognized operator")
        }
    } else if len(tokens) > 0 {
        rvals, err := p.selectorProduction(tokens)
        if err != nil {
            return nil, err
        }
        results = ancestors(results, rvals)
    }

    return results, nil
}

func (p *Parser) peek(tokens []*token, typ tokenType) (interface{}, bool, error) {
    if len(tokens) < 1 {
        return nil, false, errors.New("No more tokens")
    }
    if tokens[0].typ == typ {
        return tokens[0].val, true, nil
    }
    return nil, false, nil
}

func (p *Parser) match(tokens []*token, typ tokenType) (interface{}, []*token, error) {
    value, matched, _ := p.peek(tokens, typ)
    if !matched {
        return nil, tokens, errors.New("Match not successful")
    }
    _, tokens = tokens[0], tokens[1:]
    return value, tokens, nil
}

func (p *Parser) matchNodes(validators []func(*Node)bool) ([]*Node, error) {
    var matches []*Node
    for _, node := range p.mapDocument() {
        var passed = true
        for _, validator := range validators {
            if !validator(node) {
                passed = false
                break
            }
        }
        if passed {
            matches = append(matches, node)
        }
    }
    return matches, nil
}

func (p *Parser) typeProduction(value interface{}) func(*Node)bool {
    return func(node *Node) bool {
        return string(node.typ) == value
    }
}
func (p *Parser) keyProduction(value interface{}) func(*Node)bool {
    // TODO: Verify this -- I'm not sure what this is supposed to match
    return func(node *Node) bool {
        if node.parent_key == ""{
            return false
        }
        return string(node.parent_key) == value
    }
}

func (p *Parser) pclassProduction(value interface{}) func(*Node)bool {
    pclass := value.(string)
    if pclass == "first-child" {
        return func(node *Node) bool {
            return node.idx == 1
        }
    } else if pclass == "last-child" {
        return func(node *Node) bool {
            return node.siblings > 0 && node.idx == node.siblings
        }
    } else if pclass == "only-child" {
        return func(node *Node) bool {
            return node.siblings == 1
        }
    } else if pclass == "root" {
        return func(node *Node) bool {
            return node.parent == nil
        }
    } else if pclass == "empty" {
        return func(node *Node) bool {
            return node.typ == J_ARRAY && len(node.value.(string)) < 1
        }
    }
    return func(node *Node) bool {
        return false
    }
}

func (p *Parser) nthChildProduction(value interface{}, tokens []*token) (func(*Node)bool, []*token) {
    nthChildRegexp := regexp.MustCompile(`^\s*\(\s*(?:([+\-]?)([0-9]*)n\s*(?:([+\-])\s*([0-9]))?|(odd|even)|([+\-]?[0-9]+))\s*\)`)
    args, tokens, _ := p.match(tokens, S_EXPR)
    var a int
    var b int
    var reverse bool = false

    pattern := nthChildRegexp.FindStringSubmatch(args.(string))

    if pattern[5] != "" {
        a = 2
        if pattern[5] == "odd" {
            b = 1
        } else {
            b = 0
        }
    } else if pattern[6] != ""{
        a = 0
        b_temp, _ := strconv.ParseInt(pattern[6], 10, 64)
        b = int(b_temp)
    } else {
        sign := "+"
        if pattern[1] != "" {
            sign = pattern[1]
        }
        coeff := "1"
        if pattern[2] != "" {
            coeff = pattern[2]
        }
        a, _ := strconv.ParseInt(coeff, 10, 64)
        if sign == "-" {
            a = -1 * a
        }
        g3, _ := strconv.ParseInt(pattern[3], 10, 64)
        g4, _ := strconv.ParseInt(pattern[4], 10, 64)
        if pattern[3] != "" {
            b = int(g3 + g4)
        } else {
            b = 0
        }
    }

    if value.(string) == "nth-last-child" {
        reverse = true
    }

    return func(node *Node)bool {
        if node.siblings == 0 {
            return false
        }

        idx := node.idx - 1
        if reverse {
            idx = node.siblings - idx
        } else {
            idx++
        }

        if a == 0 {
            return b == idx
        } else {
            return ((idx - b) % a) == 0 && (idx * a + b) >= 0
        }
    }, tokens
}

func (p *Parser) pclassFuncProduction(value interface{}, tokens []*token) (func(*Node)bool, []*token) {
    sargs, tokens, _ := p.match(tokens, S_EXPR)
    pclass := value.(string)

    if pclass == "expr" {
        tokens, _ := Lex(sargs.(string), EXPRESSION_SCANNER)
        return func(node *Node)bool {
            return p.parseExpression(tokens, node).(bool)
        }, tokens
    }

    args, _ := Lex(sargs.(string)[1:len(sargs.(string))-1], SCANNER)

    if pclass == "has" {
        for _, token := range args {
            if token.val == ">" {
                token.typ = S_EMPTY
            }
        }
        rvals, _ := p.selectorProduction(args)

        var ancestors []*Node
        for _, node := range rvals {
            ancestors = append(ancestors, node.parent)
        }
        return func(node *Node)bool {
            return nodeIsMemberOfList(node, ancestors)
        }, tokens
    } else if pclass == "contains" {
        return func(node *Node)bool {
            return node.typ == J_STRING && strings.Count(node.value.(string), args[0].val.(string)) > 0
        }, tokens
    } else if pclass == "val" {
        return func(node *Node)bool {
            return node.typ == J_STRING && node.value.(string) == args[0].val.(string)
        }, tokens
    }

    // If we didn't find a known pclass, do not match anything.
    return func(node *Node)bool {
        return false
    }, tokens
}

func (p *Parser) evaluateParsedExpression(tokens []*token, node *Node, cmap map[string]func(interface{}, interface{})int) interface{} {
    var matched bool
    var lhs interface{}
    var rhs interface{}

    if len(tokens) < 1 {
        return false
    }

    value, matched, _ := p.peek(tokens, S_PAREN)
    if matched && value.(string) == "(" {
        _, tokens, _ = p.match(tokens, S_PAREN)
        lhs = p.evaluateParsedExpression(tokens, node, cmap)
        return lhs
    }

    _, matched, _ = p.peek(tokens, S_PVAR)
    if matched {
        _, tokens, _ = p.match(tokens, S_PVAR)
        lhs = node.value
    } else {
        relevantTokens := []tokenType{S_STRING, S_BOOL, S_NIL, S_NUMBER}
        for _, ttype := range relevantTokens {
            _, matched, _ = p.peek(tokens, ttype)
            if matched {
                lhs, tokens, _ = p.match(tokens, ttype)
                break
            }
        }
    }

    value, matched, _ = p.peek(tokens, S_PAREN)
    if matched && value.(string) == ")" {
        _, tokens, _ = p.match(tokens, S_PAREN)
        // Short-circuit?
        return lhs
    }

    binop, tokens, _ := p.match(tokens, S_BINOP)
    comparatorFunction := cmap[binop.(string)]
    rhs = p.evaluateParsedExpression(tokens, node, cmap)

    return comparatorFunction(lhs, rhs)
}

func (p *Parser) parseExpression(tokens []*token, node *Node) interface{} {
    comparatorMap := map[string]func(lhs interface{}, rhs interface{})int{
        "*": func(lhs interface{}, rhs interface{})int {
            return lhs.(int) * rhs.(int)
        },
        "/": func(lhs interface{}, rhs interface{})int {
            return lhs.(int) / rhs.(int)
        },
        "%": func(lhs interface{}, rhs interface{})int {
            return lhs.(int) % rhs.(int)
        },
        "+": func(lhs interface{}, rhs interface{})int {
            return lhs.(int) + rhs.(int)
        },
        "-": func(lhs interface{}, rhs interface{})int {
            return lhs.(int) - rhs.(int)
        },
        "<=": func(lhs interface{}, rhs interface{})int {
            if lhs.(int) <= rhs.(int) {
                return 1
            }
            return 0
        },
        "<": func(lhs interface{}, rhs interface{})int {
            if lhs.(int) < rhs.(int) {
                return 1
            }
            return 0
        },
        ">=": func(lhs interface{}, rhs interface{})int {
            if lhs.(int) > rhs.(int) {
                return 1
            }
            return 0
        },
        ">": func(lhs interface{}, rhs interface{})int {
            if lhs.(int) > rhs.(int) {
                return 1
            }
            return 0
        },
        "$=": func(lhs interface{}, rhs interface{})int {
            lhs_str := lhs.(string)
            rhs_str := rhs.(string)
            if strings.LastIndex(lhs_str, rhs_str) == len(lhs_str) - len(rhs_str) {
                return 1
            }
            return 0
        },
        "^=": func(lhs interface{}, rhs interface{})int {
            lhs_str := lhs.(string)
            rhs_str := rhs.(string)
            if strings.Index(lhs_str, rhs_str) == 0 {
                return 1
            }
            return 0
        },
        "*=": func(lhs interface{}, rhs interface{})int {
            lhs_str := lhs.(string)
            rhs_str := rhs.(string)
            if strings.Index(lhs_str, rhs_str) != 0 {
                return 1
            }
            return 0
        },
        "=": func(lhs interface{}, rhs interface{})int {
            if lhs.(string) == rhs.(string) {
                return 1
            }
            return 0
        },
        "!=": func(lhs interface{}, rhs interface{})int {
            if lhs.(string) != rhs.(string) {
                return 1
            }
            return 0
        },
        "&&": func(lhs interface{}, rhs interface{})int {
            if lhs.(bool) && rhs.(bool) {
                return 1
            }
            return 0
        },
        "||": func(lhs interface{}, rhs interface{})int {
            if lhs.(bool) || rhs.(bool) {
                return 1
            }
            return 0
        },
    }
    return p.evaluateParsedExpression(tokens, node, comparatorMap)
}
<commit_msg>Fixing document mapping.<commit_after>package jsonselect

import (
    "errors"
    "regexp"
    "strconv"
    "strings"
    "github.com/latestrevision/go-simplejson"
)

type Parser struct {
    data *simplejson.Json
    nodes []*Node
}

func CreateParser(body string) (*Parser, error) {
    json, err := simplejson.NewJson([]byte(body))
    if err != nil {
        return nil, err
    }
    parser := Parser{json, nil}
    parser.mapDocument()
    return &parser, err
}

func (p *Parser) Parse(selector string) ([]*Node, error) {
    tokens, err := Lex(selector, SCANNER)
    if err != nil {
        return nil, err
    }

    documentMap := p.mapDocument()
    results, err := p.selectorProduction(tokens, documentMap)
    if err != nil {
        return nil, err
    }

    return results, nil
}

func (p *Parser) selectorProduction(tokens []*token, documentMap []*Node) ([]*Node, error) {
    var results []*Node
    var validators []func(*Node)bool
    var matched bool
    var value interface{}
    var validator func(*Node)bool

    _, matched, _ = p.peek(tokens, S_TYPE)
    if matched {
        value, tokens, _ = p.match(tokens, S_TYPE)
        validators = append(
            validators,
            p.typeProduction(value),
        )
    }
    _, matched, _ = p.peek(tokens, S_IDENTIFIER)
    if matched {
        value, tokens, _ = p.match(tokens, S_IDENTIFIER)
        validators = append(
            validators,
            p.keyProduction(value),
        )
    }
    _, matched, _ = p.peek(tokens, S_PCLASS)
    if matched {
        value, tokens, _ = p.match(tokens, S_PCLASS)
        validators = append(
            validators,
            p.pclassProduction(value),
        )
    }
    _, matched, _ = p.peek(tokens, S_NTH_FUNC)
    if matched {
        value, tokens, _ = p.match(tokens, S_NTH_FUNC)
        validator, tokens = p.nthChildProduction(value, tokens)
        validators = append(validators, validator)
    }
    _, matched, _ = p.peek(tokens, S_PCLASS_FUNC)
    if matched {
        value, tokens, _ = p.match(tokens, S_PCLASS_FUNC)
        validator, tokens = p.pclassFuncProduction(value, tokens, documentMap)
        validators = append(validators, validator)
    }

    if len(validators) < 1 {
        return nil, errors.New("No selector recognized")
    }

    results, err := p.matchNodes(validators, documentMap)
    if err != nil {
        return nil, err
    }

    _, matched, _ = p.peek(tokens, S_OPER)
    if matched {
        value, tokens, _ = p.match(tokens, S_OPER)
        rvals, err := p.selectorProduction(tokens, documentMap)
        if err != nil {
            return nil, err
        }
        switch value {
            case ",":
                for _, val := range rvals {
                    // TODO: This is quite slow
                    // it seems like it's probably quite easy to expand
                    // the list just once
                    results = append(results, val)
                }
            case ">":
                results = parents(results, rvals)
            case "~":
                results = siblings(results, rvals)
            case " ":
                results = ancestors(results, rvals)
            default:
                return nil, errors.New("Unrecognized operator")
        }
    } else if len(tokens) > 0 {
        rvals, err := p.selectorProduction(tokens, documentMap)
        if err != nil {
            return nil, err
        }
        results = ancestors(results, rvals)
    }

    return results, nil
}

func (p *Parser) peek(tokens []*token, typ tokenType) (interface{}, bool, error) {
    if len(tokens) < 1 {
        return nil, false, errors.New("No more tokens")
    }
    if tokens[0].typ == typ {
        return tokens[0].val, true, nil
    }
    return nil, false, nil
}

func (p *Parser) match(tokens []*token, typ tokenType) (interface{}, []*token, error) {
    value, matched, _ := p.peek(tokens, typ)
    if !matched {
        return nil, tokens, errors.New("Match not successful")
    }
    _, tokens = tokens[0], tokens[1:]
    return value, tokens, nil
}

func (p *Parser) matchNodes(validators []func(*Node)bool, documentMap []*Node) ([]*Node, error) {
    var matches []*Node
    for _, node := range documentMap {
        var passed = true
        for _, validator := range validators {
            if !validator(node) {
                passed = false
                break
            }
        }
        if passed {
            matches = append(matches, node)
        }
    }
    return matches, nil
}

func (p *Parser) typeProduction(value interface{}) func(*Node)bool {
    return func(node *Node) bool {
        return string(node.typ) == value
    }
}
func (p *Parser) keyProduction(value interface{}) func(*Node)bool {
    // TODO: Verify this -- I'm not sure what this is supposed to match
    return func(node *Node) bool {
        if node.parent_key == ""{
            return false
        }
        return string(node.parent_key) == value
    }
}

func (p *Parser) pclassProduction(value interface{}) func(*Node)bool {
    pclass := value.(string)
    if pclass == "first-child" {
        return func(node *Node) bool {
            return node.idx == 1
        }
    } else if pclass == "last-child" {
        return func(node *Node) bool {
            return node.siblings > 0 && node.idx == node.siblings
        }
    } else if pclass == "only-child" {
        return func(node *Node) bool {
            return node.siblings == 1
        }
    } else if pclass == "root" {
        return func(node *Node) bool {
            return node.parent == nil
        }
    } else if pclass == "empty" {
        return func(node *Node) bool {
            return node.typ == J_ARRAY && len(node.value.(string)) < 1
        }
    }
    return func(node *Node) bool {
        return false
    }
}

func (p *Parser) nthChildProduction(value interface{}, tokens []*token) (func(*Node)bool, []*token) {
    nthChildRegexp := regexp.MustCompile(`^\s*\(\s*(?:([+\-]?)([0-9]*)n\s*(?:([+\-])\s*([0-9]))?|(odd|even)|([+\-]?[0-9]+))\s*\)`)
    args, tokens, _ := p.match(tokens, S_EXPR)
    var a int
    var b int
    var reverse bool = false

    pattern := nthChildRegexp.FindStringSubmatch(args.(string))

    if pattern[5] != "" {
        a = 2
        if pattern[5] == "odd" {
            b = 1
        } else {
            b = 0
        }
    } else if pattern[6] != ""{
        a = 0
        b_temp, _ := strconv.ParseInt(pattern[6], 10, 64)
        b = int(b_temp)
    } else {
        sign := "+"
        if pattern[1] != "" {
            sign = pattern[1]
        }
        coeff := "1"
        if pattern[2] != "" {
            coeff = pattern[2]
        }
        a, _ := strconv.ParseInt(coeff, 10, 64)
        if sign == "-" {
            a = -1 * a
        }
        g3, _ := strconv.ParseInt(pattern[3], 10, 64)
        g4, _ := strconv.ParseInt(pattern[4], 10, 64)
        if pattern[3] != "" {
            b = int(g3 + g4)
        } else {
            b = 0
        }
    }

    if value.(string) == "nth-last-child" {
        reverse = true
    }

    return func(node *Node)bool {
        if node.siblings == 0 {
            return false
        }

        idx := node.idx - 1
        if reverse {
            idx = node.siblings - idx
        } else {
            idx++
        }

        if a == 0 {
            return b == idx
        } else {
            return ((idx - b) % a) == 0 && (idx * a + b) >= 0
        }
    }, tokens
}

func (p *Parser) pclassFuncProduction(value interface{}, tokens []*token, documentMap []*Node) (func(*Node)bool, []*token) {
    sargs, tokens, _ := p.match(tokens, S_EXPR)
    pclass := value.(string)

    if pclass == "expr" {
        tokens, _ := Lex(sargs.(string), EXPRESSION_SCANNER)
        return func(node *Node)bool {
            return p.parseExpression(tokens, node).(bool)
        }, tokens
    }

    args, _ := Lex(sargs.(string)[1:len(sargs.(string))-1], SCANNER)

    if pclass == "has" {
        for _, token := range args {
            if token.val == ">" {
                token.typ = S_EMPTY
            }
        }
        rvals, _ := p.selectorProduction(args, documentMap)

        var ancestors []*Node
        for _, node := range rvals {
            ancestors = append(ancestors, node.parent)
        }
        return func(node *Node)bool {
            return nodeIsMemberOfList(node, ancestors)
        }, tokens
    } else if pclass == "contains" {
        return func(node *Node)bool {
            return node.typ == J_STRING && strings.Count(node.value.(string), args[0].val.(string)) > 0
        }, tokens
    } else if pclass == "val" {
        return func(node *Node)bool {
            return node.typ == J_STRING && node.value.(string) == args[0].val.(string)
        }, tokens
    }

    // If we didn't find a known pclass, do not match anything.
    return func(node *Node)bool {
        return false
    }, tokens
}

func (p *Parser) evaluateParsedExpression(tokens []*token, node *Node, cmap map[string]func(interface{}, interface{})int) interface{} {
    var matched bool
    var lhs interface{}
    var rhs interface{}

    if len(tokens) < 1 {
        return false
    }

    value, matched, _ := p.peek(tokens, S_PAREN)
    if matched && value.(string) == "(" {
        _, tokens, _ = p.match(tokens, S_PAREN)
        lhs = p.evaluateParsedExpression(tokens, node, cmap)
        return lhs
    }

    _, matched, _ = p.peek(tokens, S_PVAR)
    if matched {
        _, tokens, _ = p.match(tokens, S_PVAR)
        lhs = node.value
    } else {
        relevantTokens := []tokenType{S_STRING, S_BOOL, S_NIL, S_NUMBER}
        for _, ttype := range relevantTokens {
            _, matched, _ = p.peek(tokens, ttype)
            if matched {
                lhs, tokens, _ = p.match(tokens, ttype)
                break
            }
        }
    }

    value, matched, _ = p.peek(tokens, S_PAREN)
    if matched && value.(string) == ")" {
        _, tokens, _ = p.match(tokens, S_PAREN)
        // Short-circuit?
        return lhs
    }

    binop, tokens, _ := p.match(tokens, S_BINOP)
    comparatorFunction := cmap[binop.(string)]
    rhs = p.evaluateParsedExpression(tokens, node, cmap)

    return comparatorFunction(lhs, rhs)
}

func (p *Parser) parseExpression(tokens []*token, node *Node) interface{} {
    comparatorMap := map[string]func(lhs interface{}, rhs interface{})int{
        "*": func(lhs interface{}, rhs interface{})int {
            return lhs.(int) * rhs.(int)
        },
        "/": func(lhs interface{}, rhs interface{})int {
            return lhs.(int) / rhs.(int)
        },
        "%": func(lhs interface{}, rhs interface{})int {
            return lhs.(int) % rhs.(int)
        },
        "+": func(lhs interface{}, rhs interface{})int {
            return lhs.(int) + rhs.(int)
        },
        "-": func(lhs interface{}, rhs interface{})int {
            return lhs.(int) - rhs.(int)
        },
        "<=": func(lhs interface{}, rhs interface{})int {
            if lhs.(int) <= rhs.(int) {
                return 1
            }
            return 0
        },
        "<": func(lhs interface{}, rhs interface{})int {
            if lhs.(int) < rhs.(int) {
                return 1
            }
            return 0
        },
        ">=": func(lhs interface{}, rhs interface{})int {
            if lhs.(int) > rhs.(int) {
                return 1
            }
            return 0
        },
        ">": func(lhs interface{}, rhs interface{})int {
            if lhs.(int) > rhs.(int) {
                return 1
            }
            return 0
        },
        "$=": func(lhs interface{}, rhs interface{})int {
            lhs_str := lhs.(string)
            rhs_str := rhs.(string)
            if strings.LastIndex(lhs_str, rhs_str) == len(lhs_str) - len(rhs_str) {
                return 1
            }
            return 0
        },
        "^=": func(lhs interface{}, rhs interface{})int {
            lhs_str := lhs.(string)
            rhs_str := rhs.(string)
            if strings.Index(lhs_str, rhs_str) == 0 {
                return 1
            }
            return 0
        },
        "*=": func(lhs interface{}, rhs interface{})int {
            lhs_str := lhs.(string)
            rhs_str := rhs.(string)
            if strings.Index(lhs_str, rhs_str) != 0 {
                return 1
            }
            return 0
        },
        "=": func(lhs interface{}, rhs interface{})int {
            if lhs.(string) == rhs.(string) {
                return 1
            }
            return 0
        },
        "!=": func(lhs interface{}, rhs interface{})int {
            if lhs.(string) != rhs.(string) {
                return 1
            }
            return 0
        },
        "&&": func(lhs interface{}, rhs interface{})int {
            if lhs.(bool) && rhs.(bool) {
                return 1
            }
            return 0
        },
        "||": func(lhs interface{}, rhs interface{})int {
            if lhs.(bool) || rhs.(bool) {
                return 1
            }
            return 0
        },
    }
    return p.evaluateParsedExpression(tokens, node, comparatorMap)
}
<|endoftext|>


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

package logger_demo

import (
	"bytes"
	"io"
	"net/http"
	"strconv"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
	"github.com/venicegeo/pz-gocommon/gocommon"
	"github.com/venicegeo/pz-logger/logger"
)

func sleep() {
	time.Sleep(1 * time.Second)
}

type LoggerTester struct {
	suite.Suite
	client *logger.Client
	url    string
}

func (suite *LoggerTester) setupFixture() {
	t := suite.T()
	assert := assert.New(t)

	suite.url = "https://pz-logger.int.geointservices.io"
	client, err := logger.NewClient2(suite.url)
	assert.NoError(err)
	suite.client = client
}

func (suite *LoggerTester) teardownFixture() {
}

func TestRunSuite(t *testing.T) {
	s := &LoggerTester{}
	suite.Run(t, s)
}

func (suite *LoggerTester) getMessages(key string) []logger.Message {
	t := suite.T()
	assert := assert.New(t)

	client := suite.client

	format := &piazza.JsonPagination{
		PerPage: 100,
		Page:    0,
		Order:   piazza.PaginationOrderAscending,
		SortBy:  "createdOn",
	}
	ms, _, err := client.GetMessages(format, nil)
	assert.NoError(err)
	assert.True(len(ms) > 0)

	ret := make([]logger.Message, 0)
	for _, m := range ms {
		if m.Message == key {
			ret = append(ret, m)
		}
	}

	return ret
}

func (suite *LoggerTester) verifyMessageExists(expected *logger.Message) bool {
	t := suite.T()
	assert := assert.New(t)

	client := suite.client

	format := &piazza.JsonPagination{
		PerPage: 10,
		Page:    0,
		Order:   piazza.PaginationOrderAscending,
		SortBy:  "createdOn",
	}
	ms, _, err := client.GetMessages(format, nil)
	assert.NoError(err)
	assert.Len(ms, format.PerPage)

	for _, m := range ms {
		if m.String() == expected.String() {
			return true
		}
	}
	return false
}

func (suite *LoggerTester) TestRawGet() {
	t := suite.T()
	assert := assert.New(t)

	resp, err := http.Get(suite.url + "/message?perPage=13&page=&0")
	assert.NoError(err)
	assert.True(resp.ContentLength >= 0)
	if resp.ContentLength == -1 {
		assert.FailNow("bonk")
	}
	assert.True(resp.ContentLength > 0)

	raw := make([]byte, resp.ContentLength)
	_, err = io.ReadFull(resp.Body, raw)
	defer resp.Body.Close()
	if err != nil && err != io.EOF {
		assert.NoError(err)
	}

	assert.Equal(200, resp.StatusCode)
}

func (suite *LoggerTester) TestRawPost() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	jsn := `
	{
		"address":"XXX",
		"createdOn":"2016-07-22T16:44:58.065583138-04:00",
		"message":"XXX",
		"service":"XXX",
		"severity":"XXX"
	}`
	reader := bytes.NewReader([]byte(jsn))

	resp, err := http.Post(suite.url+"/message",
		piazza.ContentTypeJSON, reader)
	assert.NoError(err)

	raw := make([]byte, resp.ContentLength)
	_, err = io.ReadFull(resp.Body, raw)
	defer resp.Body.Close()
	if err != nil && err != io.EOF {
		assert.NoError(err)
	}

	//err = json.Unmarshal(raw, output)
	//assett.NoError(err)

	assert.Equal(200, resp.StatusCode)
}

func (suite *LoggerTester) TestGet() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	format := &piazza.JsonPagination{
		PerPage: 12,
		Page:    0,
		Order:   piazza.PaginationOrderAscending,
		SortBy:  "createdOn",
	}
	ms, _, err := client.GetMessages(format, nil)
	assert.NoError(err)
	assert.Len(ms, format.PerPage)
}

func (suite *LoggerTester) TestPost() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	var err error

	key := time.Now().String()

	data := &logger.Message{
		Service:   "log-tester",
		Address:   "128.1.2.3",
		CreatedOn: time.Now(),
		Severity:  "Info",
		Message:   key,
	}

	err = client.PostMessage(data)
	assert.NoError(err, "PostToMessages")

	sleep()

	assert.True(suite.verifyMessageExists(data))
}

func (suite *LoggerTester) TestAdmin() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	stats, err := client.GetStats()
	assert.NoError(err, "GetFromAdminStats")
	assert.NotZero(stats.NumMessages)
}

func (suite *LoggerTester) TestPagination() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	format := &piazza.JsonPagination{
		PerPage: 1,
		Page:    0,
		SortBy:  "createdOn",
		Order:   piazza.PaginationOrderAscending,
	}
	params := &piazza.HttpQueryParams{}

	ms, _, err := client.GetMessages(format, params)
	assert.NoError(err)
	assert.Len(ms, format.PerPage)
}

func (suite *LoggerTester) TestParams() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	uniqService := strconv.Itoa(time.Now().Nanosecond())
	uniqDebug := strconv.Itoa(time.Now().Nanosecond() * 3)
	uniqError := strconv.Itoa(time.Now().Nanosecond() * 5)
	uniqFatal := strconv.Itoa(time.Now().Nanosecond() * 7)

	client.SetService(piazza.ServiceName(uniqService), "1.2.3.4")

	now := time.Now()
	sec3 := time.Second * 3
	tstart := now.Add(-sec3)

	client.Debug(uniqDebug)
	client.Error(uniqError)
	client.Fatal(uniqFatal)

	sleep()

	format := &piazza.JsonPagination{
		PerPage: 256,
		Page:    0,
		Order:   piazza.PaginationOrderDescending,
		SortBy:  "createdOn",
	}

	// test date range params
	{
		tend := now.Add(sec3)

		params := &piazza.HttpQueryParams{}
		params.AddTime("after", tstart)
		params.AddTime("before", tend)

		msgs, cnt, err := client.GetMessages(format, params)

		assert.NoError(err)
		assert.True(cnt >= 3)
		assert.True(len(msgs) >= 3)
	}

	// test service param
	{
		params := &piazza.HttpQueryParams{}
		params.AddString("service", uniqService)

		msgs, _, err := client.GetMessages(format, params)

		assert.NoError(err)
		assert.Len(msgs, 3)
	}

	// test contains param
	{
		params := &piazza.HttpQueryParams{}
		params.AddString("contains", uniqDebug)

		msgs, _, err := client.GetMessages(format, params)

		assert.NoError(err)
		assert.True(len(msgs) == 1)
	}
}
<commit_msg>fully functional now!<commit_after>// Copyright 2016, RadiantBlue Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package logger_demo

import (
	"bytes"
	"io"
	"net/http"
	"strconv"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
	"github.com/venicegeo/pz-gocommon/gocommon"
	"github.com/venicegeo/pz-logger/logger"
)

func sleep() {
	time.Sleep(1 * time.Second)
}

type LoggerTester struct {
	suite.Suite
	client *logger.Client
	url    string
}

func (suite *LoggerTester) setupFixture() {
	t := suite.T()
	assert := assert.New(t)

	suite.url = "https://pz-logger.int.geointservices.io"
	client, err := logger.NewClient2(suite.url)
	assert.NoError(err)
	suite.client = client
}

func (suite *LoggerTester) teardownFixture() {
}

func TestRunSuite(t *testing.T) {
	s := &LoggerTester{}
	suite.Run(t, s)
}

func (suite *LoggerTester) verifyMessageExists(expected *logger.Message) bool {
	t := suite.T()
	assert := assert.New(t)

	client := suite.client

	format := &piazza.JsonPagination{
		PerPage: 10,
		Page:    0,
		Order:   piazza.SortOrderDescending,
		SortBy:  "createdOn",
	}
	ms, _, err := client.GetMessages(format, nil)
	assert.NoError(err)
	assert.Len(ms, format.PerPage)

	for _, m := range ms {
		if m.String() == expected.String() {
			return true
		}
	}
	return false
}

// test assume at least two mssgs in the log file, so do this first
func (suite *LoggerTester) Test0000() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	key := strconv.Itoa(time.Now().Nanosecond())

	data := &logger.Message{
		Service:   "log-tester",
		Address:   "128.1.2.3",
		CreatedOn: time.Now(),
		Severity:  "Info",
		Message:   key + "A",
	}

	err := client.PostMessage(data)
	assert.NoError(err)

	data.Message = key + "B"
	data.CreatedOn = time.Now()

	err = client.PostMessage(data)
	assert.NoError(err)
}

func (suite *LoggerTester) TestRawGet() {
	t := suite.T()
	assert := assert.New(t)

	resp, err := http.Get(suite.url + "/message?perPage=13&page=&0")
	assert.NoError(err)
	assert.True(resp.ContentLength >= 0)
	if resp.ContentLength == -1 {
		assert.FailNow("bonk")
	}
	assert.True(resp.ContentLength > 0)

	raw := make([]byte, resp.ContentLength)
	_, err = io.ReadFull(resp.Body, raw)
	defer resp.Body.Close()
	if err != nil && err != io.EOF {
		assert.NoError(err)
	}

	assert.Equal(200, resp.StatusCode)
}

func (suite *LoggerTester) TestRawPost() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	jsn := `
	{
		"address":"XXX",
		"createdOn":"2016-07-22T16:44:58.065583138-04:00",
		"message":"XXX",
		"service":"XXX",
		"severity":"XXX"
	}`
	reader := bytes.NewReader([]byte(jsn))

	resp, err := http.Post(suite.url+"/message",
		piazza.ContentTypeJSON, reader)
	assert.NoError(err)

	raw := make([]byte, resp.ContentLength)
	_, err = io.ReadFull(resp.Body, raw)
	defer resp.Body.Close()
	if err != nil && err != io.EOF {
		assert.NoError(err)
	}

	//err = json.Unmarshal(raw, output)
	//assett.NoError(err)

	assert.Equal(200, resp.StatusCode)
}

func (suite *LoggerTester) TestGet() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	format := &piazza.JsonPagination{
		PerPage: 12,
		Page:    0,
		Order:   piazza.SortOrderAscending,
		SortBy:  "createdOn",
	}
	ms, _, err := client.GetMessages(format, nil)
	assert.NoError(err)
	assert.Len(ms, format.PerPage)
}

func (suite *LoggerTester) TestPost() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	var err error

	key := time.Now().String()

	data := &logger.Message{
		Service:   "log-tester",
		Address:   "128.1.2.3",
		CreatedOn: time.Now(),
		Severity:  "Info",
		Message:   key,
	}

	err = client.PostMessage(data)
	assert.NoError(err, "PostToMessages")

	sleep()

	assert.True(suite.verifyMessageExists(data))
}

func (suite *LoggerTester) TestAdmin() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	stats, err := client.GetStats()
	assert.NoError(err, "GetFromAdminStats")

	assert.NotZero(stats.NumMessages)
}

func (suite *LoggerTester) TestPagination() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	format := &piazza.JsonPagination{
		PerPage: 10,
		Page:    0,
		SortBy:  "createdOn",
		Order:   piazza.SortOrderAscending,
	}
	params := &piazza.HttpQueryParams{}

	// check per-page
	{
		format.PerPage = 17
		ms, _, err := client.GetMessages(format, params)
		assert.NoError(err)
		assert.Len(ms, 17)
	}

	// check sort order
	{
		format.PerPage = 10
		format.Order = piazza.SortOrderAscending
		ms, _, err := client.GetMessages(format, params)
		assert.NoError(err)
		last := len(ms) - 1
		assert.True(last <= 9)

		// we can't check "before", because two createdOn's might be the same
		isBefore := ms[0].CreatedOn.Before(ms[last].CreatedOn)
		isEqual := ms[0].CreatedOn.Equal(ms[last].CreatedOn)
		assert.True(isBefore || isEqual)

		format.Order = piazza.SortOrderDescending
		ms, _, err = client.GetMessages(format, params)
		assert.NoError(err)
		last = len(ms) - 1
		assert.True(last <= 9)

		isAfter := ms[0].CreatedOn.After(ms[last].CreatedOn)
		isEqual = ms[0].CreatedOn.Equal(ms[last].CreatedOn)
		assert.True(isAfter || isEqual)
	}

	// check sort-by
	{
		format.Order = piazza.SortOrderAscending
		format.SortBy = "severity"
		format.PerPage = 100
		format.Page = 0
		ms, _, err := client.GetMessages(format, params)
		if err != nil {
			panic(88)
		}
		assert.NoError(err)

		last := len(ms) - 1
		for i := 0; i < last; i++ {
			a, b := string(ms[i].Severity), string(ms[i+1].Severity)
			isBefore := (a <= b)
			assert.True(isBefore)
		}
	}
}

func (suite *LoggerTester) TestParams() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	client := suite.client

	uniqService := strconv.Itoa(time.Now().Nanosecond())
	uniqDebug := strconv.Itoa(time.Now().Nanosecond() * 3)
	uniqError := strconv.Itoa(time.Now().Nanosecond() * 5)
	uniqFatal := strconv.Itoa(time.Now().Nanosecond() * 7)

	client.SetService(piazza.ServiceName(uniqService), "1.2.3.4")

	now := time.Now()
	sec3 := time.Second * 3
	tstart := now.Add(-sec3)

	client.Debug(uniqDebug)
	client.Error(uniqError)
	client.Fatal(uniqFatal)

	sleep()

	format := &piazza.JsonPagination{
		PerPage: 256,
		Page:    0,
		Order:   piazza.SortOrderDescending,
		SortBy:  "createdOn",
	}

	// test date range params
	{
		tend := now.Add(sec3)

		params := &piazza.HttpQueryParams{}
		params.AddTime("after", tstart)
		params.AddTime("before", tend)

		msgs, cnt, err := client.GetMessages(format, params)

		assert.NoError(err)
		assert.True(cnt >= 3)
		assert.True(len(msgs) >= 3)
	}

	// test service param
	{
		params := &piazza.HttpQueryParams{}
		params.AddString("service", uniqService)

		msgs, _, err := client.GetMessages(format, params)

		assert.NoError(err)
		assert.Len(msgs, 3)
	}

	// test contains param
	{
		params := &piazza.HttpQueryParams{}
		params.AddString("contains", uniqDebug)

		msgs, _, err := client.GetMessages(format, params)

		assert.NoError(err)
		assert.True(len(msgs) == 1)
	}
}
<|endoftext|>


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

package config

import (
	"context"

	"google.golang.org/protobuf/proto"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/config"
	"go.chromium.org/luci/config/cfgclient"
	"go.chromium.org/luci/gae/service/datastore"
	"go.chromium.org/luci/server/tq"

	pb "go.chromium.org/luci/cv/api/config/v2"
)

const configFileName = "commit-queue.cfg"

func init() {
	tq.RegisterTaskClass(tq.TaskClass{
		ID:        "refresh-project-config",
		Prototype: &RefreshProjectConfigTask{},
		Queue:     "refresh-project-config",
		Handler: func(ctx context.Context, payload proto.Message) error {
			task := payload.(*RefreshProjectConfigTask)
			action, actionFn := "update", updateProject
			if task.GetDisable() {
				action, actionFn = "disable", disableProject
			}
			project := task.GetProject()
			if err := actionFn(ctx, project); err != nil {
				errors.Log(ctx, err)
				// Explicitly not forwarding transient tag so that tq won't retry
				// because the refresh task is submitted every minute by AppEngine
				// Cron.
				return errors.Reason("failed to %s project %q: %s", action, project, err).Err()
			}
			return nil
		},
	})
}

// SubmitRefreshTasks submits tasks that update config for LUCI projects
// or disable projects that do not have CV config in LUCI Config.
func SubmitRefreshTasks(ctx context.Context) error {
	projects, err := cfgclient.ProjectsWithConfig(ctx, configFileName)
	if err != nil {
		return errors.Annotate(err, "failed to get projects that have %q from LUCI Config", configFileName).Tag(transient.Tag).Err()
	}
	tasks := make([]*tq.Task, len(projects))
	for i, p := range projects {
		tasks[i] = &tq.Task{
			Payload: &RefreshProjectConfigTask{
				Project: p,
			},
		}
	}

	curEnabledProjects, err := getAllProjectIDs(ctx, true)
	if err != nil {
		return err
	}
	projectsInLUCIConfig := stringset.NewFromSlice(projects...)
	for _, p := range curEnabledProjects {
		if !projectsInLUCIConfig.Has(p) {
			tasks = append(tasks, &tq.Task{
				Payload: &RefreshProjectConfigTask{
					Project: p,
					Disable: true,
				},
			})
		}
	}

	err = parallel.WorkPool(32, func(workCh chan<- func() error) {
		for _, task := range tasks {
			workCh <- func() error {
				if err := tq.AddTask(ctx, task); err != nil {
					logging.WithError(err).Errorf(ctx, "Failed to submit refresh task for project %q", task.Payload.(*RefreshProjectConfigTask).GetProject())
					return err
				}
				return nil
			}
		}
	})

	if err != nil {
		return err.(errors.MultiError).First()
	}
	return nil
}

// updateProject imports the latest CV Config for a given LUCI Project
// from LUCI Config if the config in CV is outdated.
func updateProject(ctx context.Context, project string) error {
	externalHash, err := getExternalContentHash(ctx, project)
	if err != nil {
		return err
	}
	existingPC := ProjectConfig{Project: project}
	switch err := datastore.Get(ctx, &existingPC); {
	case err != nil && err != datastore.ErrNoSuchEntity:
		return errors.Annotate(err, "failed to get ProjectConfig(project=%q)", project).Tag(transient.Tag).Err()
	case !existingPC.Enabled:
		// Go through update process to ensure all configs are present.
	case existingPC.ExternalHash == externalHash:
		return nil // up-to-date
	}

	cfg, err := fetchCfg(ctx, project, externalHash)
	if err != nil {
		return err
	}

	// Write out ConfigHashInfo if missing and all ConfigGroups.
	localHash := computeHash(cfg)
	cgNames := make([]string, len(cfg.GetConfigGroups()))
	for i, cg := range cfg.GetConfigGroups() {
		cgNames[i] = makeConfigGroupName(cg.GetName(), i)
	}
	targetEVersion := existingPC.EVersion + 1

	err = datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		hashInfo := ConfigHashInfo{
			Hash:    localHash,
			Project: datastore.MakeKey(ctx, projectConfigKind, project),
		}
		switch err := datastore.Get(ctx, &hashInfo); {
		case err != nil && err != datastore.ErrNoSuchEntity:
			return errors.Annotate(err, "failed to get ConfigHashInfo(Hash=%q)", localHash).Tag(transient.Tag).Err()
		case err == nil && hashInfo.ProjectEVersion >= targetEVersion:
			return nil // Do not go backwards.
		default:
			hashInfo.ProjectEVersion = targetEVersion
			hashInfo.UpdateTime = datastore.RoundTime(clock.Now(ctx)).UTC()
			hashInfo.ConfigGroupNames = cgNames
			return errors.Annotate(datastore.Put(ctx, &hashInfo), "failed to put ConfigHashInfo(Hash=%q)", localHash).Tag(transient.Tag).Err()
		}
	}, nil)
	if err != nil {
		return errors.Annotate(err, "failed to run transaction to update ConfigHashInfo").Tag(transient.Tag).Err()
	}

	if err := putConfigGroups(ctx, cfg, project, localHash); err != nil {
		return err
	}

	updated := false
	err = datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		updated = false
		pc := ProjectConfig{Project: project}
		switch err := datastore.Get(ctx, &pc); {
		case err != nil && err != datastore.ErrNoSuchEntity:
			return errors.Annotate(err, "failed to get ProjectConfig(project=%q)", project).Tag(transient.Tag).Err()
		case pc.EVersion != existingPC.EVersion:
			return nil // Already updated by concurrent updateProject.
		default:
			pc = ProjectConfig{
				Project:          project,
				Enabled:          true,
				UpdateTime:       datastore.RoundTime(clock.Now(ctx)).UTC(),
				EVersion:         targetEVersion,
				Hash:             localHash,
				ExternalHash:     externalHash,
				ConfigGroupNames: cgNames,
			}
			updated = true
			return errors.Annotate(datastore.Put(ctx, &pc), "failed to put ProjectConfig(project=%q)", project).Tag(transient.Tag).Err()
			// TODO(yiwzhang): Notify ProjectManager
		}
	}, nil)

	switch {
	case err != nil:
		return errors.Annotate(err, "failed to run transaction to update ProjectConfig").Tag(transient.Tag).Err()
	case updated:
		logging.Infof(ctx, "updated project %q to rev %s hash %s ", project, "TODO", localHash)
	}
	return nil
}

func getExternalContentHash(ctx context.Context, project string) (string, error) {
	var meta config.Meta
	switch err := cfgclient.Get(ctx, config.ProjectSet(project), configFileName, nil, &meta); {
	case err != nil:
		return "", errors.Annotate(err, "failed to fetch meta from LUCI Config").Tag(transient.Tag).Err()
	case meta.ContentHash == "":
		return "", errors.Reason("LUCI Config returns empty content hash for project %q", project).Err()
	default:
		return meta.ContentHash, nil
	}
}

func fetchCfg(ctx context.Context, project string, contentHash string) (*pb.Config, error) {
	content, err := cfgclient.Client(ctx).GetConfigByHash(ctx, contentHash)
	if err != nil {
		return nil, errors.Annotate(err, "failed to get config by content hash").Tag(transient.Tag).Err()
	}
	ret := &pb.Config{}
	if err := cfgclient.ProtoText(ret)(content); err != nil {
		return nil, errors.Annotate(err, "failed to deserialize config content").Err()
	}
	// TODO(yiwzhang): validate the config here again to prevent ingesting a bad
	// version of config that accidentally slips into LUCI Config.
	// See: go.chromium.org/luci/cq/appengine/config
	return ret, nil
}

// disableProject disables the given LUCI Project if it is currently enabled.
func disableProject(ctx context.Context, project string) error {
	disabled := false

	err := datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		disabled = false
		pc := ProjectConfig{Project: project}
		switch err := datastore.Get(ctx, &pc); {
		case datastore.IsErrNoSuchEntity(err):
			return nil // No-op when disabling non-existent Project
		case err != nil:
			return errors.Annotate(err, "failed to get existing ProjectConfig").Tag(transient.Tag).Err()
		case !pc.Enabled:
			return nil // Already disabled
		}
		pc.Enabled = false
		pc.UpdateTime = datastore.RoundTime(clock.Now(ctx)).UTC()
		pc.EVersion++
		if err := datastore.Put(ctx, &pc); err != nil {
			return errors.Annotate(err, "failed to put ProjectConfig").Tag(transient.Tag).Err()
		}
		disabled = true
		// TODO(yiwzhang): Notify ProjectManager
		return nil
	}, nil)

	switch {
	case err != nil:
		return errors.Annotate(err, "failed to run transaction to disable project %q", project).Tag(transient.Tag).Err()
	case disabled:
		logging.Infof(ctx, "disabled project %q", project)
	}
	return nil
}
<commit_msg>[cv] Use local variable to fix datarace<commit_after>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package config

import (
	"context"

	"google.golang.org/protobuf/proto"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/config"
	"go.chromium.org/luci/config/cfgclient"
	"go.chromium.org/luci/gae/service/datastore"
	"go.chromium.org/luci/server/tq"

	pb "go.chromium.org/luci/cv/api/config/v2"
)

const configFileName = "commit-queue.cfg"

func init() {
	tq.RegisterTaskClass(tq.TaskClass{
		ID:        "refresh-project-config",
		Prototype: &RefreshProjectConfigTask{},
		Queue:     "refresh-project-config",
		Handler: func(ctx context.Context, payload proto.Message) error {
			task := payload.(*RefreshProjectConfigTask)
			action, actionFn := "update", updateProject
			if task.GetDisable() {
				action, actionFn = "disable", disableProject
			}
			project := task.GetProject()
			if err := actionFn(ctx, project); err != nil {
				errors.Log(ctx, err)
				// Explicitly not forwarding transient tag so that tq won't retry
				// because the refresh task is submitted every minute by AppEngine
				// Cron.
				return errors.Reason("failed to %s project %q: %s", action, project, err).Err()
			}
			return nil
		},
	})
}

// SubmitRefreshTasks submits tasks that update config for LUCI projects
// or disable projects that do not have CV config in LUCI Config.
func SubmitRefreshTasks(ctx context.Context) error {
	projects, err := cfgclient.ProjectsWithConfig(ctx, configFileName)
	if err != nil {
		return errors.Annotate(err, "failed to get projects that have %q from LUCI Config", configFileName).Tag(transient.Tag).Err()
	}
	tasks := make([]*tq.Task, len(projects))
	for i, p := range projects {
		tasks[i] = &tq.Task{
			Payload: &RefreshProjectConfigTask{
				Project: p,
			},
		}
	}

	curEnabledProjects, err := getAllProjectIDs(ctx, true)
	if err != nil {
		return err
	}
	projectsInLUCIConfig := stringset.NewFromSlice(projects...)
	for _, p := range curEnabledProjects {
		if !projectsInLUCIConfig.Has(p) {
			tasks = append(tasks, &tq.Task{
				Payload: &RefreshProjectConfigTask{
					Project: p,
					Disable: true,
				},
			})
		}
	}

	err = parallel.WorkPool(32, func(workCh chan<- func() error) {
		for _, task := range tasks {
			task := task
			workCh <- func() error {
				if err := tq.AddTask(ctx, task); err != nil {
					logging.WithError(err).Errorf(ctx, "Failed to submit refresh task for project %q", task.Payload.(*RefreshProjectConfigTask).GetProject())
					return err
				}
				return nil
			}
		}
	})

	if err != nil {
		return err.(errors.MultiError).First()
	}
	return nil
}

// updateProject imports the latest CV Config for a given LUCI Project
// from LUCI Config if the config in CV is outdated.
func updateProject(ctx context.Context, project string) error {
	externalHash, err := getExternalContentHash(ctx, project)
	if err != nil {
		return err
	}
	existingPC := ProjectConfig{Project: project}
	switch err := datastore.Get(ctx, &existingPC); {
	case err != nil && err != datastore.ErrNoSuchEntity:
		return errors.Annotate(err, "failed to get ProjectConfig(project=%q)", project).Tag(transient.Tag).Err()
	case !existingPC.Enabled:
		// Go through update process to ensure all configs are present.
	case existingPC.ExternalHash == externalHash:
		return nil // up-to-date
	}

	cfg, err := fetchCfg(ctx, project, externalHash)
	if err != nil {
		return err
	}

	// Write out ConfigHashInfo if missing and all ConfigGroups.
	localHash := computeHash(cfg)
	cgNames := make([]string, len(cfg.GetConfigGroups()))
	for i, cg := range cfg.GetConfigGroups() {
		cgNames[i] = makeConfigGroupName(cg.GetName(), i)
	}
	targetEVersion := existingPC.EVersion + 1

	err = datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		hashInfo := ConfigHashInfo{
			Hash:    localHash,
			Project: datastore.MakeKey(ctx, projectConfigKind, project),
		}
		switch err := datastore.Get(ctx, &hashInfo); {
		case err != nil && err != datastore.ErrNoSuchEntity:
			return errors.Annotate(err, "failed to get ConfigHashInfo(Hash=%q)", localHash).Tag(transient.Tag).Err()
		case err == nil && hashInfo.ProjectEVersion >= targetEVersion:
			return nil // Do not go backwards.
		default:
			hashInfo.ProjectEVersion = targetEVersion
			hashInfo.UpdateTime = datastore.RoundTime(clock.Now(ctx)).UTC()
			hashInfo.ConfigGroupNames = cgNames
			return errors.Annotate(datastore.Put(ctx, &hashInfo), "failed to put ConfigHashInfo(Hash=%q)", localHash).Tag(transient.Tag).Err()
		}
	}, nil)
	if err != nil {
		return errors.Annotate(err, "failed to run transaction to update ConfigHashInfo").Tag(transient.Tag).Err()
	}

	if err := putConfigGroups(ctx, cfg, project, localHash); err != nil {
		return err
	}

	updated := false
	err = datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		updated = false
		pc := ProjectConfig{Project: project}
		switch err := datastore.Get(ctx, &pc); {
		case err != nil && err != datastore.ErrNoSuchEntity:
			return errors.Annotate(err, "failed to get ProjectConfig(project=%q)", project).Tag(transient.Tag).Err()
		case pc.EVersion != existingPC.EVersion:
			return nil // Already updated by concurrent updateProject.
		default:
			pc = ProjectConfig{
				Project:          project,
				Enabled:          true,
				UpdateTime:       datastore.RoundTime(clock.Now(ctx)).UTC(),
				EVersion:         targetEVersion,
				Hash:             localHash,
				ExternalHash:     externalHash,
				ConfigGroupNames: cgNames,
			}
			updated = true
			return errors.Annotate(datastore.Put(ctx, &pc), "failed to put ProjectConfig(project=%q)", project).Tag(transient.Tag).Err()
			// TODO(yiwzhang): Notify ProjectManager
		}
	}, nil)

	switch {
	case err != nil:
		return errors.Annotate(err, "failed to run transaction to update ProjectConfig").Tag(transient.Tag).Err()
	case updated:
		logging.Infof(ctx, "updated project %q to rev %s hash %s ", project, "TODO", localHash)
	}
	return nil
}

func getExternalContentHash(ctx context.Context, project string) (string, error) {
	var meta config.Meta
	switch err := cfgclient.Get(ctx, config.ProjectSet(project), configFileName, nil, &meta); {
	case err != nil:
		return "", errors.Annotate(err, "failed to fetch meta from LUCI Config").Tag(transient.Tag).Err()
	case meta.ContentHash == "":
		return "", errors.Reason("LUCI Config returns empty content hash for project %q", project).Err()
	default:
		return meta.ContentHash, nil
	}
}

func fetchCfg(ctx context.Context, project string, contentHash string) (*pb.Config, error) {
	content, err := cfgclient.Client(ctx).GetConfigByHash(ctx, contentHash)
	if err != nil {
		return nil, errors.Annotate(err, "failed to get config by content hash").Tag(transient.Tag).Err()
	}
	ret := &pb.Config{}
	if err := cfgclient.ProtoText(ret)(content); err != nil {
		return nil, errors.Annotate(err, "failed to deserialize config content").Err()
	}
	// TODO(yiwzhang): validate the config here again to prevent ingesting a bad
	// version of config that accidentally slips into LUCI Config.
	// See: go.chromium.org/luci/cq/appengine/config
	return ret, nil
}

// disableProject disables the given LUCI Project if it is currently enabled.
func disableProject(ctx context.Context, project string) error {
	disabled := false

	err := datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		disabled = false
		pc := ProjectConfig{Project: project}
		switch err := datastore.Get(ctx, &pc); {
		case datastore.IsErrNoSuchEntity(err):
			return nil // No-op when disabling non-existent Project
		case err != nil:
			return errors.Annotate(err, "failed to get existing ProjectConfig").Tag(transient.Tag).Err()
		case !pc.Enabled:
			return nil // Already disabled
		}
		pc.Enabled = false
		pc.UpdateTime = datastore.RoundTime(clock.Now(ctx)).UTC()
		pc.EVersion++
		if err := datastore.Put(ctx, &pc); err != nil {
			return errors.Annotate(err, "failed to put ProjectConfig").Tag(transient.Tag).Err()
		}
		disabled = true
		// TODO(yiwzhang): Notify ProjectManager
		return nil
	}, nil)

	switch {
	case err != nil:
		return errors.Annotate(err, "failed to run transaction to disable project %q", project).Tag(transient.Tag).Err()
	case disabled:
		logging.Infof(ctx, "disabled project %q", project)
	}
	return nil
}
<|endoftext|>


<commit_before>package controllers

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"getmelange.com/app/framework"
	"getmelange.com/app/models"
	"github.com/gorilla/websocket"
	gdb "github.com/huntaub/go-db"
)

var wsUpgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return r.Header["Origin"][0] == "http://app.melange:7776"
	},
}
var quitChan = make(chan struct{})

// RealtimeHandler connects one websocket-enabled client to the backend go server.
type RealtimeHandler struct {
	Store  *models.Store
	Tables map[string]gdb.Table

	messageChan chan *models.JSONMessage
	dataChan    chan interface{}
}

// UpgradeConnection will upgrade a normal HTTP Request to the Websocket
// protocol.
func (r *RealtimeHandler) UpgradeConnection(res http.ResponseWriter, req *http.Request) {
	fmt.Println("Requesting a WebSocket Upgrade")
	conn, err := wsUpgrader.Upgrade(res, req, nil)
	if err != nil {
		fmt.Println("WebSocket Upgrade Error", err)

		// Send a 500 on Error
		framework.WriteView(
			&framework.HTTPError{
				ErrorCode: 500,
				Message:   "Websocket upgrade failed.",
			}, res,
		)
		return
	}

	r.dataChan = make(chan interface{})
	r.messageChan = make(chan *models.JSONMessage)

	// Attach the websocket to the fetcher.
	models.AddFetchWatcher(r.messageChan)

	go func(mes <-chan *models.JSONMessage, data <-chan interface{}, quit <-chan struct{}) {
		for {
			select {
			case m := <-data:
				err := conn.WriteJSON(m)
				if err != nil {
					fmt.Println("Error writing to WS", err)
				}
			case m := <-mes:
				err := conn.WriteJSON(map[string]interface{}{
					"type": "message",
					"data": m,
				})
				if err != nil {
					fmt.Println("Error writing to WS", err)
				}
			case <-quit:
				return
			}
		}
	}(r.messageChan, r.dataChan, quitChan)

	go func(conn *websocket.Conn, dataChan chan interface{}) {
		for {
			v := make(map[string]interface{})

			fmt.Println("Waiting for new message.")
			_, p, err := conn.ReadMessage()
			if err != nil {
				if err == io.EOF {
					fmt.Println("Lost connection from websocket.")
					return
				}

				fmt.Println("Error reading from websocket.", err)
				continue
			}

			err = json.Unmarshal(p, &v)
			if err != nil {
				fmt.Println("Error unmarshalling from websocket", err)
				continue
			}

			typ, ok1 := v["type"]
			data, ok2 := v["data"]
			if !ok1 || !ok2 {
				fmt.Println("Websocket doesn't have correct message type.")
				continue
			}

			typString, ok := typ.(string)
			if !ok {
				fmt.Println("Websocket message type has incorrect format.")
				continue
			}

			outTyp, outData := r.HandleWSRequest(typString, data)

			dataChan <- map[string]interface{}{
				"type": outTyp,
				"data": outData,
			}
		}
	}(conn, r.dataChan)
}

// HandleWSRequest will handle one WS request from the client.
func (r *RealtimeHandler) HandleWSRequest(t string, d interface{}) (string, interface{}) {
	if t == "startup" {
		// We need to send all current messages to the client. Initiate transfer.
		m, err := constructManager(r.Store, r.Tables)
		if err != nil {
			fmt.Println("Unable to construct WS Message Manager", err)
		}

		m.GetAllMessages(r.messageChan)
		return "initDone", nil
	}
	fmt.Println("Got message", t, "from ws.")
	return "gotIt", nil
}
<commit_msg>Added Extra Error Checking for Websockets<commit_after>package controllers

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"getmelange.com/app/framework"
	"getmelange.com/app/models"
	"github.com/gorilla/websocket"
	gdb "github.com/huntaub/go-db"
)

var wsUpgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return r.Header["Origin"][0] == "http://app.melange:7776"
	},
}
var quitChan = make(chan struct{})

// RealtimeHandler connects one websocket-enabled client to the backend go server.
type RealtimeHandler struct {
	Store  *models.Store
	Tables map[string]gdb.Table

	messageChan chan *models.JSONMessage
	dataChan    chan interface{}
}

// UpgradeConnection will upgrade a normal HTTP Request to the Websocket
// protocol.
func (r *RealtimeHandler) UpgradeConnection(res http.ResponseWriter, req *http.Request) {
	fmt.Println("Requesting a WebSocket Upgrade")
	conn, err := wsUpgrader.Upgrade(res, req, nil)
	if err != nil {
		fmt.Println("WebSocket Upgrade Error", err)

		// Send a 500 on Error
		framework.WriteView(
			&framework.HTTPError{
				ErrorCode: 500,
				Message:   "Websocket upgrade failed.",
			}, res,
		)
		return
	}

	r.dataChan = make(chan interface{})
	r.messageChan = make(chan *models.JSONMessage)

	// Attach the websocket to the fetcher.
	models.AddFetchWatcher(r.messageChan)

	go func(mes <-chan *models.JSONMessage, data <-chan interface{}, quit <-chan struct{}) {
		for {
			select {
			case m := <-data:
				err := conn.WriteJSON(m)
				if err != nil {
					fmt.Println("Error writing to WS", err)
				}
			case m := <-mes:
				err := conn.WriteJSON(map[string]interface{}{
					"type": "message",
					"data": m,
				})
				if err != nil {
					fmt.Println("Error writing to WS", err)
				}
			case <-quit:
				return
			}
		}
	}(r.messageChan, r.dataChan, quitChan)

	go func(conn *websocket.Conn, dataChan chan interface{}) {
		for {
			v := make(map[string]interface{})

			fmt.Println("Waiting for new message.")
			_, p, err := conn.ReadMessage()
			if err != nil {
				if err == io.EOF {
					fmt.Println("Lost connection from websocket.")
					return
				} else if err.Error() == "unexpected EOF" {
					fmt.Println("Lost connection from websocket.")
					return
				}

				fmt.Println("Error reading from websocket.", err)
				continue
			}

			err = json.Unmarshal(p, &v)
			if err != nil {
				fmt.Println("Error unmarshalling from websocket", err)
				continue
			}

			typ, ok1 := v["type"]
			data, ok2 := v["data"]
			if !ok1 || !ok2 {
				fmt.Println("Websocket doesn't have correct message type.")
				continue
			}

			typString, ok := typ.(string)
			if !ok {
				fmt.Println("Websocket message type has incorrect format.")
				continue
			}

			outTyp, outData := r.HandleWSRequest(typString, data)

			dataChan <- map[string]interface{}{
				"type": outTyp,
				"data": outData,
			}
		}
	}(conn, r.dataChan)
}

// HandleWSRequest will handle one WS request from the client.
func (r *RealtimeHandler) HandleWSRequest(t string, d interface{}) (string, interface{}) {
	if t == "startup" {
		// We need to send all current messages to the client. Initiate transfer.
		m, err := constructManager(r.Store, r.Tables)
		if err != nil {
			fmt.Println("Unable to construct WS Message Manager", err)
		}

		m.GetAllMessages(r.messageChan)
		return "initDone", nil
	}
	fmt.Println("Got message", t, "from ws.")
	return "gotIt", nil
}
<|endoftext|>


<commit_before>package rhynock

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

// Some defaults for pinging
// Needs to be settable from outside
var (
	writeWait = 10 * time.Second
	pongWait = 60 * time.Second
	pingPeriod = (pongWait * 9) / 10
	maxMessageSize int64 = 512
)

// Conn encapsulates our websocket
type Conn struct {
	// Exported actual webscoket so can be messed with from outside
	Ws        *websocket.Conn
	Send      chan []byte
	Dst       BottleDst
	Quit      chan []byte
	wait      chan bool
	valid     bool

	// Ping Calculation
	// TODO make sure pings pair off, right now really bad lagg would case ping ping pong to give a false ping value
	lastping  int64
	ping      int64
	pinglock  sync.Mutex
}

//
// Change global properties of rhynock
// Options are: writeWait, pongWait, pingPeriod, maxMessageSize
// Value is a string that will be parsed into the correct type.
// For writeWait, pongWait, and pingPeriod a duration string is expected, for maxMessageSize an int string.
//
func SetProperty(key, value string) {
	switch key {
	case "writeWait":
		val, err := time.ParseDuration(value)
		if err == nil {
			writeWait = val
		} else {
			log.Println(err)
		}
	case "pongWait":
		val, err := time.ParseDuration(value)
		if err == nil {
			pongWait = val
		} else {
			log.Println(err)
		}
	case "pingPeriod":
		val, err := time.ParseDuration(value)
		if err == nil {
			pingPeriod = val
		} else {
			log.Println(err)
		}
	case "maxMessageSize":
		val, err := strconv.ParseInt(value, 10, 64)
		if err == nil {
			maxMessageSize = val
		} else {
			log.Println(err)
		}
	}
}

//
// Convenience function so you dont have to use the Send channel
//
func (c *Conn) SendMsg(message string) {
	// Basically just typecasting for convenience
	if c.valid {
		c.Send <- []byte(message)
	} else {
		log.Println("rhynock.Conn.SendMsg: Socket closed.")
	}
}

//
// Convenience function to call the quit channel with a message
//
func (c *Conn) CloseMsg(message string) {
	if c.valid {
		c.Quit <- []byte(message)
	} else {
		log.Println("rhynock.Conn.CloseMsg: Socket already closed.")
	}
}

//
// This function chews through the power cables
//
func (c *Conn) Close() {
	// Send ourself the quit signal with no message
	if c.valid {
		c.Quit <- []byte("")
	} else {
		log.Println("rhynock.Conn.Close: Socket already closed.")
	}
}

//
// Get the current ping of this connection
//
func (c *Conn) GetPing() int64 {
	c.pinglock.Lock()
	defer c.pinglock.Unlock()
	return c.ping
}

//
// Used to write a single message to the client and report any errors
//
func (c *Conn) write(t int, payload []byte) error {
	c.Ws.SetWriteDeadline(time.Now().Add(writeWait))
	return c.Ws.WriteMessage(t, payload)
}

//
// Maintains both a reader and a writer, cleans up both if one fails
//
func (c *Conn) read_write() {
	// Ping timer
	ticker := time.NewTicker(pingPeriod)

	// Clean up Connection and Connection resources
	defer func() {
		ticker.Stop()
		c.valid = false
		close(c.Send)
		close(c.Quit)
		c.Ws.Close()
		close(c.wait) // Unblocks the ConnectionHandler
	}()

	// Config websocket settings
	c.Ws.SetReadLimit(maxMessageSize)
	c.Ws.SetReadDeadline(time.Now().Add(pongWait))
	c.Ws.SetPongHandler(func(string) error {
		// Give each client pongWait seconds after the ping to respond
		c.Ws.SetReadDeadline(time.Now().Add(pongWait))
		c.pinglock.Lock()
		c.ping = (time.Now().UnixNano() - c.lastping) / 2000000
		c.pinglock.Unlock()
		return nil
	})

	// Start a reading goroutine
	// The reader will stop when the c.Ws.Close is called at
	// in the defered cleanup function, so we do not manually
	// have to close the reader
	go func() {
		for {
			// This blocks until it reads EOF or an error
			// occurs trying to read, the error can be
			// used to detect when the client closes the Connection
			_, message, err := c.Ws.ReadMessage()
			if err != nil {
				break // If we get an error escape the loop
			}

			// Bottle the message with its sender
			bottle := &Bottle{
				Sender: c,
				Message: message,
			}

			// Send to the destination for processing
			c.Dst.GetBottleChan() <- bottle
		}
		// The reader has been terminated, alert the writer if needed
		if c.valid {
			c.Close()
		}
	}()

	// Main writing loop
	for {
		select {
		case message, ok := <- c.Send:
			// Our send channel has something in it or the channel closed
			if !ok {
				// Our channel was closed, gracefully close socket Conn
				c.write(websocket.CloseMessage, []byte{})
				return
			}
			// Attempt to write the message to the websocket
			if err := c.write(websocket.TextMessage, message); err != nil {
				// If we get an error we can no longer communcate with client
				// return, no need to send CloseMessage since that would
				// just yield another error
				return
			}

		case <- ticker.C:
			// Record timestamp of ping
			c.lastping = time.Now().UnixNano()

			// Ping ticker went off. We need to ping to check for connectivity.
			if err := c.write(websocket.PingMessage, []byte{}); err != nil {
				// We got an error pinging, return and call defer
				// defer will close the socket which will kill the reader
				return
			}

		case bytes := <- c.Quit:
			// Close connection and send a final message
			c.write(websocket.TextMessage, bytes)
			c.write(websocket.CloseMessage, []byte{})
			return
		}
	}
}

var upgrader = &websocket.Upgrader{ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: func(r* http.Request) bool { return true }}

//
// Hanlder function to start a websocket connection
//
func ConnectionHandler(w http.ResponseWriter, r *http.Request, dst BottleDst) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	// Create new connection object
	c := &Conn{
		Send: make(chan []byte, 256),
		Ws:       ws,
		Dst:      dst,
		Quit:     make(chan []byte),
		valid:    true,
		wait:     make(chan bool),
		lastping: 0,
		ping:     0,
	}

	// Alert the destination that a new connection has opened
	dst.ConnectionOpened(c)

	// Make sure connections that are opened will get closed
	defer func() {
		dst.ConnectionClosed(c)
	}()

	// Start read write loop blocking
	c.read_write()
}
<commit_msg>Reroute for hydra<commit_after>package rhynock

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

// Some defaults for pinging
// Needs to be settable from outside
var (
	writeWait = 10 * time.Second
	pongWait = 60 * time.Second
	pingPeriod = (pongWait * 9) / 10
	maxMessageSize int64 = 512
)

// Conn encapsulates our websocket
type Conn struct {
	// Exported actual webscoket so can be messed with from outside
	Ws        *websocket.Conn
	Send      chan []byte
	Dst       BottleDst
	Quit      chan []byte
	valid     bool

	// Ping Calculation
	// TODO make sure pings pair off, right now really bad lagg would case ping ping pong to give a false ping value
	lastping  int64
	ping      int64
	pinglock  sync.Mutex
}

//
// Change global properties of rhynock
// Options are: writeWait, pongWait, pingPeriod, maxMessageSize
// Value is a string that will be parsed into the correct type.
// For writeWait, pongWait, and pingPeriod a duration string is expected, for maxMessageSize an int string.
//
func SetProperty(key, value string) {
	switch key {
	case "writeWait":
		val, err := time.ParseDuration(value)
		if err == nil {
			writeWait = val
		} else {
			log.Println(err)
		}
	case "pongWait":
		val, err := time.ParseDuration(value)
		if err == nil {
			pongWait = val
		} else {
			log.Println(err)
		}
	case "pingPeriod":
		val, err := time.ParseDuration(value)
		if err == nil {
			pingPeriod = val
		} else {
			log.Println(err)
		}
	case "maxMessageSize":
		val, err := strconv.ParseInt(value, 10, 64)
		if err == nil {
			maxMessageSize = val
		} else {
			log.Println(err)
		}
	}
}

//
// Convenience function so you dont have to use the Send channel
//
func (c *Conn) SendMsg(message string) {
	// Basically just typecasting for convenience
	if c.valid {
		c.Send <- []byte(message)
	} else {
		log.Println("rhynock.Conn.SendMsg: Socket closed.")
	}
}

//
// Convenience function to call the quit channel with a message
//
func (c *Conn) CloseMsg(message string) {
	if c.valid {
		c.Quit <- []byte(message)
	} else {
		log.Println("rhynock.Conn.CloseMsg: Socket already closed.")
	}
}

//
// This function chews through the power cables
//
func (c *Conn) Close() {
	// Send ourself the quit signal with no message
	if c.valid {
		c.Quit <- []byte("")
	} else {
		log.Println("rhynock.Conn.Close: Socket already closed.")
	}
}

//
// Get the current ping of this connection
//
func (c *Conn) GetPing() int64 {
	c.pinglock.Lock()
	defer c.pinglock.Unlock()
	return c.ping
}

//
// Used to write a single message to the client and report any errors
//
func (c *Conn) write(t int, payload []byte) error {
	c.Ws.SetWriteDeadline(time.Now().Add(writeWait))
	return c.Ws.WriteMessage(t, payload)
}

//
// Maintains both a reader and a writer, cleans up both if one fails
//
func (c *Conn) read_write() {
	// Ping timer
	ticker := time.NewTicker(pingPeriod)

	// Clean up Connection and Connection resources
	defer func() {
		ticker.Stop()
		c.valid = false
		c.Dst.ConnectionClosed(c)
		close(c.Send)
		close(c.Quit)
		c.Ws.Close()
	}()

	// Config websocket settings
	c.Ws.SetReadLimit(maxMessageSize)
	c.Ws.SetReadDeadline(time.Now().Add(pongWait))
	c.Ws.SetPongHandler(func(string) error {
		// Give each client pongWait seconds after the ping to respond
		c.Ws.SetReadDeadline(time.Now().Add(pongWait))
		c.pinglock.Lock()
		c.ping = (time.Now().UnixNano() - c.lastping) / 2000000
		c.pinglock.Unlock()
		return nil
	})

	// Start a reading goroutine
	// The reader will stop when the c.Ws.Close is called at
	// in the defered cleanup function, so we do not manually
	// have to close the reader
	go func() {
		for {
			// This blocks until it reads EOF or an error
			// occurs trying to read, the error can be
			// used to detect when the client closes the Connection
			_, message, err := c.Ws.ReadMessage()
			if err != nil {
				break // If we get an error escape the loop
			}

			// Bottle the message with its sender
			bottle := &Bottle{
				Sender: c,
				Message: message,
			}

			// Send to the destination for processing
			c.Dst.GetBottleChan() <- bottle
		}
		// The reader has been terminated, alert the writer if needed
		if c.valid {
			c.Close()
		}
	}()

	// Main writing loop
	for {
		select {
		case message, ok := <- c.Send:
			// Our send channel has something in it or the channel closed
			if !ok {
				// Our channel was closed, gracefully close socket Conn
				c.write(websocket.CloseMessage, []byte{})
				return
			}
			// Attempt to write the message to the websocket
			if err := c.write(websocket.TextMessage, message); err != nil {
				// If we get an error we can no longer communcate with client
				// return, no need to send CloseMessage since that would
				// just yield another error
				return
			}

		case <- ticker.C:
			// Record timestamp of ping
			c.lastping = time.Now().UnixNano()

			// Ping ticker went off. We need to ping to check for connectivity.
			if err := c.write(websocket.PingMessage, []byte{}); err != nil {
				// We got an error pinging, return and call defer
				// defer will close the socket which will kill the reader
				return
			}

		case bytes := <- c.Quit:
			// Close connection and send a final message
			c.write(websocket.TextMessage, bytes)
			c.write(websocket.CloseMessage, []byte{})
			return
		}
	}
}

var upgrader = &websocket.Upgrader{ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: func(r* http.Request) bool { return true }}

//
// Hanlder function to start a websocket connection
//
func ConnectionHandler(w http.ResponseWriter, r *http.Request, dst BottleDst) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	// Create new connection object
	c := &Conn{
		Send:     make(chan []byte, 256),
		Ws:       ws,
		Dst:      dst,
		Quit:     make(chan []byte),
		valid:    true,
		lastping: 0,
		ping:     0,
	}

	// Alert the destination that a new connection has opened
	dst.ConnectionOpened(c)

	// Start read write loop blocking
	c.read_write()
}

//
// Wrap a websocket in a rhynock handler
//
func WrapWebsocket(ws *websocket.Conn, dst BottleDst)  *Conn {
	// Create new connection object
	c := &Conn{
		Send:     make(chan []byte, 256),
		Ws:       ws,
		Dst:      dst,
		Quit:     make(chan []byte),
		valid:    true,
		lastping: 0,
		ping:     0,
	}

	// Alert the destination that a new connection has opened
	dst.ConnectionOpened(c)

	// Start read write loop blocking
	go c.read_write()

	return c
}

//
// Wraps a websocket in a rhynock.conn
//
func (c *Conn) Reroute(dst BottleDst) {
	c.Dst.ConnectionClosed(c)
	c.Dst = dst
	c.Dst.ConnectionOpened(c)
}
<|endoftext|>


<commit_before>package neptulon

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/binary"
	"encoding/pem"
	"errors"
	"fmt"
	"log"
	"net"
	"time"

	"github.com/neptulon/cmap"
)

// Conn is a client connection.
type Conn struct {
	ID                 string     // Randomly generated unique connection ID
	Data               *cmap.CMap // Thread-safe data store for storing arbitrary data for this connection session
	conn               *tls.Conn
	headerSize         int
	maxMsgSize         int
	readDeadline       time.Duration
	debug              bool
	err                error
	clientDisconnected bool // Whether the client disconnected from server before server closed connection
}

// NewConn creates a new server-side connection object.
// Default values for headerSize, maxMsgSize, and readDeadline are 4 bytes, 4294967295 bytes (4GB), and 300 seconds, respectively.
// Debug mode logs all raw TCP communication.
func NewConn(conn *tls.Conn, headerSize, maxMsgSize, readDeadline int, debug bool) (*Conn, error) {
	if headerSize == 0 {
		headerSize = 4
	}
	if maxMsgSize == 0 {
		maxMsgSize = 4294967295
	}
	if readDeadline == 0 {
		readDeadline = 300
	}

	id, err := GenID()
	if err != nil {
		return nil, err
	}

	return &Conn{
		ID:           id,
		Data:         cmap.New(),
		conn:         conn,
		headerSize:   headerSize,
		maxMsgSize:   maxMsgSize,
		readDeadline: time.Second * time.Duration(readDeadline),
		debug:        debug,
	}, nil
}

// Dial creates a new client side connection to a given network address with optional CA and/or a client certificate (PEM encoded X.509 cert/key).
// Debug mode logs all raw TCP communication.
func Dial(addr string, ca []byte, clientCert []byte, clientCertKey []byte, debug bool) (*Conn, error) {
	var cas *x509.CertPool
	var certs []tls.Certificate
	if ca != nil {
		cas = x509.NewCertPool()
		ok := cas.AppendCertsFromPEM(ca)
		if !ok {
			return nil, errors.New("failed to parse the CA certificate")
		}
	}
	if clientCert != nil {
		tlsCert, err := tls.X509KeyPair(clientCert, clientCertKey)
		if err != nil {
			return nil, fmt.Errorf("failed to parse the client certificate: %v", err)
		}

		c, _ := pem.Decode(clientCert)
		if tlsCert.Leaf, err = x509.ParseCertificate(c.Bytes); err != nil {
			return nil, fmt.Errorf("failed to parse the client certificate: %v", err)
		}

		certs = []tls.Certificate{tlsCert}
	}

	// todo: dial timeout like that of net.Conn.DialTimeout
	c, err := tls.Dial("tcp", addr, &tls.Config{RootCAs: cas, Certificates: certs})
	if err != nil {
		return nil, err
	}

	return NewConn(c, 0, 0, 0, debug)
}

// SetReadDeadline set the read deadline for the connection in seconds.
func (c *Conn) SetReadDeadline(seconds int) {
	c.readDeadline = time.Second * time.Duration(seconds)
}

// Read waits for and reads the next incoming message from the TLS connection.
func (c *Conn) Read() (msg []byte, err error) {
	if err = c.conn.SetReadDeadline(time.Now().Add(c.readDeadline)); err != nil {
		return
	}

	// read the content length header
	h := make([]byte, c.headerSize)
	var n int
	n, err = c.conn.Read(h)
	if err != nil {
		return
	}
	if n != c.headerSize {
		err = fmt.Errorf("expected to read header size %v bytes but instead read %v bytes", c.headerSize, n)
		return
	}

	// calculate the content length
	n = readHeaderBytes(h)

	// read the message content
	msg = make([]byte, n)
	total := 0
	for total < n {
		// todo: log here in case it gets stuck, or there is a dos attack, pumping up cpu usage!
		i, err := c.conn.Read(msg[total:])
		if err != nil {
			err = fmt.Errorf("errored while reading incoming message: %v", err)
			break
		}
		total += i
	}
	if total != n {
		err = fmt.Errorf("expected to read %v bytes instead read %v bytes", n, total)
	}

	if c.debug {
		log.Println("Incoming message:", string(msg))
	}

	return
}

// Write writes given message to the connection.
func (c *Conn) Write(msg []byte) error {
	l := len(msg)
	h := makeHeaderBytes(l, c.headerSize)

	// write the header
	n, err := c.conn.Write(h)
	if err != nil {
		return err
	}
	if n != c.headerSize {
		err = fmt.Errorf("expected to write %v bytes but only wrote %v bytes", l, n)
	}

	// write the body
	// todo: do we need a loop? bufio uses a loop but it might be due to buff length limitation
	n, err = c.conn.Write(msg)
	if err != nil {
		return err
	}
	if n != l {
		err = fmt.Errorf("expected to write %v bytes but only wrote %v bytes", l, n)
	}

	return nil
}

// RemoteAddr returns the remote network address.
func (c *Conn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

// ConnectionState returns basic TLS details about the connection.
func (c *Conn) ConnectionState() tls.ConnectionState {
	return c.conn.ConnectionState()
}

// Close closes a connection.
// Note: TCP/IP stack does not guarantee delivery of messages before the connection is closed.
func (c *Conn) Close() error {
	return c.conn.Close() // todo: if conn.err is nil, send a close req and wait ack then close? (or even wait for everything else to finish?)
}

func makeHeaderBytes(h, size int) []byte {
	b := make([]byte, size)
	binary.LittleEndian.PutUint32(b, uint32(h))
	return b
}

func readHeaderBytes(h []byte) int {
	return int(binary.LittleEndian.Uint32(h))
}
<commit_msg>add note<commit_after>package neptulon

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/binary"
	"encoding/pem"
	"errors"
	"fmt"
	"log"
	"net"
	"time"

	"github.com/neptulon/cmap"
)

// Conn is a client connection.
type Conn struct {
	ID                 string     // Randomly generated unique connection ID
	Data               *cmap.CMap // Thread-safe data store for storing arbitrary data for this connection session
	conn               *tls.Conn
	headerSize         int
	maxMsgSize         int
	readDeadline       time.Duration
	debug              bool
	err                error
	clientDisconnected bool // hack: Whether the client disconnected from server before server closed connection
}

// NewConn creates a new server-side connection object.
// Default values for headerSize, maxMsgSize, and readDeadline are 4 bytes, 4294967295 bytes (4GB), and 300 seconds, respectively.
// Debug mode logs all raw TCP communication.
func NewConn(conn *tls.Conn, headerSize, maxMsgSize, readDeadline int, debug bool) (*Conn, error) {
	if headerSize == 0 {
		headerSize = 4
	}
	if maxMsgSize == 0 {
		maxMsgSize = 4294967295
	}
	if readDeadline == 0 {
		readDeadline = 300
	}

	id, err := GenID()
	if err != nil {
		return nil, err
	}

	return &Conn{
		ID:           id,
		Data:         cmap.New(),
		conn:         conn,
		headerSize:   headerSize,
		maxMsgSize:   maxMsgSize,
		readDeadline: time.Second * time.Duration(readDeadline),
		debug:        debug,
	}, nil
}

// Dial creates a new client side connection to a given network address with optional CA and/or a client certificate (PEM encoded X.509 cert/key).
// Debug mode logs all raw TCP communication.
func Dial(addr string, ca []byte, clientCert []byte, clientCertKey []byte, debug bool) (*Conn, error) {
	var cas *x509.CertPool
	var certs []tls.Certificate
	if ca != nil {
		cas = x509.NewCertPool()
		ok := cas.AppendCertsFromPEM(ca)
		if !ok {
			return nil, errors.New("failed to parse the CA certificate")
		}
	}
	if clientCert != nil {
		tlsCert, err := tls.X509KeyPair(clientCert, clientCertKey)
		if err != nil {
			return nil, fmt.Errorf("failed to parse the client certificate: %v", err)
		}

		c, _ := pem.Decode(clientCert)
		if tlsCert.Leaf, err = x509.ParseCertificate(c.Bytes); err != nil {
			return nil, fmt.Errorf("failed to parse the client certificate: %v", err)
		}

		certs = []tls.Certificate{tlsCert}
	}

	// todo: dial timeout like that of net.Conn.DialTimeout
	c, err := tls.Dial("tcp", addr, &tls.Config{RootCAs: cas, Certificates: certs})
	if err != nil {
		return nil, err
	}

	return NewConn(c, 0, 0, 0, debug)
}

// SetReadDeadline set the read deadline for the connection in seconds.
func (c *Conn) SetReadDeadline(seconds int) {
	c.readDeadline = time.Second * time.Duration(seconds)
}

// Read waits for and reads the next incoming message from the TLS connection.
func (c *Conn) Read() (msg []byte, err error) {
	if err = c.conn.SetReadDeadline(time.Now().Add(c.readDeadline)); err != nil {
		return
	}

	// read the content length header
	h := make([]byte, c.headerSize)
	var n int
	n, err = c.conn.Read(h)
	if err != nil {
		return
	}
	if n != c.headerSize {
		err = fmt.Errorf("expected to read header size %v bytes but instead read %v bytes", c.headerSize, n)
		return
	}

	// calculate the content length
	n = readHeaderBytes(h)

	// read the message content
	msg = make([]byte, n)
	total := 0
	for total < n {
		// todo: log here in case it gets stuck, or there is a dos attack, pumping up cpu usage!
		i, err := c.conn.Read(msg[total:])
		if err != nil {
			err = fmt.Errorf("errored while reading incoming message: %v", err)
			break
		}
		total += i
	}
	if total != n {
		err = fmt.Errorf("expected to read %v bytes instead read %v bytes", n, total)
	}

	if c.debug {
		log.Println("Incoming message:", string(msg))
	}

	return
}

// Write writes given message to the connection.
func (c *Conn) Write(msg []byte) error {
	l := len(msg)
	h := makeHeaderBytes(l, c.headerSize)

	// write the header
	n, err := c.conn.Write(h)
	if err != nil {
		return err
	}
	if n != c.headerSize {
		err = fmt.Errorf("expected to write %v bytes but only wrote %v bytes", l, n)
	}

	// write the body
	// todo: do we need a loop? bufio uses a loop but it might be due to buff length limitation
	n, err = c.conn.Write(msg)
	if err != nil {
		return err
	}
	if n != l {
		err = fmt.Errorf("expected to write %v bytes but only wrote %v bytes", l, n)
	}

	return nil
}

// RemoteAddr returns the remote network address.
func (c *Conn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

// ConnectionState returns basic TLS details about the connection.
func (c *Conn) ConnectionState() tls.ConnectionState {
	return c.conn.ConnectionState()
}

// Close closes a connection.
// Note: TCP/IP stack does not guarantee delivery of messages before the connection is closed.
func (c *Conn) Close() error {
	return c.conn.Close() // todo: if conn.err is nil, send a close req and wait ack then close? (or even wait for everything else to finish?)
}

func makeHeaderBytes(h, size int) []byte {
	b := make([]byte, size)
	binary.LittleEndian.PutUint32(b, uint32(h))
	return b
}

func readHeaderBytes(h []byte) int {
	return int(binary.LittleEndian.Uint32(h))
}
<|endoftext|>