<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 graph_test

import (
	"fmt"
	"testing"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/internal/graph"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

func TestReverseTopsort(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type ReverseTopsortTreeTest struct {
	ctx context.Context
}

var _ SetUpInterface = &ReverseTopsortTreeTest{}

func init() { RegisterTestSuite(&ReverseTopsortTreeTest{}) }

func (t *ReverseTopsortTreeTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
}

func (t *ReverseTopsortTreeTest) run(
	root string,
	edges map[string][]string) (nodes []string, err error) {
	// Set up a successor finder.
	sf := &successorFinder{
		F: func(ctx context.Context, n string) (successors []string, err error) {
			successors = edges[n]
			return
		},
	}

	// Call through.
	c := make(chan graph.Node, 10e3)
	err = graph.ReverseTopsortTree(
		t.ctx,
		sf,
		root,
		c)

	if err != nil {
		err = fmt.Errorf("ReverseTopsortTree: %v", err)
		return
	}

	// Convert the nodes returned in the channel.
	close(c)
	for n := range c {
		nodes = append(nodes, n.(string))
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *ReverseTopsortTreeTest) SingleNode() {
	// Graph structure:
	//
	//        A
	//
	root := "A"
	edges := map[string][]string{}

	// Call
	nodes, err := t.run(root, edges)
	AssertEq(nil, err)

	ExpectThat(nodes, ElementsAre("A"))
}

func (t *ReverseTopsortTreeTest) NoBranching() {
	// Graph structure:
	//
	//        A
	//        |
	//        B
	//        |
	//        C
	//
	root := "A"
	edges := map[string][]string{
		"A": {"B"},
		"B": {"C"},
	}

	// Call
	nodes, err := t.run(root, edges)
	AssertEq(nil, err)

	AssertThat(
		sortNodes(nodes),
		ElementsAre(
			"A",
			"B",
			"C",
		))

	nodeIndex := indexNodes(nodes)
	for p, successors := range edges {
		for _, s := range successors {
			ExpectLt(nodeIndex[s], nodeIndex[p], "%q -> %q", p, s)
		}
	}
}

func (t *ReverseTopsortTreeTest) LittleBranching() {
	// Graph structure:
	//
	//        A
	//      / |
	//     B  D
	//     |  |
	//     C  E
	//
	root := "A"
	edges := map[string][]string{
		"A": {"B", "D"},
		"B": {"C"},
		"D": {"E"},
	}

	// Call
	nodes, err := t.run(root, edges)
	AssertEq(nil, err)

	AssertThat(
		sortNodes(nodes),
		ElementsAre(
			"A",
			"B",
			"C",
			"D",
			"E",
		))

	nodeIndex := indexNodes(nodes)
	for p, successors := range edges {
		for _, s := range successors {
			ExpectLt(nodeIndex[s], nodeIndex[p], "%q -> %q", p, s)
		}
	}
}

func (t *ReverseTopsortTreeTest) LotsOfBranching() {
	// Graph structure:
	//
	//        A
	//      / |
	//     B  C
	//      / | \
	//     D  E  F
	//      / |  |
	//     G  H  I
	//           | \
	//           J  K
	//
	root := "A"
	edges := map[string][]string{
		"A": []string{"B", "C"},
		"C": []string{"D", "E", "F"},
		"E": []string{"G", "H"},
		"F": []string{"I"},
		"I": []string{"J", "K"},
	}

	// Call
	nodes, err := t.run(root, edges)
	AssertEq(nil, err)

	AssertThat(
		sortNodes(nodes),
		ElementsAre(
			"A",
			"B",
			"C",
			"D",
			"E",
			"F",
			"G",
			"H",
			"I",
			"J",
			"K",
		))

	nodeIndex := indexNodes(nodes)
	for p, successors := range edges {
		for _, s := range successors {
			ExpectLt(nodeIndex[s], nodeIndex[p], "%q -> %q", p, s)
		}
	}
}

func (t *ReverseTopsortTreeTest) LargeTree() {
	const depth = 10
	edges := randomTree(depth)
	root := "root"

	// Call
	nodes, err := t.run(root, edges)
	AssertEq(nil, err)

	// All nodes should be represented.
	AssertEq(len(edges), len(nodes))
	for _, n := range nodes {
		_, ok := edges[n]
		AssertTrue(ok, "Unexpected node: %q", n)
	}

	// Edge order should be respected.
	nodeIndex := indexNodes(nodes)
	for p, successors := range edges {
		for _, s := range successors {
			ExpectLt(nodeIndex[s], nodeIndex[p], "%q -> %q", p, s)
		}
	}
}
<commit_msg>Fixed a test bug.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 graph_test

import (
	"fmt"
	"testing"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/internal/graph"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

func TestReverseTopsort(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type ReverseTopsortTreeTest struct {
	ctx context.Context
}

var _ SetUpInterface = &ReverseTopsortTreeTest{}

func init() { RegisterTestSuite(&ReverseTopsortTreeTest{}) }

func (t *ReverseTopsortTreeTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
}

func (t *ReverseTopsortTreeTest) run(
	root string,
	edges map[string][]string) (nodes []string, err error) {
	// Set up a successor finder.
	sf := &successorFinder{
		F: func(ctx context.Context, n string) (successors []string, err error) {
			successors = edges[n]
			return
		},
	}

	// Call through.
	c := make(chan graph.Node, len(edges))
	err = graph.ReverseTopsortTree(
		t.ctx,
		sf,
		root,
		c)

	if err != nil {
		err = fmt.Errorf("ReverseTopsortTree: %v", err)
		return
	}

	// Convert the nodes returned in the channel.
	close(c)
	for n := range c {
		nodes = append(nodes, n.(string))
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *ReverseTopsortTreeTest) SingleNode() {
	// Graph structure:
	//
	//        A
	//
	root := "A"
	edges := map[string][]string{
		"A": {},
	}

	// Call
	nodes, err := t.run(root, edges)
	AssertEq(nil, err)

	ExpectThat(nodes, ElementsAre("A"))
}

func (t *ReverseTopsortTreeTest) NoBranching() {
	// Graph structure:
	//
	//        A
	//        |
	//        B
	//        |
	//        C
	//
	root := "A"
	edges := map[string][]string{
		"A": {"B"},
		"B": {"C"},
		"C": {},
	}

	// Call
	nodes, err := t.run(root, edges)
	AssertEq(nil, err)

	AssertThat(
		sortNodes(nodes),
		ElementsAre(
			"A",
			"B",
			"C",
		))

	nodeIndex := indexNodes(nodes)
	for p, successors := range edges {
		for _, s := range successors {
			ExpectLt(nodeIndex[s], nodeIndex[p], "%q -> %q", p, s)
		}
	}
}

func (t *ReverseTopsortTreeTest) LittleBranching() {
	// Graph structure:
	//
	//        A
	//      / |
	//     B  D
	//     |  |
	//     C  E
	//
	root := "A"
	edges := map[string][]string{
		"A": {"B", "D"},
		"B": {"C"},
		"C": {},
		"D": {"E"},
		"E": {},
	}

	// Call
	nodes, err := t.run(root, edges)
	AssertEq(nil, err)

	AssertThat(
		sortNodes(nodes),
		ElementsAre(
			"A",
			"B",
			"C",
			"D",
			"E",
		))

	nodeIndex := indexNodes(nodes)
	for p, successors := range edges {
		for _, s := range successors {
			ExpectLt(nodeIndex[s], nodeIndex[p], "%q -> %q", p, s)
		}
	}
}

func (t *ReverseTopsortTreeTest) LotsOfBranching() {
	// Graph structure:
	//
	//        A
	//      / |
	//     B  C
	//      / | \
	//     D  E  F
	//      / |  |
	//     G  H  I
	//           | \
	//           J  K
	//
	root := "A"
	edges := map[string][]string{
		"A": {"B", "C"},
		"B": {},
		"C": {"D", "E", "F"},
		"D": {},
		"E": {"G", "H"},
		"F": {"I"},
		"G": {},
		"H": {},
		"I": {"J", "K"},
		"J": {},
		"K": {},
	}

	// Call
	nodes, err := t.run(root, edges)
	AssertEq(nil, err)

	AssertThat(
		sortNodes(nodes),
		ElementsAre(
			"A",
			"B",
			"C",
			"D",
			"E",
			"F",
			"G",
			"H",
			"I",
			"J",
			"K",
		))

	nodeIndex := indexNodes(nodes)
	for p, successors := range edges {
		for _, s := range successors {
			ExpectLt(nodeIndex[s], nodeIndex[p], "%q -> %q", p, s)
		}
	}
}

func (t *ReverseTopsortTreeTest) LargeTree() {
	const depth = 10
	edges := randomTree(depth)
	root := "root"

	// Call
	nodes, err := t.run(root, edges)
	AssertEq(nil, err)

	// All nodes should be represented.
	AssertEq(len(edges), len(nodes))
	for _, n := range nodes {
		_, ok := edges[n]
		AssertTrue(ok, "Unexpected node: %q", n)
	}

	// Edge order should be respected.
	nodeIndex := indexNodes(nodes)
	for p, successors := range edges {
		for _, s := range successors {
			ExpectLt(nodeIndex[s], nodeIndex[p], "%q -> %q", p, s)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"time"
)

type jsonRPCRequest struct {
	JsonRPC string      `json:"jsonrpc"`
	Id      string      `json:"id"`
	Method  string      `json:"method"`
	Params  interface{} `json:"params",omitempty`
}

func NewJSONRPCRequest(id string) *jsonRPCRequest {
	var req = new(jsonRPCRequest)
	req.JsonRPC = "2.0"
	req.Id = id
	if id == "" {
		req.Id = strconv.FormatInt(time.Now().UnixNano(), 10)
	}
	return req
}

func (r *jsonRPCRequest) setMethodAndParams(args []string) error {
	switch {
	case len(args) == 0:
		return fmt.Errorf("missing 'method' (first argument)")
	case len(args) == 1:
		r.Method = args[0]
		r.Params = []string{}
	case len(args) > 1:
		r.Method = args[0]

		if strings.ContainsRune(args[1], '=') {
			// if first argument conatins a '=' i assume
			// a list of key-value pairs, forming a map
			// or an 'object' in javascript terms
			r.Params = kvPairsToMap(args[1:])

		} else if len(args[1]) > 1 && args[1][0] == '{' {
			var obj, err = stringToJSObject(args[1])
			if err != nil {
				return err
			}
			r.Params = obj
		} else {
			r.Params = args[1:]
		}
	}
	return nil
}

func (r *jsonRPCRequest) toJSON() (string, error) {
	var buf = bytes.NewBuffer(nil)
	var encoder = json.NewEncoder(buf)
	if err := encoder.Encode(r); err != nil {
		return "", err
	}
	return buf.String(), nil
}

func kvPairsToMap(pairs []string) map[string]interface{} {
	var m = make(map[string]interface{})
	for _, elem := range pairs {
		var (
			parts = strings.SplitN(elem, "=", 2)
			key   = parts[0]
			val   interface{}
		)

		if len(parts) == 2 {
			// if the value can be parsed as a float, we assume
			// a json.Number
			if _, err := strconv.ParseFloat(parts[1], 64); err == nil {
				val = json.Number(parts[1])
			}
		}
		m[key] = val
	}
	return m
}

func stringToJSObject(s string) (interface{}, error) {
	var obj interface{}
	var decoder = json.NewDecoder(strings.NewReader(s))
	if err := decoder.Decode(&obj); err != nil {
		return nil, err
	}
	return obj, nil
}
<commit_msg>fix bug: assign non-numbers as well<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"time"
)

type jsonRPCRequest struct {
	JsonRPC string      `json:"jsonrpc"`
	Id      string      `json:"id"`
	Method  string      `json:"method"`
	Params  interface{} `json:"params",omitempty`
}

func NewJSONRPCRequest(id string) *jsonRPCRequest {
	var req = new(jsonRPCRequest)
	req.JsonRPC = "2.0"
	req.Id = id
	if id == "" {
		req.Id = strconv.FormatInt(time.Now().UnixNano(), 10)
	}
	return req
}

func (r *jsonRPCRequest) setMethodAndParams(args []string) error {
	switch {
	case len(args) == 0:
		return fmt.Errorf("missing 'method' (first argument)")
	case len(args) == 1:
		r.Method = args[0]
		r.Params = []string{}
	case len(args) > 1:
		r.Method = args[0]

		if strings.ContainsRune(args[1], '=') {
			// if first argument conatins a '=' i assume
			// a list of key-value pairs, forming a map
			// or an 'object' in javascript terms
			r.Params = kvPairsToMap(args[1:])

		} else if len(args[1]) > 1 && args[1][0] == '{' {
			var obj, err = stringToJSObject(args[1])
			if err != nil {
				return err
			}
			r.Params = obj
		} else {
			r.Params = args[1:]
		}
	}
	return nil
}

func (r *jsonRPCRequest) toJSON() (string, error) {
	var buf = bytes.NewBuffer(nil)
	var encoder = json.NewEncoder(buf)
	if err := encoder.Encode(r); err != nil {
		return "", err
	}
	return buf.String(), nil
}

func kvPairsToMap(pairs []string) map[string]interface{} {
	var m = make(map[string]interface{})
	for _, elem := range pairs {
		var (
			parts = strings.SplitN(elem, "=", 2)
			key   = parts[0]
			val   interface{}
		)

		if len(parts) == 2 {
			// if the value can be parsed as a float, we assume
			// a json.Number
			if _, err := strconv.ParseFloat(parts[1], 64); err == nil {
				val = json.Number(parts[1])
			} else {
				val = parts[1]
			}
		}
		m[key] = val
	}
	return m
}

func stringToJSObject(s string) (interface{}, error) {
	var obj interface{}
	var decoder = json.NewDecoder(strings.NewReader(s))
	if err := decoder.Decode(&obj); err != nil {
		return nil, err
	}
	return obj, nil
}
<|endoftext|>


<commit_before>package checklog

import (
	"bufio"
	"crypto/md5"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/Songmu/postailer"
	"github.com/jessevdk/go-flags"
	"github.com/mackerelio/checkers"
	"github.com/mackerelio/golib/pluginutil"
	"github.com/mattn/go-encoding"
	"github.com/mattn/go-zglob"
	enc "golang.org/x/text/encoding"
)

type logOpts struct {
	LogFile             string   `short:"f" long:"file" value-name:"FILE" description:"Path to log file"`
	Pattern             []string `short:"p" long:"pattern" required:"true" value-name:"PAT" description:"Pattern to search for. If specified multiple, they will be treated together with the AND operator"`
	SuppressPattern     bool     `long:"suppress-pattern" description:"Suppress pattern display"`
	Exclude             string   `short:"E" long:"exclude" value-name:"PAT" description:"Pattern to exclude from matching"`
	WarnOver            int64    `short:"w" long:"warning-over" description:"Trigger a warning if matched lines is over a number"`
	CritOver            int64    `short:"c" long:"critical-over" description:"Trigger a critical if matched lines is over a number"`
	WarnLevel           float64  `long:"warning-level" value-name:"N" description:"Warning level if pattern has a group"`
	CritLevel           float64  `long:"critical-level" value-name:"N" description:"Critical level if pattern has a group"`
	ReturnContent       bool     `short:"r" long:"return" description:"Return matched line"`
	FilePattern         string   `short:"F" long:"file-pattern" value-name:"FILE" description:"Check a pattern of files, instead of one file"`
	CaseInsensitive     bool     `short:"i" long:"icase" description:"Run a case insensitive match"`
	StateDir            string   `short:"s" long:"state-dir" value-name:"DIR" description:"Dir to keep state files under"`
	NoState             bool     `long:"no-state" description:"Don't use state file and read whole logs"`
	Encoding            string   `long:"encoding" description:"Encoding of log file"`
	Missing             string   `long:"missing" default:"UNKNOWN" value-name:"(CRITICAL|WARNING|OK|UNKNOWN)" description:"Exit status when log files missing"`
	CheckFirst          bool     `long:"check-first" description:"Check the log on the first run"`
	patternReg          []*regexp.Regexp
	excludeReg          *regexp.Regexp
	fileListFromGlob    []string
	fileListFromPattern []string
	origArgs            []string
	decoder             *enc.Decoder
}

func (opts *logOpts) prepare() error {
	if opts.LogFile == "" && opts.FilePattern == "" {
		return fmt.Errorf("No log file specified")
	}

	var err error
	var reg *regexp.Regexp
	for _, ptn := range opts.Pattern {
		if reg, err = regCompileWithCase(ptn, opts.CaseInsensitive); err != nil {
			return fmt.Errorf("pattern is invalid")
		}
		opts.patternReg = append(opts.patternReg, reg)
	}

	if len(opts.patternReg) > 1 && (opts.WarnLevel > 0 || opts.CritLevel > 0) {
		return fmt.Errorf("When multiple patterns specified, --warning-level --critical-level can not be used")
	}

	if opts.Exclude != "" {
		opts.excludeReg, err = regCompileWithCase(opts.Exclude, opts.CaseInsensitive)
		if err != nil {
			return fmt.Errorf("exclude pattern is invalid")
		}
	}

	if opts.LogFile != "" {
		opts.fileListFromGlob, err = zglob.Glob(opts.LogFile)
		// unless --missing specified, we should ignore file not found error
		if err != nil && err != os.ErrNotExist {
			return fmt.Errorf("invalid glob for --file")
		}
	}

	if opts.FilePattern != "" {
		dirStr := filepath.Dir(opts.FilePattern)
		filePat := filepath.Base(opts.FilePattern)
		reg, err := regCompileWithCase(filePat, opts.CaseInsensitive)
		if err != nil {
			return fmt.Errorf("file-pattern is invalid")
		}

		fileInfos, err := ioutil.ReadDir(dirStr)
		if err != nil {
			return fmt.Errorf("cannot read the directory:" + err.Error())
		}

		for _, fileInfo := range fileInfos {
			if fileInfo.IsDir() {
				continue
			}
			fname := fileInfo.Name()
			if reg.MatchString(fname) {
				opts.fileListFromPattern = append(opts.fileListFromPattern, dirStr+string(filepath.Separator)+fileInfo.Name())
			}
		}
	}
	if !validateMissing(opts.Missing) {
		return fmt.Errorf("missing option is invalid")
	}
	return nil
}

// Do the plugin
func Do() {
	ckr := run(os.Args[1:])
	ckr.Name = "LOG"
	ckr.Exit()
}

func regCompileWithCase(ptn string, caseInsensitive bool) (*regexp.Regexp, error) {
	if caseInsensitive {
		ptn = "(?i)" + ptn
	}
	return regexp.Compile(ptn)
}

func validateMissing(missing string) bool {
	switch missing {
	case "CRITICAL", "WARNING", "OK", "UNKNOWN", "":
		return true
	default:
		return false
	}
}

func parseArgs(args []string) (*logOpts, error) {
	origArgs := make([]string, len(args))
	copy(origArgs, args)
	opts := &logOpts{}
	_, err := flags.ParseArgs(opts, args)
	opts.origArgs = origArgs
	if opts.StateDir == "" {
		workdir := pluginutil.PluginWorkDir()
		opts.StateDir = filepath.Join(workdir, "check-log")
	}
	return opts, err
}

func run(args []string) *checkers.Checker {
	opts, err := parseArgs(args)
	if err != nil {
		os.Exit(1)
	}

	err = opts.prepare()
	if err != nil {
		return checkers.Unknown(err.Error())
	}

	warnNum := int64(0)
	critNum := int64(0)
	var missingFiles []string
	errorOverall := ""

	if opts.LogFile != "" && len(opts.fileListFromGlob) == 0 {
		missingFiles = append(missingFiles, opts.LogFile)
	}

	for _, f := range append(opts.fileListFromGlob, opts.fileListFromPattern...) {
		_, err := os.Stat(f)
		if err != nil {
			missingFiles = append(missingFiles, f)
			continue
		}
		w, c, errLines, err := opts.searchLog(f)
		if err != nil {
			return checkers.Unknown(err.Error())
		}
		warnNum += w
		critNum += c
		if opts.ReturnContent && errLines != "" {
			errorOverall += "[" + f + "]\n" + errLines
		}
	}

	var patterns []string
	for _, ptn := range opts.Pattern {
		patterns = append(patterns, fmt.Sprintf("/%s/", ptn))
	}
	var msg string
	if opts.SuppressPattern {
		msg = fmt.Sprintf("%d warnings, %d criticals.", warnNum, critNum)
	} else {
		msg = fmt.Sprintf("%d warnings, %d criticals for pattern %s.", warnNum, critNum, strings.Join(patterns, " and "))
	}
	if errorOverall != "" {
		msg += "\n" + errorOverall
	}
	checkSt := checkers.OK
	if len(missingFiles) > 0 {
		switch opts.Missing {
		case "OK":
		case "WARNING":
			checkSt = checkers.WARNING
		case "CRITICAL":
			checkSt = checkers.CRITICAL
		default:
			checkSt = checkers.UNKNOWN
		}
		msg += "\n" + fmt.Sprintf("The following %d files are missing.", len(missingFiles))
		for _, f := range missingFiles {
			msg += "\n" + f
		}
	}
	if warnNum > opts.WarnOver {
		checkSt = checkers.WARNING
	}
	if critNum > opts.CritOver {
		checkSt = checkers.CRITICAL
	}
	return checkers.NewChecker(checkSt, msg)
}

func (opts *logOpts) searchLog(logFile string) (int64, int64, string, error) {
	stateFile := getStateFile(opts.StateDir, logFile, opts.origArgs)
	skipBytes := int64(0)
	if !opts.NoState {
		s, err := getBytesToSkip(stateFile)
		if err != nil {
			return 0, 0, "", err
		}
		skipBytes = s
	}

	var r io.ReadSeeker
	if opts.NoState || opts.CheckFirst {
		f, err := os.Open(logFile)
		if err != nil {
			return 0, 0, "", err
		}
		defer f.Close()
		r = f
	} else {
		// Trace old logfile after logrotate with the inode number of the logfile.
		posfile := stateFile + ".pos.json"
		f, err := postailer.Open(logFile, posfile)
		if err != nil {
			return 0, 0, "", err
		}
		defer f.Close()
		r = f
	}

	stat, err := os.Stat(logFile)
	if err != nil {
		return 0, 0, "", err
	}

	if !opts.NoState && !opts.CheckFirst {
		if _, err = os.Stat(stateFile); os.IsNotExist(err) {
			skipBytes = stat.Size()
		}
	}

	rotated := false
	if stat.Size() < skipBytes {
		rotated = true
	} else if skipBytes > 0 {
		r.Seek(skipBytes, 0)
	}

	if opts.Encoding != "" {
		e := encoding.GetEncoding(opts.Encoding)
		if e == nil {
			return 0, 0, "", fmt.Errorf("unknown encoding:" + opts.Encoding)
		}
		opts.decoder = e.NewDecoder()
	}

	warnNum, critNum, readBytes, errLines, err := opts.searchReader(r)
	if err != nil {
		return warnNum, critNum, errLines, err
	}

	if rotated {
		skipBytes = readBytes
	} else {
		skipBytes += readBytes
	}

	if !opts.NoState {
		err = writeBytesToSkip(stateFile, skipBytes)
		if err != nil {
			log.Printf("writeByteToSkip failed: %s\n", err.Error())
		}
	}
	return warnNum, critNum, errLines, nil
}

func (opts *logOpts) searchReader(rdr io.Reader) (warnNum, critNum, readBytes int64, errLines string, err error) {
	r := bufio.NewReader(rdr)
	for {
		lineBytes, rErr := r.ReadBytes('\n')
		if rErr != nil {
			if rErr != io.EOF {
				err = rErr
			}
			break
		}
		readBytes += int64(len(lineBytes))

		if opts.decoder != nil {
			lineBytes, err = opts.decoder.Bytes(lineBytes)
			if err != nil {
				break
			}
		}
		line := strings.Trim(string(lineBytes), "\r\n")
		if matched, matches := opts.match(line); matched {
			if len(matches) > 1 && (opts.WarnLevel > 0 || opts.CritLevel > 0) {
				level, err := strconv.ParseFloat(matches[1], 64)
				if err != nil {
					warnNum++
					critNum++
					errLines += line + "\n"
				} else {
					levelOver := false
					if level > opts.WarnLevel {
						levelOver = true
						warnNum++
					}
					if level > opts.CritLevel {
						levelOver = true
						critNum++
					}
					if levelOver {
						errLines += line + "\n"
					}
				}
			} else {
				warnNum++
				critNum++
				errLines += line + "\n"
			}
		}
	}
	return
}

func (opts *logOpts) match(line string) (bool, []string) {
	var matches []string
	for _, pReg := range opts.patternReg {
		eReg := opts.excludeReg

		matches = pReg.FindStringSubmatch(line)
		if len(matches) == 0 || (eReg != nil && eReg.MatchString(line)) {
			return false, nil
		}
	}
	return true, matches
}

var stateRe = regexp.MustCompile(`^([a-zA-Z]):[/\\]`)

func getStateFile(stateDir, f string, args []string) string {
	return filepath.Join(
		stateDir,
		fmt.Sprintf(
			"%s-%x",
			stateRe.ReplaceAllString(f, `$1`+string(filepath.Separator)),
			md5.Sum([]byte(strings.Join(args, " "))),
		),
	)
}

func getBytesToSkip(f string) (int64, error) {
	_, err := os.Stat(f)
	if err != nil {
		return 0, nil
	}
	b, err := ioutil.ReadFile(f)
	if err != nil {
		return 0, err
	}
	i, err := strconv.ParseInt(strings.Trim(string(b), " \r\n"), 10, 64)
	if err != nil {
		log.Printf("failed to getBytesToSkip (ignoring): %s", err)
	}
	return i, nil
}

func writeBytesToSkip(f string, num int64) error {
	err := os.MkdirAll(filepath.Dir(f), 0755)
	if err != nil {
		return err
	}
	return writeFileAtomically(f, []byte(fmt.Sprintf("%d", num)))
}

func writeFileAtomically(f string, contents []byte) error {
	// MUST be located on same disk partition
	tmpf, err := ioutil.TempFile(filepath.Dir(f), "tmp")
	if err != nil {
		return err
	}
	// os.Remove here works successfully when tmpf.Write fails or os.Rename fails.
	// In successful case, os.Remove fails because the temporary file is already renamed.
	defer os.Remove(tmpf.Name())
	_, err = tmpf.Write(contents)
	tmpf.Close() // should be called before rename
	if err != nil {
		return err
	}
	return os.Rename(tmpf.Name(), f)
}
<commit_msg>fix bug it doesn't trace the old file if specified --check-first<commit_after>package checklog

import (
	"bufio"
	"crypto/md5"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/Songmu/postailer"
	"github.com/jessevdk/go-flags"
	"github.com/mackerelio/checkers"
	"github.com/mackerelio/golib/pluginutil"
	"github.com/mattn/go-encoding"
	"github.com/mattn/go-zglob"
	enc "golang.org/x/text/encoding"
)

type logOpts struct {
	LogFile             string   `short:"f" long:"file" value-name:"FILE" description:"Path to log file"`
	Pattern             []string `short:"p" long:"pattern" required:"true" value-name:"PAT" description:"Pattern to search for. If specified multiple, they will be treated together with the AND operator"`
	SuppressPattern     bool     `long:"suppress-pattern" description:"Suppress pattern display"`
	Exclude             string   `short:"E" long:"exclude" value-name:"PAT" description:"Pattern to exclude from matching"`
	WarnOver            int64    `short:"w" long:"warning-over" description:"Trigger a warning if matched lines is over a number"`
	CritOver            int64    `short:"c" long:"critical-over" description:"Trigger a critical if matched lines is over a number"`
	WarnLevel           float64  `long:"warning-level" value-name:"N" description:"Warning level if pattern has a group"`
	CritLevel           float64  `long:"critical-level" value-name:"N" description:"Critical level if pattern has a group"`
	ReturnContent       bool     `short:"r" long:"return" description:"Return matched line"`
	FilePattern         string   `short:"F" long:"file-pattern" value-name:"FILE" description:"Check a pattern of files, instead of one file"`
	CaseInsensitive     bool     `short:"i" long:"icase" description:"Run a case insensitive match"`
	StateDir            string   `short:"s" long:"state-dir" value-name:"DIR" description:"Dir to keep state files under"`
	NoState             bool     `long:"no-state" description:"Don't use state file and read whole logs"`
	Encoding            string   `long:"encoding" description:"Encoding of log file"`
	Missing             string   `long:"missing" default:"UNKNOWN" value-name:"(CRITICAL|WARNING|OK|UNKNOWN)" description:"Exit status when log files missing"`
	CheckFirst          bool     `long:"check-first" description:"Check the log on the first run"`
	patternReg          []*regexp.Regexp
	excludeReg          *regexp.Regexp
	fileListFromGlob    []string
	fileListFromPattern []string
	origArgs            []string
	decoder             *enc.Decoder
}

func (opts *logOpts) prepare() error {
	if opts.LogFile == "" && opts.FilePattern == "" {
		return fmt.Errorf("No log file specified")
	}

	var err error
	var reg *regexp.Regexp
	for _, ptn := range opts.Pattern {
		if reg, err = regCompileWithCase(ptn, opts.CaseInsensitive); err != nil {
			return fmt.Errorf("pattern is invalid")
		}
		opts.patternReg = append(opts.patternReg, reg)
	}

	if len(opts.patternReg) > 1 && (opts.WarnLevel > 0 || opts.CritLevel > 0) {
		return fmt.Errorf("When multiple patterns specified, --warning-level --critical-level can not be used")
	}

	if opts.Exclude != "" {
		opts.excludeReg, err = regCompileWithCase(opts.Exclude, opts.CaseInsensitive)
		if err != nil {
			return fmt.Errorf("exclude pattern is invalid")
		}
	}

	if opts.LogFile != "" {
		opts.fileListFromGlob, err = zglob.Glob(opts.LogFile)
		// unless --missing specified, we should ignore file not found error
		if err != nil && err != os.ErrNotExist {
			return fmt.Errorf("invalid glob for --file")
		}
	}

	if opts.FilePattern != "" {
		dirStr := filepath.Dir(opts.FilePattern)
		filePat := filepath.Base(opts.FilePattern)
		reg, err := regCompileWithCase(filePat, opts.CaseInsensitive)
		if err != nil {
			return fmt.Errorf("file-pattern is invalid")
		}

		fileInfos, err := ioutil.ReadDir(dirStr)
		if err != nil {
			return fmt.Errorf("cannot read the directory:" + err.Error())
		}

		for _, fileInfo := range fileInfos {
			if fileInfo.IsDir() {
				continue
			}
			fname := fileInfo.Name()
			if reg.MatchString(fname) {
				opts.fileListFromPattern = append(opts.fileListFromPattern, dirStr+string(filepath.Separator)+fileInfo.Name())
			}
		}
	}
	if !validateMissing(opts.Missing) {
		return fmt.Errorf("missing option is invalid")
	}
	return nil
}

// Do the plugin
func Do() {
	ckr := run(os.Args[1:])
	ckr.Name = "LOG"
	ckr.Exit()
}

func regCompileWithCase(ptn string, caseInsensitive bool) (*regexp.Regexp, error) {
	if caseInsensitive {
		ptn = "(?i)" + ptn
	}
	return regexp.Compile(ptn)
}

func validateMissing(missing string) bool {
	switch missing {
	case "CRITICAL", "WARNING", "OK", "UNKNOWN", "":
		return true
	default:
		return false
	}
}

func parseArgs(args []string) (*logOpts, error) {
	origArgs := make([]string, len(args))
	copy(origArgs, args)
	opts := &logOpts{}
	_, err := flags.ParseArgs(opts, args)
	opts.origArgs = origArgs
	if opts.StateDir == "" {
		workdir := pluginutil.PluginWorkDir()
		opts.StateDir = filepath.Join(workdir, "check-log")
	}
	return opts, err
}

func run(args []string) *checkers.Checker {
	opts, err := parseArgs(args)
	if err != nil {
		os.Exit(1)
	}

	err = opts.prepare()
	if err != nil {
		return checkers.Unknown(err.Error())
	}

	warnNum := int64(0)
	critNum := int64(0)
	var missingFiles []string
	errorOverall := ""

	if opts.LogFile != "" && len(opts.fileListFromGlob) == 0 {
		missingFiles = append(missingFiles, opts.LogFile)
	}

	for _, f := range append(opts.fileListFromGlob, opts.fileListFromPattern...) {
		_, err := os.Stat(f)
		if err != nil {
			missingFiles = append(missingFiles, f)
			continue
		}
		w, c, errLines, err := opts.searchLog(f)
		if err != nil {
			return checkers.Unknown(err.Error())
		}
		warnNum += w
		critNum += c
		if opts.ReturnContent && errLines != "" {
			errorOverall += "[" + f + "]\n" + errLines
		}
	}

	var patterns []string
	for _, ptn := range opts.Pattern {
		patterns = append(patterns, fmt.Sprintf("/%s/", ptn))
	}
	var msg string
	if opts.SuppressPattern {
		msg = fmt.Sprintf("%d warnings, %d criticals.", warnNum, critNum)
	} else {
		msg = fmt.Sprintf("%d warnings, %d criticals for pattern %s.", warnNum, critNum, strings.Join(patterns, " and "))
	}
	if errorOverall != "" {
		msg += "\n" + errorOverall
	}
	checkSt := checkers.OK
	if len(missingFiles) > 0 {
		switch opts.Missing {
		case "OK":
		case "WARNING":
			checkSt = checkers.WARNING
		case "CRITICAL":
			checkSt = checkers.CRITICAL
		default:
			checkSt = checkers.UNKNOWN
		}
		msg += "\n" + fmt.Sprintf("The following %d files are missing.", len(missingFiles))
		for _, f := range missingFiles {
			msg += "\n" + f
		}
	}
	if warnNum > opts.WarnOver {
		checkSt = checkers.WARNING
	}
	if critNum > opts.CritOver {
		checkSt = checkers.CRITICAL
	}
	return checkers.NewChecker(checkSt, msg)
}

func (opts *logOpts) searchLog(logFile string) (int64, int64, string, error) {
	stateFile := getStateFile(opts.StateDir, logFile, opts.origArgs)
	skipBytes := int64(0)
	if !opts.NoState {
		s, err := getBytesToSkip(stateFile)
		if err != nil {
			return 0, 0, "", err
		}
		skipBytes = s
	}

	var r io.ReadSeeker
	if opts.NoState {
		f, err := os.Open(logFile)
		if err != nil {
			return 0, 0, "", err
		}
		defer f.Close()
		r = f
	} else {
		// Trace old logfile after logrotate with the inode number of the logfile.
		posfile := stateFile + ".pos.json"
		f, err := postailer.Open(logFile, posfile)
		if err != nil {
			return 0, 0, "", err
		}
		defer f.Close()
		r = f
	}

	stat, err := os.Stat(logFile)
	if err != nil {
		return 0, 0, "", err
	}

	if !opts.NoState && !opts.CheckFirst {
		if _, err = os.Stat(stateFile); os.IsNotExist(err) {
			skipBytes = stat.Size()
		}
	}

	rotated := false
	if stat.Size() < skipBytes {
		rotated = true
	} else if skipBytes > 0 {
		r.Seek(skipBytes, 0)
	}

	if opts.Encoding != "" {
		e := encoding.GetEncoding(opts.Encoding)
		if e == nil {
			return 0, 0, "", fmt.Errorf("unknown encoding:" + opts.Encoding)
		}
		opts.decoder = e.NewDecoder()
	}

	warnNum, critNum, readBytes, errLines, err := opts.searchReader(r)
	if err != nil {
		return warnNum, critNum, errLines, err
	}

	if rotated {
		skipBytes = readBytes
	} else {
		skipBytes += readBytes
	}

	if !opts.NoState {
		err = writeBytesToSkip(stateFile, skipBytes)
		if err != nil {
			log.Printf("writeByteToSkip failed: %s\n", err.Error())
		}
	}
	return warnNum, critNum, errLines, nil
}

func (opts *logOpts) searchReader(rdr io.Reader) (warnNum, critNum, readBytes int64, errLines string, err error) {
	r := bufio.NewReader(rdr)
	for {
		lineBytes, rErr := r.ReadBytes('\n')
		if rErr != nil {
			if rErr != io.EOF {
				err = rErr
			}
			break
		}
		readBytes += int64(len(lineBytes))

		if opts.decoder != nil {
			lineBytes, err = opts.decoder.Bytes(lineBytes)
			if err != nil {
				break
			}
		}
		line := strings.Trim(string(lineBytes), "\r\n")
		if matched, matches := opts.match(line); matched {
			if len(matches) > 1 && (opts.WarnLevel > 0 || opts.CritLevel > 0) {
				level, err := strconv.ParseFloat(matches[1], 64)
				if err != nil {
					warnNum++
					critNum++
					errLines += line + "\n"
				} else {
					levelOver := false
					if level > opts.WarnLevel {
						levelOver = true
						warnNum++
					}
					if level > opts.CritLevel {
						levelOver = true
						critNum++
					}
					if levelOver {
						errLines += line + "\n"
					}
				}
			} else {
				warnNum++
				critNum++
				errLines += line + "\n"
			}
		}
	}
	return
}

func (opts *logOpts) match(line string) (bool, []string) {
	var matches []string
	for _, pReg := range opts.patternReg {
		eReg := opts.excludeReg

		matches = pReg.FindStringSubmatch(line)
		if len(matches) == 0 || (eReg != nil && eReg.MatchString(line)) {
			return false, nil
		}
	}
	return true, matches
}

var stateRe = regexp.MustCompile(`^([a-zA-Z]):[/\\]`)

func getStateFile(stateDir, f string, args []string) string {
	return filepath.Join(
		stateDir,
		fmt.Sprintf(
			"%s-%x",
			stateRe.ReplaceAllString(f, `$1`+string(filepath.Separator)),
			md5.Sum([]byte(strings.Join(args, " "))),
		),
	)
}

func getBytesToSkip(f string) (int64, error) {
	_, err := os.Stat(f)
	if err != nil {
		return 0, nil
	}
	b, err := ioutil.ReadFile(f)
	if err != nil {
		return 0, err
	}
	i, err := strconv.ParseInt(strings.Trim(string(b), " \r\n"), 10, 64)
	if err != nil {
		log.Printf("failed to getBytesToSkip (ignoring): %s", err)
	}
	return i, nil
}

func writeBytesToSkip(f string, num int64) error {
	err := os.MkdirAll(filepath.Dir(f), 0755)
	if err != nil {
		return err
	}
	return writeFileAtomically(f, []byte(fmt.Sprintf("%d", num)))
}

func writeFileAtomically(f string, contents []byte) error {
	// MUST be located on same disk partition
	tmpf, err := ioutil.TempFile(filepath.Dir(f), "tmp")
	if err != nil {
		return err
	}
	// os.Remove here works successfully when tmpf.Write fails or os.Rename fails.
	// In successful case, os.Remove fails because the temporary file is already renamed.
	defer os.Remove(tmpf.Name())
	_, err = tmpf.Write(contents)
	tmpf.Close() // should be called before rename
	if err != nil {
		return err
	}
	return os.Rename(tmpf.Name(), f)
}
<|endoftext|>


<commit_before>package hal

/*
 * Copyright 2016-2017 Netflix, 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.
 */

import (
	"sync"

	"github.com/juju/errors"
)

// directory is a simple graph of directory-style information that can be linked
// and queried on those links. Goals:
// 1. avoid coupling between plugins
// 2. make it easy to share data between plugins
// 3. make it easier to link data across various systems (e.g. Pagerduty, company directory, Slack)
type directory struct {
	initOnce sync.Once
}

const DirNodeTable = `
CREATE TABLE IF NOT EXISTS dir_node (
	pkey    VARCHAR(191) NOT NULL,
	kind    VARCHAR(191) NOT NULL,
	ts      TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
	PRIMARY KEY(pkey, kind)
)`

type DirNode struct {
	Key  string `json:"key"`
	Kind string `json:"kind"`
	Ts   int    `json:"ts"`
}

const DirNodeAttrTable = `
CREATE TABLE IF NOT EXISTS dir_node_attr (
	pkey    VARCHAR(191) NOT NULL,
	kind    VARCHAR(191) NOT NULL,
	attr    VARCHAR(191) NOT NULL,
	value   MEDIUMTEXT,
	ts      TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
	PRIMARY KEY(pkey, kind, attr),
	INDEX (pkey, kind),
	FOREIGN KEY (pkey, kind) REFERENCES dir_node(pkey, kind) ON UPDATE CASCADE
)`

type DirNodeAttr struct {
	Key   string `json:"key"`
	Kind  string `json:"kind"`
	Attr  string `json:"attr"`
	Value string `json:"value"`
	Ts    int    `json:"ts"`
}

const DirEdgeTable = `
CREATE TABLE IF NOT EXISTS dir_edge (
	keyA    VARCHAR(191) NOT NULL,
	kindA   VARCHAR(191) NOT NULL,
	keyB    VARCHAR(191) NOT NULL,
	kindB   VARCHAR(191) NOT NULL,
	ts      TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
	PRIMARY KEY(keyA, kindA, keyB, kindB),
	INDEX (keyA, kindA),
	INDEX (keyB, kindB),
	FOREIGN KEY (keyA, kindA) REFERENCES dir_node(pkey, kind) ON UPDATE CASCADE,
	FOREIGN KEY (keyB, kindB) REFERENCES dir_node(pkey, kind) ON UPDATE CASCADE
)`

type DirEdge struct {
	KeyA  string `json:"key_a"`
	KindA string `json:"kind_a"`
	KeyB  string `json:"key_b"`
	KindB string `json:"kind_b"`
	Ts    int    `json:"ts"`
}

var dirSingleton directory

func Directory() *directory {
	dirSingleton.initOnce.Do(func() {
		SqlInit(DirNodeTable)
		SqlInit(DirNodeAttrTable)
		SqlInit(DirEdgeTable)
	})

	return &dirSingleton
}

func (dir *directory) exec(sql string, params ...interface{}) error {
	db := SqlDB()

	_, err := db.Exec(sql, params...)
	if err != nil {
		return errors.Annotatef(err, "SQL: %q, Values: %+q", sql, params)
	}

	return nil
}

func (dir *directory) query(sql string, params ...interface{}) ([][]string, error) {
	db := SqlDB()
	out := make([][]string, 0)

	rows, err := db.Query(sql, params...)
	if err != nil {
		return out, errors.Annotatef(err, "SQL: %q, Values: %+q", sql, params)
	}

	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		return out, errors.Annotate(err, "rows.Columns()")
	}

	for rows.Next() {
		irow := make([]interface{}, len(cols))
		row := make([]string, len(irow))

		for i, _ := range irow {
			irow[i] = &row[i]
		}

		err := rows.Scan(irow...)
		if err != nil {
			return out, errors.Annotate(err, "rows.Scan()")
		}

		out = append(out, row)
	}

	return out, nil
}

// Put adds/updates a node for the given key/attr and creates edges for
// the attributes where possible.
func (dir *directory) Put(key, kind string, attrs map[string]string, edgeAttrs []string) error {
	err := dir.PutNode(key, kind)
	if err != nil {
		return err
	}

	if attrs == nil {
		return errors.NotValidf("attrs cannot be nil")
	}

	for attr, value := range attrs {
		err := dir.PutNodeAttr(key, kind, attr, value)
		if err != nil {
			return err
		}
	}

	// experimental: use the provided list of keys to try to create edges based on attributes
	for _, ea := range edgeAttrs {
		if value, exists := attrs[ea]; exists {
			neighbors, err := dir.GetAttrNodes(ea, value)
			if err != nil {
				return errors.Annotate(err, "GetAttrNodes failed")
			}
			for _, neighbor := range neighbors {
				dir.PutEdge(key, kind, neighbor[0], neighbor[1])
			}
		}
	}

	return nil
}

func (dir *directory) PutNode(key, kind string) error {
	sql := `INSERT INTO dir_node (pkey, kind) VALUES (?, ?) ON DUPLICATE KEY UPDATE ts=NOW()`
	return dir.exec(sql, key, kind)
}

func (dir *directory) HasNode(key, kind string) (bool, error) {
	sql := `SELECT pkey, kind FROM dir_node WHERE pkey=? AND kind=?`

	data, err := dir.query(sql, &key, &kind)
	if err != nil {
		return false, err
	}

	if len(data) > 0 {
		return true, nil
	}

	return false, nil
}

func (dir *directory) DelNode(key, kind string) error {
	sql := `DELETE FROM dir_node WHERE key=? AND kind=?`
	return dir.exec(sql, key, kind)
}

func (dir *directory) PutNodeAttr(key, kind, attr, value string) error {
	sql := `INSERT INTO dir_node_attr (pkey, kind, attr, value) VALUES (?, ?, ?, ?) ON DUPLICATE KEY UPDATE value=?`
	return dir.exec(sql, key, kind, attr, value, value)
}

// GetAttrNodes takes an attribute and value and returns a list of nodes
// that have (exactly) matching attributes.
func (dir *directory) GetAttrNodes(attr, value string) ([][2]string, error) {
	out := make([][2]string, 0)
	sql := `SELECT pkey, kind FROM dir_node_attr WHERE attr=? AND value=? GROUP BY pkey, kind`

	data, err := dir.query(sql, &attr, &value)
	if err != nil {
		return out, err
	}

	for _, item := range data {
		// item is []string, return must be [2]string
		out = append(out, [2]string{item[0], item[1]})
	}

	return out, nil
}

func (dir *directory) HasEdge(keyA, kindA, keyB, kindB string) (bool, error) {
	sql := `SELECT "y" FROM dir_edge WHERE keyA=? AND kindA=? AND keyB=? AND kindB=?`

	data, err := dir.query(sql, &keyA, &kindA, &keyB, &kindB)
	if err != nil {
		return false, err
	}

	if len(data) > 0 {
		return true, nil
	}

	return false, nil
}

func (dir *directory) PutEdge(keyA, kindA, keyB, kindB string) error {
	sql := `INSERT INTO dir_edge (keyA, kindA, keyB, kindB) VALUES (?, ?, ?, ?) ON DUPLICATE KEY UPDATE ts=NOW()`
	return dir.exec(sql, keyA, kindA, keyB, kindB)
}

func (dir *directory) DelEdge(keyA, kindA, keyB, kindB string) error {
	sql := `DELETE FROM dir_edge WHERE keyA=? AND kindA=? AND keyB=? AND kindB=?`
	return dir.exec(sql, keyA, kindA, keyB, kindB)
}

func (dir *directory) GetNeighbors(key, kind string) ([][2]string, error) {
	out := make([][2]string, 0)

	sql := `SELECT keyA, kindA, keyB, kindB FROM dir_edge WHERE (keyA=? AND kindA=?) OR (keyB=? AND kindB=?)`

	edges, err := dir.query(sql, &key, &kind, &key, &kind)
	if err != nil {
		return out, err
	}

	for _, e := range edges {
		if e[0] == key && e[1] == kind {
			out = append(out, [2]string{e[2], e[3]})
		} else {
			out = append(out, [2]string{e[0], e[1]})
		}
	}

	return out, nil
}

func (dir *directory) GetEdges() ([]DirEdge, error) {
	sql := `SELECT keyA, kindA, keyB, kindB, ts FROM dir_edge`
	db := SqlDB()
	out := make([]DirEdge, 0)

	rows, err := db.Query(sql)
	if err != nil {
		return out, errors.Annotatef(err, "SQL: %q", sql)
	}

	defer rows.Close()

	for rows.Next() {
		row := DirEdge{}
		err := rows.Scan(&row.KeyA, &row.KindA, &row.KeyB, &row.KindB, &row.Ts)
		if err != nil {
			return out, errors.Annotate(err, "rows.Scan()")
		}

		out = append(out, row)
	}

	return out, nil
}

func (dir *directory) GetNodes() ([]DirNode, error) {
	sql := `SELECT pkey, kind ts FROM dir_node`
	db := SqlDB()
	out := make([]DirNode, 0)

	rows, err := db.Query(sql)
	if err != nil {
		return out, errors.Annotatef(err, "SQL: %q", sql)
	}

	defer rows.Close()

	for rows.Next() {
		row := DirNode{}
		err := rows.Scan(&row.Key, &row.Kind, &row.Ts)
		if err != nil {
			return out, errors.Annotate(err, "rows.Scan()")
		}

		out = append(out, row)
	}

	return out, nil
}

func (dir *directory) GetNodeAttrs() ([]DirNodeAttr, error) {
	sql := `SELECT pkey, kind, attr, value, ts FROM dir_node_attr`
	db := SqlDB()
	out := make([]DirNodeAttr, 0)

	rows, err := db.Query(sql)
	if err != nil {
		return out, errors.Annotatef(err, "SQL: %q", sql)
	}

	defer rows.Close()

	for rows.Next() {
		row := DirNodeAttr{}
		err := rows.Scan(&row.Key, &row.Kind, &row.Attr, &row.Value, &row.Ts)
		if err != nil {
			return out, errors.Annotate(err, "rows.Scan()")
		}

		out = append(out, row)
	}

	return out, nil
}

/* test data

INSERT INTO dir_node (kind, pkey) VALUES ("AD", "angua@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("AD", "carrot@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("AD", "aching@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("AD", "granny@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("AD", "vetinari@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("AD", "vimes@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("AD", "nobbs@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Angua");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Carrot");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Tiffany");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Mistress Weatherwax");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Patrician");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Sam");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Nobby");

INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "angua@dwmail.com", "slack", "Angua");
INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "carrot@dwmail.com", "slack", "Carrot");
INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "aching@dwmail.com", "slack", "Tiffany");
INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "granny@dwmail.com", "slack", "Mistress Weatherwax");
INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "vetinari@dwmail.com", "slack", "Patrician");
INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "vimes@dwmail.com", "slack", "Sam");
INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "nobbs@dwmail.com", "slack", "Nobby");

INSERT INTO dir_node_attr (kind, pkey, attr, value) VALUES ("AD", "angua@dwmail.com", "email", "angua@dwmail.com");
INSERT INTO dir_node_attr (kind, pkey, attr, value) VALUES ("AD", "angua@dwmail.com", "sms", "5551234567"
INSERT INTO dir_node_attr (kind, pkey, attr, value) VALUES ("AD", "carrot@dwmail.com", "sms", "5555555555");
*/
<commit_msg>use UNIX_TIMESTAMP(ts) in select queries<commit_after>package hal

/*
 * Copyright 2016-2017 Netflix, 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.
 */

import (
	"sync"

	"github.com/juju/errors"
)

// directory is a simple graph of directory-style information that can be linked
// and queried on those links. Goals:
// 1. avoid coupling between plugins
// 2. make it easy to share data between plugins
// 3. make it easier to link data across various systems (e.g. Pagerduty, company directory, Slack)
type directory struct {
	initOnce sync.Once
}

const DirNodeTable = `
CREATE TABLE IF NOT EXISTS dir_node (
	pkey    VARCHAR(191) NOT NULL,
	kind    VARCHAR(191) NOT NULL,
	ts      TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
	PRIMARY KEY(pkey, kind)
)`

type DirNode struct {
	Key  string `json:"key"`
	Kind string `json:"kind"`
	Ts   int    `json:"ts"`
}

const DirNodeAttrTable = `
CREATE TABLE IF NOT EXISTS dir_node_attr (
	pkey    VARCHAR(191) NOT NULL,
	kind    VARCHAR(191) NOT NULL,
	attr    VARCHAR(191) NOT NULL,
	value   MEDIUMTEXT,
	ts      TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
	PRIMARY KEY(pkey, kind, attr),
	INDEX (pkey, kind),
	FOREIGN KEY (pkey, kind) REFERENCES dir_node(pkey, kind) ON UPDATE CASCADE
)`

type DirNodeAttr struct {
	Key   string `json:"key"`
	Kind  string `json:"kind"`
	Attr  string `json:"attr"`
	Value string `json:"value"`
	Ts    int    `json:"ts"`
}

const DirEdgeTable = `
CREATE TABLE IF NOT EXISTS dir_edge (
	keyA    VARCHAR(191) NOT NULL,
	kindA   VARCHAR(191) NOT NULL,
	keyB    VARCHAR(191) NOT NULL,
	kindB   VARCHAR(191) NOT NULL,
	ts      TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
	PRIMARY KEY(keyA, kindA, keyB, kindB),
	INDEX (keyA, kindA),
	INDEX (keyB, kindB),
	FOREIGN KEY (keyA, kindA) REFERENCES dir_node(pkey, kind) ON UPDATE CASCADE,
	FOREIGN KEY (keyB, kindB) REFERENCES dir_node(pkey, kind) ON UPDATE CASCADE
)`

type DirEdge struct {
	KeyA  string `json:"key_a"`
	KindA string `json:"kind_a"`
	KeyB  string `json:"key_b"`
	KindB string `json:"kind_b"`
	Ts    int    `json:"ts"`
}

var dirSingleton directory

func Directory() *directory {
	dirSingleton.initOnce.Do(func() {
		SqlInit(DirNodeTable)
		SqlInit(DirNodeAttrTable)
		SqlInit(DirEdgeTable)
	})

	return &dirSingleton
}

func (dir *directory) exec(sql string, params ...interface{}) error {
	db := SqlDB()

	_, err := db.Exec(sql, params...)
	if err != nil {
		return errors.Annotatef(err, "SQL: %q, Values: %+q", sql, params)
	}

	return nil
}

func (dir *directory) query(sql string, params ...interface{}) ([][]string, error) {
	db := SqlDB()
	out := make([][]string, 0)

	rows, err := db.Query(sql, params...)
	if err != nil {
		return out, errors.Annotatef(err, "SQL: %q, Values: %+q", sql, params)
	}

	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		return out, errors.Annotate(err, "rows.Columns()")
	}

	for rows.Next() {
		irow := make([]interface{}, len(cols))
		row := make([]string, len(irow))

		for i, _ := range irow {
			irow[i] = &row[i]
		}

		err := rows.Scan(irow...)
		if err != nil {
			return out, errors.Annotate(err, "rows.Scan()")
		}

		out = append(out, row)
	}

	return out, nil
}

// Put adds/updates a node for the given key/attr and creates edges for
// the attributes where possible.
func (dir *directory) Put(key, kind string, attrs map[string]string, edgeAttrs []string) error {
	err := dir.PutNode(key, kind)
	if err != nil {
		return err
	}

	if attrs == nil {
		return errors.NotValidf("attrs cannot be nil")
	}

	for attr, value := range attrs {
		err := dir.PutNodeAttr(key, kind, attr, value)
		if err != nil {
			return err
		}
	}

	// experimental: use the provided list of keys to try to create edges based on attributes
	for _, ea := range edgeAttrs {
		if value, exists := attrs[ea]; exists {
			neighbors, err := dir.GetAttrNodes(ea, value)
			if err != nil {
				return errors.Annotate(err, "GetAttrNodes failed")
			}
			for _, neighbor := range neighbors {
				dir.PutEdge(key, kind, neighbor[0], neighbor[1])
			}
		}
	}

	return nil
}

func (dir *directory) PutNode(key, kind string) error {
	sql := `INSERT INTO dir_node (pkey, kind) VALUES (?, ?) ON DUPLICATE KEY UPDATE ts=NOW()`
	return dir.exec(sql, key, kind)
}

func (dir *directory) HasNode(key, kind string) (bool, error) {
	sql := `SELECT pkey, kind FROM dir_node WHERE pkey=? AND kind=?`

	data, err := dir.query(sql, &key, &kind)
	if err != nil {
		return false, err
	}

	if len(data) > 0 {
		return true, nil
	}

	return false, nil
}

func (dir *directory) DelNode(key, kind string) error {
	sql := `DELETE FROM dir_node WHERE key=? AND kind=?`
	return dir.exec(sql, key, kind)
}

func (dir *directory) PutNodeAttr(key, kind, attr, value string) error {
	sql := `INSERT INTO dir_node_attr (pkey, kind, attr, value) VALUES (?, ?, ?, ?) ON DUPLICATE KEY UPDATE value=?`
	return dir.exec(sql, key, kind, attr, value, value)
}

// GetAttrNodes takes an attribute and value and returns a list of nodes
// that have (exactly) matching attributes.
func (dir *directory) GetAttrNodes(attr, value string) ([][2]string, error) {
	out := make([][2]string, 0)
	sql := `SELECT pkey, kind FROM dir_node_attr WHERE attr=? AND value=? GROUP BY pkey, kind`

	data, err := dir.query(sql, &attr, &value)
	if err != nil {
		return out, err
	}

	for _, item := range data {
		// item is []string, return must be [2]string
		out = append(out, [2]string{item[0], item[1]})
	}

	return out, nil
}

func (dir *directory) HasEdge(keyA, kindA, keyB, kindB string) (bool, error) {
	sql := `SELECT "y" FROM dir_edge WHERE keyA=? AND kindA=? AND keyB=? AND kindB=?`

	data, err := dir.query(sql, &keyA, &kindA, &keyB, &kindB)
	if err != nil {
		return false, err
	}

	if len(data) > 0 {
		return true, nil
	}

	return false, nil
}

func (dir *directory) PutEdge(keyA, kindA, keyB, kindB string) error {
	sql := `INSERT INTO dir_edge (keyA, kindA, keyB, kindB) VALUES (?, ?, ?, ?) ON DUPLICATE KEY UPDATE ts=NOW()`
	return dir.exec(sql, keyA, kindA, keyB, kindB)
}

func (dir *directory) DelEdge(keyA, kindA, keyB, kindB string) error {
	sql := `DELETE FROM dir_edge WHERE keyA=? AND kindA=? AND keyB=? AND kindB=?`
	return dir.exec(sql, keyA, kindA, keyB, kindB)
}

func (dir *directory) GetNeighbors(key, kind string) ([][2]string, error) {
	out := make([][2]string, 0)

	sql := `SELECT keyA, kindA, keyB, kindB FROM dir_edge WHERE (keyA=? AND kindA=?) OR (keyB=? AND kindB=?)`

	edges, err := dir.query(sql, &key, &kind, &key, &kind)
	if err != nil {
		return out, err
	}

	for _, e := range edges {
		if e[0] == key && e[1] == kind {
			out = append(out, [2]string{e[2], e[3]})
		} else {
			out = append(out, [2]string{e[0], e[1]})
		}
	}

	return out, nil
}

func (dir *directory) GetEdges() ([]DirEdge, error) {
	sql := `SELECT keyA, kindA, keyB, kindB, UNIX_TIMESTAMP(ts) FROM dir_edge`
	db := SqlDB()
	out := make([]DirEdge, 0)

	rows, err := db.Query(sql)
	if err != nil {
		return out, errors.Annotatef(err, "SQL: %q", sql)
	}

	defer rows.Close()

	for rows.Next() {
		row := DirEdge{}
		err := rows.Scan(&row.KeyA, &row.KindA, &row.KeyB, &row.KindB, &row.Ts)
		if err != nil {
			return out, errors.Annotate(err, "rows.Scan()")
		}

		out = append(out, row)
	}

	return out, nil
}

func (dir *directory) GetNodes() ([]DirNode, error) {
	sql := `SELECT pkey, kind UNIX_TIMESTAMP(ts) FROM dir_node`
	db := SqlDB()
	out := make([]DirNode, 0)

	rows, err := db.Query(sql)
	if err != nil {
		return out, errors.Annotatef(err, "SQL: %q", sql)
	}

	defer rows.Close()

	for rows.Next() {
		row := DirNode{}
		err := rows.Scan(&row.Key, &row.Kind, &row.Ts)
		if err != nil {
			return out, errors.Annotate(err, "rows.Scan()")
		}

		out = append(out, row)
	}

	return out, nil
}

func (dir *directory) GetNodeAttrs() ([]DirNodeAttr, error) {
	sql := `SELECT pkey, kind, attr, value, UNIX_TIMESTAMP(ts) FROM dir_node_attr`
	db := SqlDB()
	out := make([]DirNodeAttr, 0)

	rows, err := db.Query(sql)
	if err != nil {
		return out, errors.Annotatef(err, "SQL: %q", sql)
	}

	defer rows.Close()

	for rows.Next() {
		row := DirNodeAttr{}
		err := rows.Scan(&row.Key, &row.Kind, &row.Attr, &row.Value, &row.Ts)
		if err != nil {
			return out, errors.Annotate(err, "rows.Scan()")
		}

		out = append(out, row)
	}

	return out, nil
}

/* test data

INSERT INTO dir_node (kind, pkey) VALUES ("AD", "angua@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("AD", "carrot@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("AD", "aching@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("AD", "granny@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("AD", "vetinari@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("AD", "vimes@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("AD", "nobbs@dwmail.com");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Angua");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Carrot");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Tiffany");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Mistress Weatherwax");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Patrician");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Sam");
INSERT INTO dir_node (kind, pkey) VALUES ("slack", "Nobby");

INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "angua@dwmail.com", "slack", "Angua");
INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "carrot@dwmail.com", "slack", "Carrot");
INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "aching@dwmail.com", "slack", "Tiffany");
INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "granny@dwmail.com", "slack", "Mistress Weatherwax");
INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "vetinari@dwmail.com", "slack", "Patrician");
INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "vimes@dwmail.com", "slack", "Sam");
INSERT INTO dir_edge (kindA, keyA, kindB, keyB) VALUES ("AD", "nobbs@dwmail.com", "slack", "Nobby");

INSERT INTO dir_node_attr (kind, pkey, attr, value) VALUES ("AD", "angua@dwmail.com", "email", "angua@dwmail.com");
INSERT INTO dir_node_attr (kind, pkey, attr, value) VALUES ("AD", "angua@dwmail.com", "sms", "5551234567"
INSERT INTO dir_node_attr (kind, pkey, attr, value) VALUES ("AD", "carrot@dwmail.com", "sms", "5555555555");
*/
<|endoftext|>


<commit_before>package irc

import (
	"errors"
	"regexp"
	"time"
)

var (
	// debugging flags
	DEBUG_NET     = false
	DEBUG_CLIENT  = false
	DEBUG_CHANNEL = false
	DEBUG_SERVER  = false

	// errors
	ErrAlreadyDestroyed = errors.New("already destroyed")

	// regexps
	ChannelNameExpr = regexp.MustCompile(`^[&!#+][\pL\pN]{1,63}$`)
	NicknameExpr    = regexp.MustCompile(
		"^[\\pL\\[\\]{}^`][\\pL\\pN\\[\\]{}^`]{1,31}$")
)

const (
	SEM_VER       = "ergonomadic-1.2.11"
	CRLF          = "\r\n"
	MAX_REPLY_LEN = 512 - len(CRLF)

	LOGIN_TIMEOUT = time.Minute / 2 // how long the client has to login
	IDLE_TIMEOUT  = time.Minute     // how long before a client is considered idle
	QUIT_TIMEOUT  = time.Minute     // how long after idle before a client is kicked

	// string codes
	AWAY    StringCode = "AWAY"
	CAP     StringCode = "CAP"
	DEBUG   StringCode = "DEBUG"
	ERROR   StringCode = "ERROR"
	INVITE  StringCode = "INVITE"
	ISON    StringCode = "ISON"
	JOIN    StringCode = "JOIN"
	KICK    StringCode = "KICK"
	KILL    StringCode = "KILL"
	LIST    StringCode = "LIST"
	MODE    StringCode = "MODE"
	MOTD    StringCode = "MOTD"
	NAMES   StringCode = "NAMES"
	NICK    StringCode = "NICK"
	NOTICE  StringCode = "NOTICE"
	OPER    StringCode = "OPER"
	PART    StringCode = "PART"
	PASS    StringCode = "PASS"
	PING    StringCode = "PING"
	PONG    StringCode = "PONG"
	PRIVMSG StringCode = "PRIVMSG"
	PROXY   StringCode = "PROXY"
	QUIT    StringCode = "QUIT"
	TIME    StringCode = "TIME"
	TOPIC   StringCode = "TOPIC"
	USER    StringCode = "USER"
	VERSION StringCode = "VERSION"
	WHO     StringCode = "WHO"
	WHOIS   StringCode = "WHOIS"

	// numeric codes
	RPL_WELCOME           NumericCode = 1
	RPL_YOURHOST          NumericCode = 2
	RPL_CREATED           NumericCode = 3
	RPL_MYINFO            NumericCode = 4
	RPL_BOUNCE            NumericCode = 5
	RPL_TRACELINK         NumericCode = 200
	RPL_TRACECONNECTING   NumericCode = 201
	RPL_TRACEHANDSHAKE    NumericCode = 202
	RPL_TRACEUNKNOWN      NumericCode = 203
	RPL_TRACEOPERATOR     NumericCode = 204
	RPL_TRACEUSER         NumericCode = 205
	RPL_TRACESERVER       NumericCode = 206
	RPL_TRACESERVICE      NumericCode = 207
	RPL_TRACENEWTYPE      NumericCode = 208
	RPL_TRACECLASS        NumericCode = 209
	RPL_TRACERECONNECT    NumericCode = 210
	RPL_STATSLINKINFO     NumericCode = 211
	RPL_STATSCOMMANDS     NumericCode = 212
	RPL_ENDOFSTATS        NumericCode = 219
	RPL_UMODEIS           NumericCode = 221
	RPL_SERVLIST          NumericCode = 234
	RPL_SERVLISTEND       NumericCode = 235
	RPL_STATSUPTIME       NumericCode = 242
	RPL_STATSOLINE        NumericCode = 243
	RPL_LUSERCLIENT       NumericCode = 251
	RPL_LUSEROP           NumericCode = 252
	RPL_LUSERUNKNOWN      NumericCode = 253
	RPL_LUSERCHANNELS     NumericCode = 254
	RPL_LUSERME           NumericCode = 255
	RPL_ADMINME           NumericCode = 256
	RPL_ADMINLOC1         NumericCode = 257
	RPL_ADMINLOC2         NumericCode = 258
	RPL_ADMINEMAIL        NumericCode = 259
	RPL_TRACELOG          NumericCode = 261
	RPL_TRACEEND          NumericCode = 262
	RPL_TRYAGAIN          NumericCode = 263
	RPL_AWAY              NumericCode = 301
	RPL_USERHOST          NumericCode = 302
	RPL_ISON              NumericCode = 303
	RPL_UNAWAY            NumericCode = 305
	RPL_NOWAWAY           NumericCode = 306
	RPL_WHOISUSER         NumericCode = 311
	RPL_WHOISSERVER       NumericCode = 312
	RPL_WHOISOPERATOR     NumericCode = 313
	RPL_WHOWASUSER        NumericCode = 314
	RPL_ENDOFWHO          NumericCode = 315
	RPL_WHOISIDLE         NumericCode = 317
	RPL_ENDOFWHOIS        NumericCode = 318
	RPL_WHOISCHANNELS     NumericCode = 319
	RPL_LIST              NumericCode = 322
	RPL_LISTEND           NumericCode = 323
	RPL_CHANNELMODEIS     NumericCode = 324
	RPL_UNIQOPIS          NumericCode = 325
	RPL_NOTOPIC           NumericCode = 331
	RPL_TOPIC             NumericCode = 332
	RPL_INVITING          NumericCode = 341
	RPL_SUMMONING         NumericCode = 342
	RPL_INVITELIST        NumericCode = 346
	RPL_ENDOFINVITELIST   NumericCode = 347
	RPL_EXCEPTLIST        NumericCode = 348
	RPL_ENDOFEXCEPTLIST   NumericCode = 349
	RPL_VERSION           NumericCode = 351
	RPL_WHOREPLY          NumericCode = 352
	RPL_NAMREPLY          NumericCode = 353
	RPL_LINKS             NumericCode = 364
	RPL_ENDOFLINKS        NumericCode = 365
	RPL_ENDOFNAMES        NumericCode = 366
	RPL_BANLIST           NumericCode = 367
	RPL_ENDOFBANLIST      NumericCode = 368
	RPL_ENDOFWHOWAS       NumericCode = 369
	RPL_INFO              NumericCode = 371
	RPL_MOTD              NumericCode = 372
	RPL_ENDOFINFO         NumericCode = 374
	RPL_MOTDSTART         NumericCode = 375
	RPL_ENDOFMOTD         NumericCode = 376
	RPL_YOUREOPER         NumericCode = 381
	RPL_REHASHING         NumericCode = 382
	RPL_YOURESERVICE      NumericCode = 383
	RPL_TIME              NumericCode = 391
	RPL_USERSSTART        NumericCode = 392
	RPL_USERS             NumericCode = 393
	RPL_ENDOFUSERS        NumericCode = 394
	RPL_NOUSERS           NumericCode = 395
	ERR_NOSUCHNICK        NumericCode = 401
	ERR_NOSUCHSERVER      NumericCode = 402
	ERR_NOSUCHCHANNEL     NumericCode = 403
	ERR_CANNOTSENDTOCHAN  NumericCode = 404
	ERR_TOOMANYCHANNELS   NumericCode = 405
	ERR_WASNOSUCHNICK     NumericCode = 406
	ERR_TOOMANYTARGETS    NumericCode = 407
	ERR_NOSUCHSERVICE     NumericCode = 408
	ERR_NOORIGIN          NumericCode = 409
	ERR_NORECIPIENT       NumericCode = 411
	ERR_NOTEXTTOSEND      NumericCode = 412
	ERR_NOTOPLEVEL        NumericCode = 413
	ERR_WILDTOPLEVEL      NumericCode = 414
	ERR_BADMASK           NumericCode = 415
	ERR_UNKNOWNCOMMAND    NumericCode = 421
	ERR_NOMOTD            NumericCode = 422
	ERR_NOADMININFO       NumericCode = 423
	ERR_FILEERROR         NumericCode = 424
	ERR_NONICKNAMEGIVEN   NumericCode = 431
	ERR_ERRONEUSNICKNAME  NumericCode = 432
	ERR_NICKNAMEINUSE     NumericCode = 433
	ERR_NICKCOLLISION     NumericCode = 436
	ERR_UNAVAILRESOURCE   NumericCode = 437
	ERR_USERNOTINCHANNEL  NumericCode = 441
	ERR_NOTONCHANNEL      NumericCode = 442
	ERR_USERONCHANNEL     NumericCode = 443
	ERR_NOLOGIN           NumericCode = 444
	ERR_SUMMONDISABLED    NumericCode = 445
	ERR_USERSDISABLED     NumericCode = 446
	ERR_NOTREGISTERED     NumericCode = 451
	ERR_NEEDMOREPARAMS    NumericCode = 461
	ERR_ALREADYREGISTRED  NumericCode = 462
	ERR_NOPERMFORHOST     NumericCode = 463
	ERR_PASSWDMISMATCH    NumericCode = 464
	ERR_YOUREBANNEDCREEP  NumericCode = 465
	ERR_YOUWILLBEBANNED   NumericCode = 466
	ERR_KEYSET            NumericCode = 467
	ERR_CHANNELISFULL     NumericCode = 471
	ERR_UNKNOWNMODE       NumericCode = 472
	ERR_INVITEONLYCHAN    NumericCode = 473
	ERR_BANNEDFROMCHAN    NumericCode = 474
	ERR_BADCHANNELKEY     NumericCode = 475
	ERR_BADCHANMASK       NumericCode = 476
	ERR_NOCHANMODES       NumericCode = 477
	ERR_BANLISTFULL       NumericCode = 478
	ERR_NOPRIVILEGES      NumericCode = 481
	ERR_CHANOPRIVSNEEDED  NumericCode = 482
	ERR_CANTKILLSERVER    NumericCode = 483
	ERR_RESTRICTED        NumericCode = 484
	ERR_UNIQOPPRIVSNEEDED NumericCode = 485
	ERR_NOOPERHOST        NumericCode = 491
	ERR_UMODEUNKNOWNFLAG  NumericCode = 501
	ERR_USERSDONTMATCH    NumericCode = 502

	Add    ModeOp = '+'
	List   ModeOp = '='
	Remove ModeOp = '-'

	Away          UserMode = 'a'
	Invisible     UserMode = 'i'
	LocalOperator UserMode = 'O'
	Operator      UserMode = 'o'
	Restricted    UserMode = 'r'
	ServerNotice  UserMode = 's' // deprecated
	WallOps       UserMode = 'w'

	Anonymous       ChannelMode = 'a' // flag
	BanMask         ChannelMode = 'b' // arg
	ChannelCreator  ChannelMode = 'O' // flag
	ChannelOperator ChannelMode = 'o' // arg
	ExceptMask      ChannelMode = 'e' // arg
	InviteMask      ChannelMode = 'I' // arg
	InviteOnly      ChannelMode = 'i' // flag
	Key             ChannelMode = 'k' // flag arg
	Moderated       ChannelMode = 'm' // flag
	NoOutside       ChannelMode = 'n' // flag
	OpOnlyTopic     ChannelMode = 't' // flag
	Persistent      ChannelMode = 'P' // flag
	Private         ChannelMode = 'p' // flag
	Quiet           ChannelMode = 'q' // flag
	ReOp            ChannelMode = 'r' // flag
	Secret          ChannelMode = 's' // flag, deprecated
	UserLimit       ChannelMode = 'l' // flag arg
	Voice           ChannelMode = 'v' // arg
)

const (
	Authorization Phase = iota
	Registration  Phase = iota
	Normal        Phase = iota
)
<commit_msg>allow underscores in nicks<commit_after>package irc

import (
	"errors"
	"regexp"
	"time"
)

var (
	// debugging flags
	DEBUG_NET     = false
	DEBUG_CLIENT  = false
	DEBUG_CHANNEL = false
	DEBUG_SERVER  = false

	// errors
	ErrAlreadyDestroyed = errors.New("already destroyed")

	// regexps
	ChannelNameExpr = regexp.MustCompile(`^[&!#+][\pL\pN]{1,63}$`)
	NicknameExpr    = regexp.MustCompile(
		"^[\\pL\\[\\]{}^`_][\\pL\\pN\\[\\]{}^`_]{1,31}$")
)

const (
	SEM_VER       = "ergonomadic-1.2.11"
	CRLF          = "\r\n"
	MAX_REPLY_LEN = 512 - len(CRLF)

	LOGIN_TIMEOUT = time.Minute / 2 // how long the client has to login
	IDLE_TIMEOUT  = time.Minute     // how long before a client is considered idle
	QUIT_TIMEOUT  = time.Minute     // how long after idle before a client is kicked

	// string codes
	AWAY    StringCode = "AWAY"
	CAP     StringCode = "CAP"
	DEBUG   StringCode = "DEBUG"
	ERROR   StringCode = "ERROR"
	INVITE  StringCode = "INVITE"
	ISON    StringCode = "ISON"
	JOIN    StringCode = "JOIN"
	KICK    StringCode = "KICK"
	KILL    StringCode = "KILL"
	LIST    StringCode = "LIST"
	MODE    StringCode = "MODE"
	MOTD    StringCode = "MOTD"
	NAMES   StringCode = "NAMES"
	NICK    StringCode = "NICK"
	NOTICE  StringCode = "NOTICE"
	OPER    StringCode = "OPER"
	PART    StringCode = "PART"
	PASS    StringCode = "PASS"
	PING    StringCode = "PING"
	PONG    StringCode = "PONG"
	PRIVMSG StringCode = "PRIVMSG"
	PROXY   StringCode = "PROXY"
	QUIT    StringCode = "QUIT"
	TIME    StringCode = "TIME"
	TOPIC   StringCode = "TOPIC"
	USER    StringCode = "USER"
	VERSION StringCode = "VERSION"
	WHO     StringCode = "WHO"
	WHOIS   StringCode = "WHOIS"

	// numeric codes
	RPL_WELCOME           NumericCode = 1
	RPL_YOURHOST          NumericCode = 2
	RPL_CREATED           NumericCode = 3
	RPL_MYINFO            NumericCode = 4
	RPL_BOUNCE            NumericCode = 5
	RPL_TRACELINK         NumericCode = 200
	RPL_TRACECONNECTING   NumericCode = 201
	RPL_TRACEHANDSHAKE    NumericCode = 202
	RPL_TRACEUNKNOWN      NumericCode = 203
	RPL_TRACEOPERATOR     NumericCode = 204
	RPL_TRACEUSER         NumericCode = 205
	RPL_TRACESERVER       NumericCode = 206
	RPL_TRACESERVICE      NumericCode = 207
	RPL_TRACENEWTYPE      NumericCode = 208
	RPL_TRACECLASS        NumericCode = 209
	RPL_TRACERECONNECT    NumericCode = 210
	RPL_STATSLINKINFO     NumericCode = 211
	RPL_STATSCOMMANDS     NumericCode = 212
	RPL_ENDOFSTATS        NumericCode = 219
	RPL_UMODEIS           NumericCode = 221
	RPL_SERVLIST          NumericCode = 234
	RPL_SERVLISTEND       NumericCode = 235
	RPL_STATSUPTIME       NumericCode = 242
	RPL_STATSOLINE        NumericCode = 243
	RPL_LUSERCLIENT       NumericCode = 251
	RPL_LUSEROP           NumericCode = 252
	RPL_LUSERUNKNOWN      NumericCode = 253
	RPL_LUSERCHANNELS     NumericCode = 254
	RPL_LUSERME           NumericCode = 255
	RPL_ADMINME           NumericCode = 256
	RPL_ADMINLOC1         NumericCode = 257
	RPL_ADMINLOC2         NumericCode = 258
	RPL_ADMINEMAIL        NumericCode = 259
	RPL_TRACELOG          NumericCode = 261
	RPL_TRACEEND          NumericCode = 262
	RPL_TRYAGAIN          NumericCode = 263
	RPL_AWAY              NumericCode = 301
	RPL_USERHOST          NumericCode = 302
	RPL_ISON              NumericCode = 303
	RPL_UNAWAY            NumericCode = 305
	RPL_NOWAWAY           NumericCode = 306
	RPL_WHOISUSER         NumericCode = 311
	RPL_WHOISSERVER       NumericCode = 312
	RPL_WHOISOPERATOR     NumericCode = 313
	RPL_WHOWASUSER        NumericCode = 314
	RPL_ENDOFWHO          NumericCode = 315
	RPL_WHOISIDLE         NumericCode = 317
	RPL_ENDOFWHOIS        NumericCode = 318
	RPL_WHOISCHANNELS     NumericCode = 319
	RPL_LIST              NumericCode = 322
	RPL_LISTEND           NumericCode = 323
	RPL_CHANNELMODEIS     NumericCode = 324
	RPL_UNIQOPIS          NumericCode = 325
	RPL_NOTOPIC           NumericCode = 331
	RPL_TOPIC             NumericCode = 332
	RPL_INVITING          NumericCode = 341
	RPL_SUMMONING         NumericCode = 342
	RPL_INVITELIST        NumericCode = 346
	RPL_ENDOFINVITELIST   NumericCode = 347
	RPL_EXCEPTLIST        NumericCode = 348
	RPL_ENDOFEXCEPTLIST   NumericCode = 349
	RPL_VERSION           NumericCode = 351
	RPL_WHOREPLY          NumericCode = 352
	RPL_NAMREPLY          NumericCode = 353
	RPL_LINKS             NumericCode = 364
	RPL_ENDOFLINKS        NumericCode = 365
	RPL_ENDOFNAMES        NumericCode = 366
	RPL_BANLIST           NumericCode = 367
	RPL_ENDOFBANLIST      NumericCode = 368
	RPL_ENDOFWHOWAS       NumericCode = 369
	RPL_INFO              NumericCode = 371
	RPL_MOTD              NumericCode = 372
	RPL_ENDOFINFO         NumericCode = 374
	RPL_MOTDSTART         NumericCode = 375
	RPL_ENDOFMOTD         NumericCode = 376
	RPL_YOUREOPER         NumericCode = 381
	RPL_REHASHING         NumericCode = 382
	RPL_YOURESERVICE      NumericCode = 383
	RPL_TIME              NumericCode = 391
	RPL_USERSSTART        NumericCode = 392
	RPL_USERS             NumericCode = 393
	RPL_ENDOFUSERS        NumericCode = 394
	RPL_NOUSERS           NumericCode = 395
	ERR_NOSUCHNICK        NumericCode = 401
	ERR_NOSUCHSERVER      NumericCode = 402
	ERR_NOSUCHCHANNEL     NumericCode = 403
	ERR_CANNOTSENDTOCHAN  NumericCode = 404
	ERR_TOOMANYCHANNELS   NumericCode = 405
	ERR_WASNOSUCHNICK     NumericCode = 406
	ERR_TOOMANYTARGETS    NumericCode = 407
	ERR_NOSUCHSERVICE     NumericCode = 408
	ERR_NOORIGIN          NumericCode = 409
	ERR_NORECIPIENT       NumericCode = 411
	ERR_NOTEXTTOSEND      NumericCode = 412
	ERR_NOTOPLEVEL        NumericCode = 413
	ERR_WILDTOPLEVEL      NumericCode = 414
	ERR_BADMASK           NumericCode = 415
	ERR_UNKNOWNCOMMAND    NumericCode = 421
	ERR_NOMOTD            NumericCode = 422
	ERR_NOADMININFO       NumericCode = 423
	ERR_FILEERROR         NumericCode = 424
	ERR_NONICKNAMEGIVEN   NumericCode = 431
	ERR_ERRONEUSNICKNAME  NumericCode = 432
	ERR_NICKNAMEINUSE     NumericCode = 433
	ERR_NICKCOLLISION     NumericCode = 436
	ERR_UNAVAILRESOURCE   NumericCode = 437
	ERR_USERNOTINCHANNEL  NumericCode = 441
	ERR_NOTONCHANNEL      NumericCode = 442
	ERR_USERONCHANNEL     NumericCode = 443
	ERR_NOLOGIN           NumericCode = 444
	ERR_SUMMONDISABLED    NumericCode = 445
	ERR_USERSDISABLED     NumericCode = 446
	ERR_NOTREGISTERED     NumericCode = 451
	ERR_NEEDMOREPARAMS    NumericCode = 461
	ERR_ALREADYREGISTRED  NumericCode = 462
	ERR_NOPERMFORHOST     NumericCode = 463
	ERR_PASSWDMISMATCH    NumericCode = 464
	ERR_YOUREBANNEDCREEP  NumericCode = 465
	ERR_YOUWILLBEBANNED   NumericCode = 466
	ERR_KEYSET            NumericCode = 467
	ERR_CHANNELISFULL     NumericCode = 471
	ERR_UNKNOWNMODE       NumericCode = 472
	ERR_INVITEONLYCHAN    NumericCode = 473
	ERR_BANNEDFROMCHAN    NumericCode = 474
	ERR_BADCHANNELKEY     NumericCode = 475
	ERR_BADCHANMASK       NumericCode = 476
	ERR_NOCHANMODES       NumericCode = 477
	ERR_BANLISTFULL       NumericCode = 478
	ERR_NOPRIVILEGES      NumericCode = 481
	ERR_CHANOPRIVSNEEDED  NumericCode = 482
	ERR_CANTKILLSERVER    NumericCode = 483
	ERR_RESTRICTED        NumericCode = 484
	ERR_UNIQOPPRIVSNEEDED NumericCode = 485
	ERR_NOOPERHOST        NumericCode = 491
	ERR_UMODEUNKNOWNFLAG  NumericCode = 501
	ERR_USERSDONTMATCH    NumericCode = 502

	Add    ModeOp = '+'
	List   ModeOp = '='
	Remove ModeOp = '-'

	Away          UserMode = 'a'
	Invisible     UserMode = 'i'
	LocalOperator UserMode = 'O'
	Operator      UserMode = 'o'
	Restricted    UserMode = 'r'
	ServerNotice  UserMode = 's' // deprecated
	WallOps       UserMode = 'w'

	Anonymous       ChannelMode = 'a' // flag
	BanMask         ChannelMode = 'b' // arg
	ChannelCreator  ChannelMode = 'O' // flag
	ChannelOperator ChannelMode = 'o' // arg
	ExceptMask      ChannelMode = 'e' // arg
	InviteMask      ChannelMode = 'I' // arg
	InviteOnly      ChannelMode = 'i' // flag
	Key             ChannelMode = 'k' // flag arg
	Moderated       ChannelMode = 'm' // flag
	NoOutside       ChannelMode = 'n' // flag
	OpOnlyTopic     ChannelMode = 't' // flag
	Persistent      ChannelMode = 'P' // flag
	Private         ChannelMode = 'p' // flag
	Quiet           ChannelMode = 'q' // flag
	ReOp            ChannelMode = 'r' // flag
	Secret          ChannelMode = 's' // flag, deprecated
	UserLimit       ChannelMode = 'l' // flag arg
	Voice           ChannelMode = 'v' // arg
)

const (
	Authorization Phase = iota
	Registration  Phase = iota
	Normal        Phase = iota
)
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"

	"github.com/buger/jsonparser"
	"github.com/garyburd/redigo/redis"
	"github.com/gorilla/mux"
	uuid "github.com/satori/go.uuid"
	"github.com/syhlion/redisocket.v2"
	"github.com/syhlion/requestwork.v2"
)

type UserHandler func(u *User) (err error)

var DefaultSubHandler = func(channel string, data []byte) (d []byte, err error) {
	return data, nil
}

type User struct {
	id       string
	channels map[string]bool
	appKey   string
	request  *http.Request
	*redisocket.Client
}

type WsManager struct {
	users map[*User]bool
	*sync.RWMutex
	pool *redis.Pool
	*redisocket.Hub
	worker *requestwork.Worker
}

func (wm *WsManager) Auth(w http.ResponseWriter, r *http.Request) {
	jwt := r.FormValue("jwt")

	v := url.Values{}
	v.Add("data", jwt)
	req, err := http.NewRequest("POST", decode_service, bytes.NewBufferString(v.Encode()))

	if err != nil {
		logger.GetRequestEntry(r).Warn(err)
		http.Error(w, "jwt decode fail", http.StatusUnauthorized)
		return
	}
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Add("Content-Length", strconv.Itoa(len(v.Encode())))

	logger.GetRequestEntry(r).Debugf("request jwt: %s", jwt)
	ctx, _ := context.WithTimeout(context.Background(), 30*time.Second)
	a := &JwtPack{}
	err = wm.worker.Execute(ctx, req, func(resp *http.Response, e error) (err error) {
		if e != nil {
			logger.Debug(e)
			return e
		}
		defer resp.Body.Close()
		err = json.NewDecoder(resp.Body).Decode(a)
		if err != nil {
			logger.Debug(err)
			return
		}
		return
	})
	if err != nil {
		logger.GetRequestEntry(r).Warn(err)
		http.Error(w, "jwt decode fail", http.StatusUnauthorized)
		return
	}
	conn := wm.pool.Get()
	defer conn.Close()
	b, err := json.Marshal(a.Gusher)
	if err != nil {
		logger.GetRequestEntry(r).Debug(err)
		http.Error(w, "jwt decode fail", http.StatusUnauthorized)
		return
	}
	uid := uuid.NewV1()
	conn.Send("SET", uid.String(), string(b))
	conn.Send("EXPIRE", uid.String(), 30)
	conn.Flush()
	w.Header().Set("Content-Type", "application/json;charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode(struct {
		Token string `json:"token"`
	}{
		Token: uid.String(),
	})

}

func (wm *WsManager) Count() int {
	return len(wm.users)
}

func (wm *WsManager) Connect(w http.ResponseWriter, r *http.Request) {
	params := mux.Vars(r)
	appKey := params["app_key"]
	token := r.FormValue("token")
	if appKey == "" || token == "" {
		logger.GetRequestEntry(r).Warn("app_key or token is nil")
		http.Error(w, "app_key is nil", http.StatusUnauthorized)
		return
	}
	conn := wm.pool.Get()
	defer conn.Close()
	reply, err := redis.Bytes(conn.Do("GET", token))
	if err != nil {
		logger.GetRequestEntry(r).Warn(err)
		http.Error(w, "token error", http.StatusUnauthorized)
		return
	}
	auth := Auth{}
	err = json.Unmarshal(reply, &auth)
	if err != nil {
		logger.GetRequestEntry(r).Warn(err)
		http.Error(w, "token error", http.StatusUnauthorized)
		return
	}
	if appKey != auth.AppKey {
		http.Error(w, "appkey error", http.StatusUnauthorized)
		return
	}
	channels := make(map[string]bool)
	for _, v := range auth.Channels {
		channels[v] = false
	}

	s, err := wm.Upgrade(w, r, nil)
	if err != nil {
		logger.GetRequestEntry(r).Warnf("upgrade ws connection %s", err)
		return
	}

	u := &User{
		appKey:   appKey,
		request:  r,
		Client:   s,
		channels: channels,
	}
	wm.Lock()
	wm.users[u] = true
	wm.Unlock()
	logger.GetRequestEntry(r).Debug("user listen start")
	err = u.Listen(func(data []byte) (err error) {
		logger.GetRequestEntry(r).Debugf("client receive command %s", data)
		h, err := CommanRouter(data)
		if err != nil {
			return
		}
		err = h(u)
		if err != nil {
			return
		}

		return
	})
	if err != nil {
		logger.GetRequestEntry(r).Debugf("disconnect %s", err)
	}
	wm.Disconnect(u)

}

func (wm *WsManager) Disconnect(u *User) {
	wm.Lock()
	delete(wm.users, u)
	wm.Unlock()

}
func (wm *WsManager) Close() {
	for u, _ := range wm.users {
		u.Close()
	}
}
func SubscribeCommand(data []byte) (h UserHandler, err error) {

	channel, err := jsonparser.GetString(data, "channel")
	if err != nil {
		return
	}
	h = func(u *User) (err error) {
		command := &ChannelCommand{}
		var reply []byte
		if _, ok := u.channels[channel]; ok {
			logger.GetRequestEntry(u.request).Debugf("sub %s@%s channel", u.appKey, channel)
			u.On(u.appKey+"@"+channel, DefaultSubHandler)
			u.channels[channel] = true
			command.Event = SubscribeReplySucceeded
			command.Data.Channel = channel
			reply, err = json.Marshal(command)
			if err != nil {
				logger.GetRequestEntry(u.request).Debugf("sub sucess reply %s", err)
			}
		} else {
			command.Event = SubscribeReplyError
			reply, err = json.Marshal(command)
			if err != nil {
				logger.GetRequestEntry(u.request).Debugf("sub error reply %s", err)
			}

		}

		u.Send(reply)
		return
	}
	return
}
func UnSubscribeCommand(data []byte) (h UserHandler, err error) {
	channel, err := jsonparser.GetString(data, "channel")
	if err != nil {
		return
	}
	h = func(u *User) (err error) {
		command := &ChannelCommand{}
		var reply []byte
		//反訂閱處理
		if _, ok := u.channels[channel]; ok {
			logger.GetRequestEntry(u.request).Debugf("unsub %s@%s channel", u.appKey, channel)
			u.Off(u.appKey + "@" + channel)
			u.channels[command.Data.Channel] = false
			command.Event = UnSubscribeReplySucceeded
			command.Data.Channel = channel
			reply, err = json.Marshal(command)
			if err != nil {
				logger.GetRequestEntry(u.request).Debugf("unsub sucess reply %s", err)
			}
		} else {
			command.Event = UnSubscribeReplyError
			reply, err = json.Marshal(command)
			if err != nil {
				logger.GetRequestEntry(u.request).Debugf("unsub error reply %s", err)
			}
		}
		u.Send(reply)
		return
	}
	return
}

func CommanRouter(data []byte) (h UserHandler, err error) {

	val, err := jsonparser.GetString(data, "event")
	if err != nil {
		return
	}
	d, _, _, err := jsonparser.Get(data, "data")
	if err != nil {
		return
	}
	switch val {
	case SubscribeEvent:
		h, err = SubscribeCommand(d)
		break
	case UnSubscribeEvent:
		h, err = UnSubscribeCommand(d)
		break
	default:
		err = errors.New("event errors")
		break
	}
	return
}
<commit_msg>add auth api allow crossdoman<commit_after>package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"

	"github.com/buger/jsonparser"
	"github.com/garyburd/redigo/redis"
	"github.com/gorilla/mux"
	uuid "github.com/satori/go.uuid"
	"github.com/syhlion/redisocket.v2"
	"github.com/syhlion/requestwork.v2"
)

type UserHandler func(u *User) (err error)

var DefaultSubHandler = func(channel string, data []byte) (d []byte, err error) {
	return data, nil
}

type User struct {
	id       string
	channels map[string]bool
	appKey   string
	request  *http.Request
	*redisocket.Client
}

type WsManager struct {
	users map[*User]bool
	*sync.RWMutex
	pool *redis.Pool
	*redisocket.Hub
	worker *requestwork.Worker
}

func (wm *WsManager) Auth(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
	jwt := r.FormValue("jwt")

	v := url.Values{}
	v.Add("data", jwt)
	req, err := http.NewRequest("POST", decode_service, bytes.NewBufferString(v.Encode()))

	if err != nil {
		logger.GetRequestEntry(r).Warn(err)
		http.Error(w, "jwt decode fail", http.StatusUnauthorized)
		return
	}
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Add("Content-Length", strconv.Itoa(len(v.Encode())))

	logger.GetRequestEntry(r).Debugf("request jwt: %s", jwt)
	ctx, _ := context.WithTimeout(context.Background(), 30*time.Second)
	a := &JwtPack{}
	err = wm.worker.Execute(ctx, req, func(resp *http.Response, e error) (err error) {
		if e != nil {
			logger.Debug(e)
			return e
		}
		defer resp.Body.Close()
		err = json.NewDecoder(resp.Body).Decode(a)
		if err != nil {
			logger.Debug(err)
			return
		}
		return
	})
	if err != nil {
		logger.GetRequestEntry(r).Warn(err)
		http.Error(w, "jwt decode fail", http.StatusUnauthorized)
		return
	}
	conn := wm.pool.Get()
	defer conn.Close()
	b, err := json.Marshal(a.Gusher)
	if err != nil {
		logger.GetRequestEntry(r).Debug(err)
		http.Error(w, "jwt decode fail", http.StatusUnauthorized)
		return
	}
	uid := uuid.NewV1()
	conn.Send("SET", uid.String(), string(b))
	conn.Send("EXPIRE", uid.String(), 60)
	conn.Flush()
	w.Header().Set("Content-Type", "application/json;charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode(struct {
		Token string `json:"token"`
	}{
		Token: uid.String(),
	})

}

func (wm *WsManager) Count() int {
	return len(wm.users)
}

func (wm *WsManager) Connect(w http.ResponseWriter, r *http.Request) {
	params := mux.Vars(r)
	appKey := params["app_key"]
	token := r.FormValue("token")
	if appKey == "" || token == "" {
		logger.GetRequestEntry(r).Warn("app_key or token is nil")
		http.Error(w, "app_key is nil", http.StatusUnauthorized)
		return
	}
	conn := wm.pool.Get()
	defer conn.Close()
	reply, err := redis.Bytes(conn.Do("GET", token))
	if err != nil {
		logger.GetRequestEntry(r).Warn(err)
		http.Error(w, "token error", http.StatusUnauthorized)
		return
	}
	auth := Auth{}
	err = json.Unmarshal(reply, &auth)
	if err != nil {
		logger.GetRequestEntry(r).Warn(err)
		http.Error(w, "token error", http.StatusUnauthorized)
		return
	}
	if appKey != auth.AppKey {
		http.Error(w, "appkey error", http.StatusUnauthorized)
		return
	}
	channels := make(map[string]bool)
	for _, v := range auth.Channels {
		channels[v] = false
	}

	s, err := wm.Upgrade(w, r, nil)
	if err != nil {
		logger.GetRequestEntry(r).Warnf("upgrade ws connection %s", err)
		return
	}

	u := &User{
		appKey:   appKey,
		request:  r,
		Client:   s,
		channels: channels,
	}
	wm.Lock()
	wm.users[u] = true
	wm.Unlock()
	logger.GetRequestEntry(r).Debug("user listen start")
	err = u.Listen(func(data []byte) (err error) {
		logger.GetRequestEntry(r).Debugf("client receive command %s", data)
		h, err := CommanRouter(data)
		if err != nil {
			return
		}
		err = h(u)
		if err != nil {
			return
		}

		return
	})
	if err != nil {
		logger.GetRequestEntry(r).Debugf("disconnect %s", err)
	}
	wm.Disconnect(u)

}

func (wm *WsManager) Disconnect(u *User) {
	wm.Lock()
	delete(wm.users, u)
	wm.Unlock()

}
func (wm *WsManager) Close() {
	for u, _ := range wm.users {
		u.Close()
	}
}
func SubscribeCommand(data []byte) (h UserHandler, err error) {

	channel, err := jsonparser.GetString(data, "channel")
	if err != nil {
		return
	}
	h = func(u *User) (err error) {
		command := &ChannelCommand{}
		var reply []byte
		if _, ok := u.channels[channel]; ok {
			logger.GetRequestEntry(u.request).Debugf("sub %s@%s channel", u.appKey, channel)
			u.On(u.appKey+"@"+channel, DefaultSubHandler)
			u.channels[channel] = true
			command.Event = SubscribeReplySucceeded
			command.Data.Channel = channel
			reply, err = json.Marshal(command)
			if err != nil {
				logger.GetRequestEntry(u.request).Debugf("sub sucess reply %s", err)
			}
		} else {
			command.Event = SubscribeReplyError
			reply, err = json.Marshal(command)
			if err != nil {
				logger.GetRequestEntry(u.request).Debugf("sub error reply %s", err)
			}

		}

		u.Send(reply)
		return
	}
	return
}
func UnSubscribeCommand(data []byte) (h UserHandler, err error) {
	channel, err := jsonparser.GetString(data, "channel")
	if err != nil {
		return
	}
	h = func(u *User) (err error) {
		command := &ChannelCommand{}
		var reply []byte
		//反訂閱處理
		if _, ok := u.channels[channel]; ok {
			logger.GetRequestEntry(u.request).Debugf("unsub %s@%s channel", u.appKey, channel)
			u.Off(u.appKey + "@" + channel)
			u.channels[command.Data.Channel] = false
			command.Event = UnSubscribeReplySucceeded
			command.Data.Channel = channel
			reply, err = json.Marshal(command)
			if err != nil {
				logger.GetRequestEntry(u.request).Debugf("unsub sucess reply %s", err)
			}
		} else {
			command.Event = UnSubscribeReplyError
			reply, err = json.Marshal(command)
			if err != nil {
				logger.GetRequestEntry(u.request).Debugf("unsub error reply %s", err)
			}
		}
		u.Send(reply)
		return
	}
	return
}

func CommanRouter(data []byte) (h UserHandler, err error) {

	val, err := jsonparser.GetString(data, "event")
	if err != nil {
		return
	}
	d, _, _, err := jsonparser.Get(data, "data")
	if err != nil {
		return
	}
	switch val {
	case SubscribeEvent:
		h, err = SubscribeCommand(d)
		break
	case UnSubscribeEvent:
		h, err = UnSubscribeCommand(d)
		break
	default:
		err = errors.New("event errors")
		break
	}
	return
}
<|endoftext|>


<commit_before>package mergo

import (
	"testing"

	"github.com/imdario/mergo"
)

func TestIssue121WithSliceDeepCopy(t *testing.T) {
	dst := map[string]interface{}{
		"inter": map[string]interface{}{
			"a": "1",
			"b": "2",
		},
	}

	src := map[string]interface{}{
		"inter": map[string]interface{}{
			"a": "3",
			"c": "4",
		},
	}

	if err := mergo.Merge(&dst, src, WithSliceDeepCopy); err != nil {
		t.Fatalf("Error during the merge: %v", err)
	}

	if dst["inter"].(map[string]interface{})["a"].(string) != "3" {
		t.Error("inter.a should equal '3'")
	}

	if dst["inter"].(map[string]interface{})["c"].(string) != "4" {
		t.Error("inter.c should equal '4'")
	}
}
<commit_msg>Fixed typo<commit_after>package mergo

import (
	"testing"
)

func TestIssue121WithSliceDeepCopy(t *testing.T) {
	dst := map[string]interface{}{
		"inter": map[string]interface{}{
			"a": "1",
			"b": "2",
		},
	}

	src := map[string]interface{}{
		"inter": map[string]interface{}{
			"a": "3",
			"c": "4",
		},
	}

	if err := Merge(&dst, src, WithSliceDeepCopy); err != nil {
		t.Fatalf("Error during the merge: %v", err)
	}

	if dst["inter"].(map[string]interface{})["a"].(string) != "3" {
		t.Error("inter.a should equal '3'")
	}

	if dst["inter"].(map[string]interface{})["c"].(string) != "4" {
		t.Error("inter.c should equal '4'")
	}
}
<|endoftext|>


<commit_before>package itchfs

import (
	"fmt"
	"net/http"
	"net/url"

	itchio "github.com/itchio/go-itchio"
	"github.com/itchio/httpkit/httpfile"
)

type ItchFS struct {
	ItchServer string
}

func (ifs *ItchFS) Scheme() string {
	return "itchfs"
}

func needsRenewal(res *http.Response, body []byte) bool {
	if res.StatusCode == 400 {
		// XXX: could parse XML / make sure it's expired URL and not something else,
		// but 400 is a good enough indicator for GCS
		return true
	}
	return false
}

func (ifs *ItchFS) MakeResource(u *url.URL) (httpfile.GetURLFunc, httpfile.NeedsRenewalFunc, error) {
	if u.Host != "" {
		return nil, nil, fmt.Errorf("invalid itchfs URL (must start with itchfs:///): %s", u.String())
	}

	vals := u.Query()

	apiKey := vals.Get("api_key")
	if apiKey == "" {
		return nil, nil, fmt.Errorf("missing API key")
	}

	itchClient := itchio.ClientWithKey(apiKey)
	if ifs.ItchServer != "" {
		itchClient.SetServer(ifs.ItchServer)
	}

	source, err := ObtainSource(itchClient, u.Path)
	if err != nil {
		return nil, nil, err
	}

	getURL, err := source.makeGetURL()
	if err != nil {
		return nil, nil, err
	}

	return getURL, needsRenewal, nil
}
<commit_msg>More compatible expired URL support<commit_after>package itchfs

import (
	"fmt"
	"net/http"
	"net/url"

	itchio "github.com/itchio/go-itchio"
	"github.com/itchio/httpkit/httpfile"
)

type ItchFS struct {
	ItchServer string
}

func (ifs *ItchFS) Scheme() string {
	return "itchfs"
}

func needsRenewal(res *http.Response, body []byte) bool {
	if res.StatusCode == 400 {
		// XXX: could parse XML / make sure it's expired URL and not something else,
		// but 400 is a good enough indicator for GCS
		return true
	}
	if res.StatusCode == 403 {
		// 403 is a good indicator for Highwinds - additionally, we could parse the URL
		// and compare the expires timestamp ourselves
	}
	return false
}

func (ifs *ItchFS) MakeResource(u *url.URL) (httpfile.GetURLFunc, httpfile.NeedsRenewalFunc, error) {
	if u.Host != "" {
		return nil, nil, fmt.Errorf("invalid itchfs URL (must start with itchfs:///): %s", u.String())
	}

	vals := u.Query()

	apiKey := vals.Get("api_key")
	if apiKey == "" {
		return nil, nil, fmt.Errorf("missing API key")
	}

	itchClient := itchio.ClientWithKey(apiKey)
	if ifs.ItchServer != "" {
		itchClient.SetServer(ifs.ItchServer)
	}

	source, err := ObtainSource(itchClient, u.Path)
	if err != nil {
		return nil, nil, err
	}

	getURL, err := source.makeGetURL()
	if err != nil {
		return nil, nil, err
	}

	return getURL, needsRenewal, nil
}
<|endoftext|>