<commit_before>package main

import (
	"github.com/jkomoros/sudoku"
	"github.com/mitchellh/go-wordwrap"
	"github.com/nsf/termbox-go"
)

type mainModel struct {
	grid     *sudoku.Grid
	selected *sudoku.Cell
	state    InputState
	//The size of the console output. Not used for much.
	outputWidth    int
	lastShownHint  *sudoku.SolveDirections
	consoleMessage string
	//if true, will zero out console message on turn of event loop.
	consoleMessageShort bool
	//If exitNow is flipped to true, we will quit at next turn of event loop.
	exitNow bool
	toggles []toggle
}

const (
	TOGGLE_SOLVED = iota
	TOGGLE_INVALID
	TOGGLE_FAST_MODE
	TOGGLE_MARK_MODE
)

type toggle struct {
	Value     func() bool
	Toggle    func()
	OnText    string
	OffText   string
	GridColor termbox.Attribute
}

func newModel() *mainModel {
	model := &mainModel{
		state: STATE_DEFAULT,
	}
	model.setUpToggles()
	model.EnsureSelected()
	return model
}

func (m *mainModel) setUpToggles() {

	//State variable for the closure
	var fastMode bool
	var markMode bool

	m.toggles = []toggle{
		//Solved
		{
			func() bool {
				return m.grid.Solved()
			},
			func() {
				//Do nothing; read only
			},
			"  SOLVED  ",
			" UNSOLVED ",
			termbox.ColorYellow,
		},
		//invalid
		{
			func() bool {
				return m.grid.Invalid()
			},
			func() {
				//Read only
			},
			" INVALID ",
			"  VALID  ",
			termbox.ColorRed,
		},
		//Fast mode
		{
			func() bool {
				return fastMode
			},
			func() {
				fastMode = !fastMode
			},
			"  FAST MODE  ",
			"             ",
			termbox.ColorBlue,
		},
		//Mark mode
		{
			func() bool {
				return markMode
			},
			func() {
				markMode = !markMode
			},
			" MARKING ",
			"         ",
			termbox.ColorCyan,
		},
	}
}

//EnterState attempts to set the model to the given state. The state object is
//given a chance to do initalization and potentially cancel the transition,
//leaving the model in the same state as before.
func (m *mainModel) EnterState(state InputState) {
	//SetState doesn't do much, it just makes it feel less weird than
	//STATE.enter(m) (which feels backward)

	if state.shouldEnter(m) {
		m.state = state
	}
}

//enterConfirmState is a special state to set
func (m *mainModel) enterConfirmState(msg string, defaultAction defaultOption, yesAction func(), noAction func()) {
	STATE_CONFIRM.msg = msg
	STATE_CONFIRM.defaultAction = defaultAction
	STATE_CONFIRM.yesAction = yesAction
	STATE_CONFIRM.noAction = noAction
	m.EnterState(STATE_CONFIRM)
}

func (m *mainModel) SetConsoleMessage(msg string, shortLived bool) {

	if m.outputWidth != 0 {
		//Wrap to fit in given size
		msg = wordwrap.WrapString(msg, uint(m.outputWidth))
	}

	m.consoleMessage = msg
	m.consoleMessageShort = shortLived
	m.lastShownHint = nil
}

func (m *mainModel) EndOfEventLoop() {
	if m.consoleMessageShort {
		m.ClearConsole()
	}
}

func (m *mainModel) ClearConsole() {
	m.consoleMessage = ""
	m.consoleMessageShort = false
	m.lastShownHint = nil
}

func (m *mainModel) StatusLine() string {
	return m.state.statusLine(m)
}

func (m *mainModel) Selected() *sudoku.Cell {
	return m.selected
}

func (m *mainModel) SetSelected(cell *sudoku.Cell) {
	if cell == m.selected {
		//Already done
		return
	}
	m.selected = cell
	m.state.newCellSelected(m)
}

func (m *mainModel) EnsureSelected() {
	m.EnsureGrid()
	//Ensures that at least one cell is selected.
	if m.Selected() == nil {
		m.SetSelected(m.grid.Cell(0, 0))
	}
}

func (m *mainModel) MoveSelectionLeft(fast bool) {
	m.EnsureSelected()
	r := m.Selected().Row()
	c := m.Selected().Col()
	for {
		c--
		if c < 0 {
			c = 0
		}
		if fast && m.grid.Cell(r, c).Number() != 0 {
			if c == 0 {
				//We're at the end and didn't find anything.
				//guess there's nothing to find.
				m.SetConsoleMessage(FAST_MODE_NO_OPEN_CELLS, true)
				return
			}
			continue
		}
		m.SetSelected(m.grid.Cell(r, c))
		break
	}
}

func (m *mainModel) MoveSelectionRight(fast bool) {
	m.EnsureSelected()
	r := m.Selected().Row()
	c := m.Selected().Col()
	for {
		c++
		if c >= sudoku.DIM {
			c = sudoku.DIM - 1
		}
		if fast && m.grid.Cell(r, c).Number() != 0 {
			if c == sudoku.DIM-1 {
				//We're at the end and didn't find anything.
				//guess there's nothing to find.
				m.SetConsoleMessage(FAST_MODE_NO_OPEN_CELLS, true)
				return
			}
			continue
		}
		m.SetSelected(m.grid.Cell(r, c))
		break
	}
}

func (m *mainModel) MoveSelectionUp(fast bool) {
	m.EnsureSelected()
	r := m.Selected().Row()
	c := m.Selected().Col()
	for {
		r--
		if r < 0 {
			r = 0
		}
		if fast && m.grid.Cell(r, c).Number() != 0 {
			if r == 0 {
				//We're at the end and didn't find anything.
				//guess there's nothing to find.
				m.SetConsoleMessage(FAST_MODE_NO_OPEN_CELLS, true)
				return
			}
			continue
		}
		m.SetSelected(m.grid.Cell(r, c))
		break
	}
}

func (m *mainModel) MoveSelectionDown(fast bool) {
	m.EnsureSelected()
	r := m.Selected().Row()
	c := m.Selected().Col()
	for {
		r++
		if r >= sudoku.DIM {
			r = sudoku.DIM - 1
		}
		if fast && m.grid.Cell(r, c).Number() != 0 {
			if r == sudoku.DIM-1 {
				//We're at the end and didn't find anything.
				//guess there's nothing to find.
				m.SetConsoleMessage(FAST_MODE_NO_OPEN_CELLS, true)
				return
			}
			continue
		}
		m.SetSelected(m.grid.Cell(r, c))
		break
	}
}

func (m *mainModel) FastMode() bool {
	return m.toggles[TOGGLE_FAST_MODE].Value()
}

func (m *mainModel) ToggleFastMode() {
	m.toggles[TOGGLE_FAST_MODE].Toggle()
}

func (m *mainModel) MarkMode() bool {
	return m.toggles[TOGGLE_MARK_MODE].Value()
}

func (m *mainModel) ToggleMarkMode() {
	m.toggles[TOGGLE_MARK_MODE].Toggle()
}

func (m *mainModel) EnsureGrid() {
	if m.grid == nil {
		m.NewGrid()
	}
}

func (m *mainModel) NewGrid() {
	oldCell := m.Selected()

	m.grid = sudoku.GenerateGrid(nil)
	//The currently selected cell is tied to the grid, so we need to fix it up.
	if oldCell != nil {
		m.SetSelected(oldCell.InGrid(m.grid))
	}
	m.grid.LockFilledCells()
}

func (m *mainModel) SetSelectedNumber(num int) {
	m.EnsureSelected()
	if m.Selected().Locked() {
		m.SetConsoleMessage(DEFAULT_MODE_FAIL_LOCKED, true)
		return
	}

	if m.Selected().Number() != num {
		m.Selected().SetNumber(num)
	} else {
		//If the number to set is already set, then empty the cell instead.
		m.Selected().SetNumber(0)
	}

	m.checkHintDone()
}

func (m *mainModel) checkHintDone() {
	if m.lastShownHint == nil {
		return
	}
	lastStep := m.lastShownHint.Steps[len(m.lastShownHint.Steps)-1]
	num := lastStep.TargetNums[0]
	cell := lastStep.TargetCells[0]
	if cell.InGrid(m.grid).Number() == num {
		m.ClearConsole()
	}
}

func (m *mainModel) ToggleSelectedMark(num int) {
	m.EnsureSelected()
	if m.Selected().Locked() {
		m.SetConsoleMessage(MARKS_MODE_FAIL_LOCKED, true)
		return
	}
	if m.Selected().Number() != 0 {
		m.SetConsoleMessage(MARKS_MODE_FAIL_NUMBER, true)
		return
	}
	m.Selected().SetMark(num, !m.Selected().Mark(num))
}

func (m *mainModel) FillSelectedWithLegalMarks() {
	m.EnsureSelected()
	m.Selected().ResetMarks()
	for _, num := range m.Selected().Possibilities() {
		m.Selected().SetMark(num, true)
	}
}

func (m *mainModel) RemoveInvalidMarksFromSelected() {
	m.EnsureSelected()
	for _, num := range m.Selected().Marks() {
		if !m.Selected().Possible(num) {
			m.Selected().SetMark(num, false)
		}
	}
}
<commit_msg>Flipped priority order of mark mode and fast mode, so mark mode's color comes first. Related to #177.<commit_after>package main

import (
	"github.com/jkomoros/sudoku"
	"github.com/mitchellh/go-wordwrap"
	"github.com/nsf/termbox-go"
)

type mainModel struct {
	grid     *sudoku.Grid
	selected *sudoku.Cell
	state    InputState
	//The size of the console output. Not used for much.
	outputWidth    int
	lastShownHint  *sudoku.SolveDirections
	consoleMessage string
	//if true, will zero out console message on turn of event loop.
	consoleMessageShort bool
	//If exitNow is flipped to true, we will quit at next turn of event loop.
	exitNow bool
	toggles []toggle
}

const (
	TOGGLE_SOLVED = iota
	TOGGLE_INVALID
	TOGGLE_MARK_MODE
	TOGGLE_FAST_MODE
)

type toggle struct {
	Value     func() bool
	Toggle    func()
	OnText    string
	OffText   string
	GridColor termbox.Attribute
}

func newModel() *mainModel {
	model := &mainModel{
		state: STATE_DEFAULT,
	}
	model.setUpToggles()
	model.EnsureSelected()
	return model
}

func (m *mainModel) setUpToggles() {

	//State variable for the closure
	var fastMode bool
	var markMode bool

	m.toggles = []toggle{
		//Solved
		{
			func() bool {
				return m.grid.Solved()
			},
			func() {
				//Do nothing; read only
			},
			"  SOLVED  ",
			" UNSOLVED ",
			termbox.ColorYellow,
		},
		//invalid
		{
			func() bool {
				return m.grid.Invalid()
			},
			func() {
				//Read only
			},
			" INVALID ",
			"  VALID  ",
			termbox.ColorRed,
		},
		//Mark mode
		{
			func() bool {
				return markMode
			},
			func() {
				markMode = !markMode
			},
			" MARKING ",
			"         ",
			termbox.ColorCyan,
		},
		//Fast mode
		{
			func() bool {
				return fastMode
			},
			func() {
				fastMode = !fastMode
			},
			"  FAST MODE  ",
			"             ",
			termbox.ColorBlue,
		},
	}
}

//EnterState attempts to set the model to the given state. The state object is
//given a chance to do initalization and potentially cancel the transition,
//leaving the model in the same state as before.
func (m *mainModel) EnterState(state InputState) {
	//SetState doesn't do much, it just makes it feel less weird than
	//STATE.enter(m) (which feels backward)

	if state.shouldEnter(m) {
		m.state = state
	}
}

//enterConfirmState is a special state to set
func (m *mainModel) enterConfirmState(msg string, defaultAction defaultOption, yesAction func(), noAction func()) {
	STATE_CONFIRM.msg = msg
	STATE_CONFIRM.defaultAction = defaultAction
	STATE_CONFIRM.yesAction = yesAction
	STATE_CONFIRM.noAction = noAction
	m.EnterState(STATE_CONFIRM)
}

func (m *mainModel) SetConsoleMessage(msg string, shortLived bool) {

	if m.outputWidth != 0 {
		//Wrap to fit in given size
		msg = wordwrap.WrapString(msg, uint(m.outputWidth))
	}

	m.consoleMessage = msg
	m.consoleMessageShort = shortLived
	m.lastShownHint = nil
}

func (m *mainModel) EndOfEventLoop() {
	if m.consoleMessageShort {
		m.ClearConsole()
	}
}

func (m *mainModel) ClearConsole() {
	m.consoleMessage = ""
	m.consoleMessageShort = false
	m.lastShownHint = nil
}

func (m *mainModel) StatusLine() string {
	return m.state.statusLine(m)
}

func (m *mainModel) Selected() *sudoku.Cell {
	return m.selected
}

func (m *mainModel) SetSelected(cell *sudoku.Cell) {
	if cell == m.selected {
		//Already done
		return
	}
	m.selected = cell
	m.state.newCellSelected(m)
}

func (m *mainModel) EnsureSelected() {
	m.EnsureGrid()
	//Ensures that at least one cell is selected.
	if m.Selected() == nil {
		m.SetSelected(m.grid.Cell(0, 0))
	}
}

func (m *mainModel) MoveSelectionLeft(fast bool) {
	m.EnsureSelected()
	r := m.Selected().Row()
	c := m.Selected().Col()
	for {
		c--
		if c < 0 {
			c = 0
		}
		if fast && m.grid.Cell(r, c).Number() != 0 {
			if c == 0 {
				//We're at the end and didn't find anything.
				//guess there's nothing to find.
				m.SetConsoleMessage(FAST_MODE_NO_OPEN_CELLS, true)
				return
			}
			continue
		}
		m.SetSelected(m.grid.Cell(r, c))
		break
	}
}

func (m *mainModel) MoveSelectionRight(fast bool) {
	m.EnsureSelected()
	r := m.Selected().Row()
	c := m.Selected().Col()
	for {
		c++
		if c >= sudoku.DIM {
			c = sudoku.DIM - 1
		}
		if fast && m.grid.Cell(r, c).Number() != 0 {
			if c == sudoku.DIM-1 {
				//We're at the end and didn't find anything.
				//guess there's nothing to find.
				m.SetConsoleMessage(FAST_MODE_NO_OPEN_CELLS, true)
				return
			}
			continue
		}
		m.SetSelected(m.grid.Cell(r, c))
		break
	}
}

func (m *mainModel) MoveSelectionUp(fast bool) {
	m.EnsureSelected()
	r := m.Selected().Row()
	c := m.Selected().Col()
	for {
		r--
		if r < 0 {
			r = 0
		}
		if fast && m.grid.Cell(r, c).Number() != 0 {
			if r == 0 {
				//We're at the end and didn't find anything.
				//guess there's nothing to find.
				m.SetConsoleMessage(FAST_MODE_NO_OPEN_CELLS, true)
				return
			}
			continue
		}
		m.SetSelected(m.grid.Cell(r, c))
		break
	}
}

func (m *mainModel) MoveSelectionDown(fast bool) {
	m.EnsureSelected()
	r := m.Selected().Row()
	c := m.Selected().Col()
	for {
		r++
		if r >= sudoku.DIM {
			r = sudoku.DIM - 1
		}
		if fast && m.grid.Cell(r, c).Number() != 0 {
			if r == sudoku.DIM-1 {
				//We're at the end and didn't find anything.
				//guess there's nothing to find.
				m.SetConsoleMessage(FAST_MODE_NO_OPEN_CELLS, true)
				return
			}
			continue
		}
		m.SetSelected(m.grid.Cell(r, c))
		break
	}
}

func (m *mainModel) FastMode() bool {
	return m.toggles[TOGGLE_FAST_MODE].Value()
}

func (m *mainModel) ToggleFastMode() {
	m.toggles[TOGGLE_FAST_MODE].Toggle()
}

func (m *mainModel) MarkMode() bool {
	return m.toggles[TOGGLE_MARK_MODE].Value()
}

func (m *mainModel) ToggleMarkMode() {
	m.toggles[TOGGLE_MARK_MODE].Toggle()
}

func (m *mainModel) EnsureGrid() {
	if m.grid == nil {
		m.NewGrid()
	}
}

func (m *mainModel) NewGrid() {
	oldCell := m.Selected()

	m.grid = sudoku.GenerateGrid(nil)
	//The currently selected cell is tied to the grid, so we need to fix it up.
	if oldCell != nil {
		m.SetSelected(oldCell.InGrid(m.grid))
	}
	m.grid.LockFilledCells()
}

func (m *mainModel) SetSelectedNumber(num int) {
	m.EnsureSelected()
	if m.Selected().Locked() {
		m.SetConsoleMessage(DEFAULT_MODE_FAIL_LOCKED, true)
		return
	}

	if m.Selected().Number() != num {
		m.Selected().SetNumber(num)
	} else {
		//If the number to set is already set, then empty the cell instead.
		m.Selected().SetNumber(0)
	}

	m.checkHintDone()
}

func (m *mainModel) checkHintDone() {
	if m.lastShownHint == nil {
		return
	}
	lastStep := m.lastShownHint.Steps[len(m.lastShownHint.Steps)-1]
	num := lastStep.TargetNums[0]
	cell := lastStep.TargetCells[0]
	if cell.InGrid(m.grid).Number() == num {
		m.ClearConsole()
	}
}

func (m *mainModel) ToggleSelectedMark(num int) {
	m.EnsureSelected()
	if m.Selected().Locked() {
		m.SetConsoleMessage(MARKS_MODE_FAIL_LOCKED, true)
		return
	}
	if m.Selected().Number() != 0 {
		m.SetConsoleMessage(MARKS_MODE_FAIL_NUMBER, true)
		return
	}
	m.Selected().SetMark(num, !m.Selected().Mark(num))
}

func (m *mainModel) FillSelectedWithLegalMarks() {
	m.EnsureSelected()
	m.Selected().ResetMarks()
	for _, num := range m.Selected().Possibilities() {
		m.Selected().SetMark(num, true)
	}
}

func (m *mainModel) RemoveInvalidMarksFromSelected() {
	m.EnsureSelected()
	for _, num := range m.Selected().Marks() {
		if !m.Selected().Possible(num) {
			m.Selected().SetMark(num, false)
		}
	}
}
<|endoftext|>


<commit_before>// +build !appengine

package main

import (
	"log"
	"net/http"
	"time"

	_ "github.com/jamesog/ip.jog.li"
)

func main() {
	s := &http.Server{
		Addr:         ":8000",
		Handler:      nil,
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 5 * time.Second,
	}
	log.Fatal(s.ListenAndServe())
}
<commit_msg>cmd/ip.jog.li: Read PORT from the environment<commit_after>// +build !appengine

package main

import (
	"log"
	"net/http"
	"os"
	"time"

	_ "github.com/jamesog/ip.jog.li"
)

func main() {
	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}
	s := &http.Server{
		Addr:         ":" + port,
		Handler:      nil,
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 5 * time.Second,
	}
	log.Fatal(s.ListenAndServe())
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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
//
//     https://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 main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"

	scheduler "cloud.google.com/go/scheduler/apiv1"
	backupfunction "github.com/cloudspannerecosystem/scheduled-backups"
	schedulerpb "google.golang.org/genproto/googleapis/cloud/scheduler/v1"
	"google.golang.org/genproto/protobuf/field_mask"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"gopkg.in/yaml.v2"
)

const defaultLocation = "us-central1"
const pubsubTopic = "cloud-spanner-scheduled-backups"
const jobPrefix = "spanner-backup"

// Project contains the information of a GCP project.
type Project struct {
	Name      string     `yaml:"name"`
	Instances []Instance `yaml:"instances"`
}

// Instance contains the information of an instance.
type Instance struct {
	Name      string     `yaml:"name"`
	Databases []Database `yaml:"databases"`
}

// Database contains the backup schedule configuration for a database.
type Database struct {
	Name     string `yaml:"name"`
	Schedule string `yaml:"schedule"`
	Expire   string `yaml:"expire"`
	Location string `yaml:"location"`
	TimeZone string `yaml:"time_zone"`
}

func main() {
	var filename string

	flag.StringVar(&filename, "config", "", "The file path of the config file in yaml format.")
	flag.Parse()

	if filename == "" {
		flag.Usage()
		os.Exit(2)
	}
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatalf("Failed to read file: %v", err)
	}

	var project Project

	err = yaml.Unmarshal(content, &project)
	if err != nil {
		log.Fatalf("Failed to parse the config file: %v", err)
	}

	ctx := context.Background()
	client, err := scheduler.NewCloudSchedulerClient(ctx)
	if err != nil {
		log.Fatalf("Failed to create a scheduler client: %v", err)
	}
	defer client.Close()

	topicPath := fmt.Sprintf("projects/%s/topics/%s", project.Name, pubsubTopic)

	for _, instance := range project.Instances {
		for _, db := range instance.Databases {
			dbPath := fmt.Sprintf("projects/%s/instances/%s/databases/%s", project.Name, instance.Name, db.Name)
			// Get the specified location. If not given, use the default one.
			loc := db.Location
			if loc == "" {
				loc = defaultLocation
			}
			locPath := fmt.Sprintf("projects/%s/locations/%s", project.Name, loc)
			jobID := fmt.Sprintf("%s-%s", jobPrefix, db.Name)
			jobName := fmt.Sprintf("%s/jobs/%s", locPath, jobID)

			err = updateJob(ctx, client, jobName, locPath, dbPath, topicPath, db)
			if err != nil {
				if errCode(err) == codes.NotFound {
					// Create a new job if the job does not exist.
					createJob(ctx, client, jobName, locPath, dbPath, topicPath, db)
				} else {
					log.Printf("Failed to update job %v: %v\n", jobName, err)
				}
			}
		}
	}
}

// errCode extracts the canonical error code from a Go error.
func errCode(err error) codes.Code {
	s, ok := status.FromError(err)
	if !ok {
		return codes.Unknown
	}
	return s.Code()
}

func updateJob(ctx context.Context, client *scheduler.CloudSchedulerClient, jobName, locPath, dbPath, topicPath string, db Database) error {
	meta := backupfunction.BackupParameters{Database: dbPath, Expire: db.Expire}
	data, err := json.Marshal(meta)
	if err != nil {
		log.Fatalf("Failed to marshal data: %v", err)
	}

	// Update a job.
	req := &schedulerpb.UpdateJobRequest{
		Job: &schedulerpb.Job{
			Name: jobName,
			Target: &schedulerpb.Job_PubsubTarget{
				PubsubTarget: &schedulerpb.PubsubTarget{
					TopicName: topicPath,
					Data:      data,
				},
			},
			Schedule: db.Schedule,
			TimeZone: db.TimeZone,
		},
		UpdateMask: &field_mask.FieldMask{
			Paths: []string{"schedule", "pubsub_target.data", "time_zone"},
		},
	}
	_, err = client.UpdateJob(ctx, req)
	if err == nil {
		log.Printf("Update the job %v.", jobName)
	}
	return err
}

func createJob(ctx context.Context, client *scheduler.CloudSchedulerClient, jobName, locPath, dbPath, topicPath string, db Database) {
	meta := backupfunction.BackupParameters{Database: dbPath, Expire: db.Expire}
	data, err := json.Marshal(meta)
	if err != nil {
		log.Fatalf("Failed to marshal data: %v", err)
	}

	// Create a new job.
	req := &schedulerpb.CreateJobRequest{
		Parent: locPath,
		Job: &schedulerpb.Job{
			Name:        jobName,
			Description: fmt.Sprintf("A scheduling job for Cloud Spanner database %s", dbPath),
			Target: &schedulerpb.Job_PubsubTarget{
				PubsubTarget: &schedulerpb.PubsubTarget{
					TopicName: topicPath,
					Data:      data,
				},
			},
			Schedule: db.Schedule,
			TimeZone: db.TimeZone,
		},
	}
	resp, err := client.CreateJob(ctx, req)
	if err != nil {
		log.Fatalf("Failed to create a cloud scheduler job: %v", err)
	}
	log.Printf("Create a scheduled backup job: %v\n", resp)
}
<commit_msg>feat: remove constants and make them configurable via flags (#12)<commit_after>// Copyright 2020 Google LLC
//
// 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
//
//     https://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 main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"

	scheduler "cloud.google.com/go/scheduler/apiv1"
	backupfunction "github.com/cloudspannerecosystem/scheduled-backups"
	schedulerpb "google.golang.org/genproto/googleapis/cloud/scheduler/v1"
	"google.golang.org/genproto/protobuf/field_mask"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"gopkg.in/yaml.v2"
)

// Project contains the information of a GCP project.
type Project struct {
	Name      string     `yaml:"name"`
	Instances []Instance `yaml:"instances"`
}

// Instance contains the information of an instance.
type Instance struct {
	Name      string     `yaml:"name"`
	Databases []Database `yaml:"databases"`
}

// Database contains the backup schedule configuration for a database.
type Database struct {
	Name     string `yaml:"name"`
	Schedule string `yaml:"schedule"`
	Expire   string `yaml:"expire"`
	Location string `yaml:"location"`
	TimeZone string `yaml:"time_zone"`
}

func main() {
	var filename, defaultLocation, pubsubTopic, jobPrefix string

	flag.StringVar(&filename, "config", "", "The file path of the config file in yaml format.")
	flag.StringVar(&defaultLocation, "location", "us-central1", "The location where the scheduler is deployed.")
	flag.StringVar(&pubsubTopic, "pubsubTopic", "cloud-spanner-scheduled-backups", "The PubSub topic where the scheduler sends to.")
	flag.StringVar(&jobPrefix, "jobPrefix", "spanner-backup", "The name prefix of a scheduler job.")
	flag.Parse()

	if filename == "" {
		flag.Usage()
		os.Exit(2)
	}
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatalf("Failed to read file: %v", err)
	}

	var project Project

	err = yaml.Unmarshal(content, &project)
	if err != nil {
		log.Fatalf("Failed to parse the config file: %v", err)
	}

	ctx := context.Background()
	client, err := scheduler.NewCloudSchedulerClient(ctx)
	if err != nil {
		log.Fatalf("Failed to create a scheduler client: %v", err)
	}
	defer client.Close()

	topicPath := fmt.Sprintf("projects/%s/topics/%s", project.Name, pubsubTopic)

	for _, instance := range project.Instances {
		for _, db := range instance.Databases {
			dbPath := fmt.Sprintf("projects/%s/instances/%s/databases/%s", project.Name, instance.Name, db.Name)
			// Get the specified location. If not given, use the default one.
			loc := db.Location
			if loc == "" {
				loc = defaultLocation
			}
			locPath := fmt.Sprintf("projects/%s/locations/%s", project.Name, loc)
			jobID := fmt.Sprintf("%s-%s", jobPrefix, db.Name)
			jobName := fmt.Sprintf("%s/jobs/%s", locPath, jobID)

			err = updateJob(ctx, client, jobName, locPath, dbPath, topicPath, db)
			if err != nil {
				if errCode(err) == codes.NotFound {
					// Create a new job if the job does not exist.
					createJob(ctx, client, jobName, locPath, dbPath, topicPath, db)
				} else {
					log.Printf("Failed to update job %v: %v\n", jobName, err)
				}
			}
		}
	}
}

// errCode extracts the canonical error code from a Go error.
func errCode(err error) codes.Code {
	s, ok := status.FromError(err)
	if !ok {
		return codes.Unknown
	}
	return s.Code()
}

func updateJob(ctx context.Context, client *scheduler.CloudSchedulerClient, jobName, locPath, dbPath, topicPath string, db Database) error {
	meta := backupfunction.BackupParameters{Database: dbPath, Expire: db.Expire}
	data, err := json.Marshal(meta)
	if err != nil {
		log.Fatalf("Failed to marshal data: %v", err)
	}

	// Update a job.
	req := &schedulerpb.UpdateJobRequest{
		Job: &schedulerpb.Job{
			Name: jobName,
			Target: &schedulerpb.Job_PubsubTarget{
				PubsubTarget: &schedulerpb.PubsubTarget{
					TopicName: topicPath,
					Data:      data,
				},
			},
			Schedule: db.Schedule,
			TimeZone: db.TimeZone,
		},
		UpdateMask: &field_mask.FieldMask{
			Paths: []string{"schedule", "pubsub_target.data", "time_zone"},
		},
	}
	_, err = client.UpdateJob(ctx, req)
	if err == nil {
		log.Printf("Update the job %v.", jobName)
	}
	return err
}

func createJob(ctx context.Context, client *scheduler.CloudSchedulerClient, jobName, locPath, dbPath, topicPath string, db Database) {
	meta := backupfunction.BackupParameters{Database: dbPath, Expire: db.Expire}
	data, err := json.Marshal(meta)
	if err != nil {
		log.Fatalf("Failed to marshal data: %v", err)
	}

	// Create a new job.
	req := &schedulerpb.CreateJobRequest{
		Parent: locPath,
		Job: &schedulerpb.Job{
			Name:        jobName,
			Description: fmt.Sprintf("A scheduling job for Cloud Spanner database %s", dbPath),
			Target: &schedulerpb.Job_PubsubTarget{
				PubsubTarget: &schedulerpb.PubsubTarget{
					TopicName: topicPath,
					Data:      data,
				},
			},
			Schedule: db.Schedule,
			TimeZone: db.TimeZone,
		},
	}
	resp, err := client.CreateJob(ctx, req)
	if err != nil {
		log.Fatalf("Failed to create a cloud scheduler job: %v", err)
	}
	log.Printf("Create a scheduled backup job: %v\n", resp)
}
<|endoftext|>


<commit_before>/*

Scollector is a metric collection agent for OpenTSDB 2.0.

tcollector (https://github.com/OpenTSDB/tcollector) is OpenTSDB's data
collection framework built for OpenTSDB 1.0. scollector aims to be tcollector
for OpenTSDB 2.0.

Unlike tcollector, scollector is a single binary with all collectors are
compiled into scollector itself. scollector supports external collectors, but
your goal should be to use those temporarily until the go version is written. It
has native collectors for Linux, Darwin, and Windows.

Usage:
	scollector [flag]

The flags are:
	-h="bosun"
		Bosun or OpenTSDB host; can optionally specify a port and scheme
		("http://bosun.example.com:8070"), but will default to
		http://bosun:80/
	-c=""
		external collectors directory
	-s=""
		SNMP host to poll of the format:
		"community@host[,community@host...]"
	-v=""
		vSphere host to poll of the format:
		"user:password@host[,user:password@host...]"
	-i=""
		ICMP host to ping of the format:
		"host[,host...]"
	-f=""
		filter collectors matching this term (regex)
	-l
		list enabled collectors
	-m
		disable sending of metadata
	-n
		disable sending of scollector self metrics
	-conf=""
		location of configuration file
	-hostname=""
		override for system hostname

Additional flags on Windows:
	-winsvc=""
		Windows Service management; can be: install, remove, start, stop

Debug flags:
	-d
		enables debug output
	-p
		print to screen instead of sending to a host
	-j
		with -p, prints JSON
	-fake=0
		generates X fake data points per second on the test.fake metric

scollector will, by default, report to the host `bosun`, which you should
configure on your local nameserver to point to your TSDB server. This makes it
possible to run scollector correctly with zero configuration or command line
flags.

Warning

scollector has not been tested outside of the Stack Exchange environment, and
thus may act incorrectly elsewhere.

scollector requires the new HTTP API of OpenTSDB 2.1 with gzip support, which is
currently only present in the "next" branch
(https://github.com/opentsdb/opentsdb/tree/next). Ensure that is in use if not
using the docker image.

Logs

If started with -p or -d, scollector logs to Stdout. Otherwise, on Unixes,
scollector logs to syslog. On Windows when started as a service, the Event Log
is used.

External Collectors

External collectors are executables that scollector invokes, collects output
from, and uses that like other collector data. The -c option specfies the
external collectors directory. It should contain numbered directories like
OpenTSDB tcollector expects. Any executable file in those directories is run
every N seconds, where N is the name of the directory. Use 0 for a program that
should be run continuously and simply pass data through to OpenTSDB (the program
will be restarted if it exits). Data output format is:

	metric timestamp value tag1=val1 tag2=val2 ...

Timestamp is in Unix format (seconds since epoch). Tags are optional. A host tag
is automatically added, but overridden if specified. Stderr output is passed to
scollector's log.

Configuration File

If scollector.conf exists in the same directory as the scollector executable, it
will set configuration flags. Configuration file values overwrite command line
flags. The configuration file is of the form key = value, one per line.
Supported keys are: host (-h), hostname (-hostname), filter (-f), coldir (-c),
snmp (-s), icmp (-i), vsphere (-v). Example:

	host = other-tsdb:1234
	filter = snmp
	snmp = com@theswitch

Windows

scollector has full Windows support. It can be run standalone, or installed as a
service (see -winsvc). The Event Log is used when installed as a service.

SNMP

By default, scollector will collect data from the current host and report it to
the TSDB server. scollector has an SNMP mode where it instead polls a given
host:

	scollector -s community@host[,community@host...]

Poll frequency currently defaults to 5 minutes. Some common OIDs regarding
interfaces are collected. Others can be added easily.

*/
package main
<commit_msg>cmd/scollector: update doc.go documentation<commit_after>/*

Scollector is a metric collection agent for OpenTSDB 2.0 and http://Bosun.org

tcollector (https://github.com/OpenTSDB/tcollector) is OpenTSDB's data
collection framework built for OpenTSDB 1.0. scollector aims to be tcollector
for OpenTSDB 2.0 and is one method of sending data to Bosun for monitoring.

Unlike tcollector, scollector is a single binary where all collectors are
compiled into scollector itself. scollector supports external collectors, but
your goal should be to use those temporarily until the go version is written or
the target system send data directly to OpenTSDB or Bosun. scollector has
native collectors for Linux, Darwin, and Windows and can pull data from other
systems such as AWS, SNMP, and vSphere.

Usage:
	scollector [flag]

The flags are:
	-aws=""
	    AWS keys and region, format: "access_key:secret_key@region".
	-b=0
	    OpenTSDB batch size. Used for debugging bad data.
	-c=""
	    External collectors directory.
	-conf=""
	    Location of configuration file. Defaults to scollector.conf in directory of
	    the scollector executable.
	-f=""
	    Filters collectors matching this term, multiple terms separated by comma.
	    Works with all other arguments.
	-freq="15"
	    Set the default frequency in seconds for most collectors.
	-h=""
	    Bosun or OpenTSDB host; can optionally specify a port and scheme
	    ("http://bosun.example.com:8070"), but will default to
	    http://bosun:80/
	-hostname=""
	    If set, use as value of host tag instead of system hostname.
	-i=""
	    ICMP host to ping of the format: "host[,host...]".
	-l
	    List available collectors.
	-m
	    Disable sending of metadata.
	-n
	    Disable sending of scollector self metrics.
	-s=""
	    SNMP host to poll of the format: "community@host[,community@host...]".
	-t=""
	    Tags to add to every datapoint in the format dc=ny,rack=3. If a collector
	    specifies the same tag key, this one will be overwritten. The host tag is
	    not supported.
	-u
	    Enables full hostnames: doesn't truncate to first ".".
	-v=""
	    vSphere host to poll of the format:
	    "user:password@host[,user:password@host...]".
	-version
	    Prints the version and exits.

Additional flags on Windows:
	-winsvc=""
		Windows Service management; can be: install, remove, start, stop

Debug flags:
	-d
		enables debug output
	-p
		print to screen instead of sending to a host
	-fake=0
		generates X fake data points per second on the test.fake metric

scollector will, by default, report to the host `bosun`, which you should
configure on your local nameserver to point to your TSDB server. This makes it
possible to run scollector correctly with zero configuration or command line
flags.

Warning

scollector has not been tested outside of the Stack Exchange environment, and
thus may act incorrectly elsewhere.

scollector requires the new HTTP API of OpenTSDB 2.1 with gzip support. Ensure
that is in use if not using the OpenTSDB docker image.

Logs

If started with -p or -d, scollector logs to Stdout. Otherwise, on Unixes,
scollector logs to syslog. On Windows when started as a service, the Event Log
is used.

External Collectors

External collectors are executables that scollector invokes, collects output
from, and uses that like other collector data. The -c option specfies the
external collectors directory. It should contain numbered directories like
OpenTSDB tcollector expects. Any executable file in those directories is run
every N seconds, where N is the name of the directory. Use 0 for a program that
should be run continuously and simply pass data through to OpenTSDB (the program
will be restarted if it exits). Data output format is:

	metric timestamp value tag1=val1 tag2=val2 ...

Timestamp is in Unix format (seconds since epoch). Tags are optional. A host tag
is automatically added, but overridden if specified. Stderr output is passed to
scollector's log.

Configuration File

If scollector.conf exists in the same directory as the scollector executable or
is specified via the -conf="" flag, it's content will be used to set
configuration flags. Configuration file values overwrite command line flags.
The configuration file is of the form key = value, one per line. Supported keys
are: host (-h), hostname (-hostname), filter (-f), coldir (-c),
snmp (-s), icmp (-i), vsphere (-v). Example:

	host = other-tsdb:1234
	filter = snmp
	snmp = com@theswitch

There also are additional values that are used to configure specific collectors.

	// Linux processes to monitor use the form "command,name,command line regex"
	process = ruby,puppet-agent,puppet
	process = java,opentsdb,opentsdb
	process = java,elastic,elasticsearch
	process = java,logstash,logstash
	process = /opt/bosun/bosun,bosun,
	process = /opt/scollector/scollector,scollector,

	// Windows processes and service monitors use the form "name regex"
	process = ^chrome
	process = ^powershell
	process = ^scollector
	process = ^WinRM
	process = (?i)^MSSQLServer //Matches are case sensitive unless specified

	// Dotnet processes to monitor use the form "process/service name regex"
	process_dotnet=^w3wp //Optional, as IIS processes are always monitored
	process_dotnet=^MyCustomService
	process_dotnet=^powershell

Windows

scollector has full Windows support. It can be run standalone, or installed as a
service (see -winsvc). The Event Log is used when installed as a service.

SNMP

By default, scollector will collect data from the current host and report it to
the TSDB server. scollector has an SNMP mode where it also polls a given host:

	scollector -s community@host[,community@host...]

Poll frequency currently defaults to 5 minutes. Some common OIDs regarding
interfaces are collected. Others can be added easily.

*/
package main
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"sync"
	"time"

	"github.com/Impyy/tox4go/bootstrap"
	"github.com/Impyy/tox4go/crypto"
	"github.com/Impyy/tox4go/dht"
	"github.com/Impyy/tox4go/dht/ping"
	"github.com/Impyy/tox4go/relay"
	"github.com/Impyy/tox4go/transport"
	"github.com/didip/tollbooth"
)

const (
	enableIpv6  = true
	probeRate   = 1 * time.Minute
	refreshRate = 5 * time.Minute
)

var (
	lastScan    int64
	lastRefresh int64
	nodes       = []*toxNode{}
	nodesMutex  = sync.Mutex{}
	tcpPorts    = []int{443, 3389, 33445}
)

func main() {
	if parseFlags() {
		return
	}

	// load state if available
	state, err := loadState()
	if err != nil {
		log.Fatalf("error loading state: %s", err.Error())
	}
	lastScan = state.LastScan
	lastRefresh = state.LastRefresh
	nodes = state.Nodes

	if err := loadCountries(); err != nil {
		log.Fatalf("error loading countries.json: %s", err.Error())
	}

	inst, err := NewInstance(":33450")
	if err != nil {
		log.Fatalf("fatal: %s", err.Error())
	}
	inst.UDPTransport.Handle(dht.PacketIDSendNodes, inst.handleSendNodesPacket)
	inst.UDPTransport.Handle(bootstrap.PacketIDBootstrapInfo, handleBootstrapInfoPacket)

	//handle stop signal
	interruptChan := make(chan os.Signal)
	signal.Notify(interruptChan, os.Interrupt)

	//setup http server
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", httpListenPort))
	if err != nil {
		log.Fatalf("error in net.Listen: %s", err.Error())
	}
	limiter := tollbooth.NewLimiter(1, 2*time.Second)
	limiter.Methods = []string{"POST"}
	limiter.IPLookups = []string{"X-Forwarded-For", "RemoteAddr", "X-Real-IP"}
	serveMux := http.NewServeMux()
	serveMux.HandleFunc("/", handleHTTPRequest)
	serveMux.Handle("/test", tollbooth.LimitFuncHandler(limiter, handleHTTPRequest))
	serveMux.HandleFunc("/json", handleJSONRequest)
	go func() {
		err := http.Serve(listener, serveMux)
		if err != nil {
			log.Printf("http server error: %s\n", err.Error())
			interruptChan <- os.Interrupt
		}
	}()

	//listen for tox packets
	go func() {
		err := inst.UDPTransport.Listen()
		if err != nil {
			log.Printf("udp transport error: %s\n", err.Error())
			interruptChan <- os.Interrupt
		}
	}()
	//go tcpTransport.Listen()

	err = refreshNodes()
	if err != nil {
		log.Fatal(err.Error())
	}
	inst.probeNodes()

	probeTicker := time.NewTicker(probeRate)
	refreshTicker := time.NewTicker(refreshRate)
	updateTicker := time.NewTicker(30 * time.Second)
	run := true

	for run {
		select {
		case <-interruptChan:
			fmt.Printf("killing routines\n")
			probeTicker.Stop()
			refreshTicker.Stop()
			updateTicker.Stop()
			inst.UDPTransport.Stop()
			//tcpTransport.Stop()
			listener.Close()
			run = false
		case <-probeTicker.C:
			// we want an empty ping list at the start of every probe
			inst.PingsMutex.Lock()
			inst.Pings.Clear(false)
			inst.PingsMutex.Unlock()

			nodesMutex.Lock()
			err := inst.probeNodes()
			nodesMutex.Unlock()
			if err != nil {
				log.Printf("error while trying to probe nodes: %s", err.Error())
			}
		case <-refreshTicker.C:
			err := refreshNodes()
			if err != nil {
				log.Printf("error while trying to refresh nodes: %s", err.Error())
			}
		case <-updateTicker.C:
			inst.PingsMutex.Lock()
			inst.Pings.Clear(true)
			inst.PingsMutex.Unlock()

			nodesMutex.Lock()
			for _, node := range nodes {
				if time.Now().Sub(time.Unix(node.LastPing, 0)) > time.Minute*2 {
					node.UDPStatus = false
				}
			}
			sort.Stable(nodeSlice(nodes))

			state := getState()
			saveState(state)
			nodesMutex.Unlock()
		}
	}
}

func refreshNodes() error {
	parsedNodes, err := parseNodes()
	if err != nil {
		return err
	}

	nodesMutex.Lock()
	for _, freshNode := range parsedNodes {
		for i, node := range nodes {
			if freshNode.PublicKey == node.PublicKey {
				freshNode.LastPing = node.LastPing
				freshNode.UDPStatus = node.UDPStatus
				freshNode.TCPStatus = node.TCPStatus
				freshNode.TCPPorts = node.TCPPorts
				freshNode.MOTD = node.MOTD
				freshNode.Version = node.Version
				nodes[i] = freshNode
				break
			}
		}
	}
	nodes = parsedNodes
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()

	lastRefresh = time.Now().Unix()
	return nil
}

func (i *instance) probeNodes() error {
	for _, node := range nodes {
		err := i.getBootstrapInfo(node)
		if err != nil {
			fmt.Println(err.Error())
		}

		p, err := i.getNodes(node)
		if err != nil {
			fmt.Println(err.Error())
		} else {
			i.PingsMutex.Lock()
			err = i.Pings.Add(p)
			if err != nil {
				fmt.Println(err.Error())
			}
			i.PingsMutex.Unlock()
		}

		ports := tcpPorts
		exists := false
		for _, i := range ports {
			if i == node.Port {
				exists = true
			}
		}
		if !exists {
			ports = append(ports, node.Port)
		}

		go i.probeNodeTCPPorts(node, ports)
	}

	lastScan = time.Now().Unix()
	return nil
}

func (i *instance) probeNodeTCPPorts(node *toxNode, ports []int) {
	c := make(chan int)
	for _, port := range ports {
		go func(p int) {
			conn, err := connectTCP(node, p)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
				c <- -1
				return
			}

			err = i.tcpHandshake(node, conn)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
				c <- -1
			} else {
				c <- p
			}
			conn.Close()
		}(port)
	}

	nodesMutex.Lock()
	node.TCPPorts = []int{}

	for i := 0; i < len(ports); i++ {
		port := <-c
		if port != -1 {
			fmt.Printf("tcp port for %s: %d\n", node.Maintainer, port)
			node.TCPPorts = append(node.TCPPorts, port)
		}
	}
	if len(node.TCPPorts) > 0 {
		node.LastPing = time.Now().Unix()
	}
	node.TCPStatus = len(node.TCPPorts) > 0
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()
}

func (i *instance) tcpHandshake(node *toxNode, conn *net.TCPConn) error {
	nodePublicKey := new([crypto.PublicKeySize]byte)
	decPublicKey, err := hex.DecodeString(node.PublicKey)
	if err != nil {
		return err
	}
	copy(nodePublicKey[:], decPublicKey)

	relayConn, err := relay.NewConnection()
	if err != nil {
		return err
	}

	req, err := relayConn.StartHandshake()
	if err != nil {
		return err
	}

	reqBytes, err := req.MarshalBinary()
	if err != nil {
		return err
	}

	encryptedReqBytes, nonce, err := i.Ident.EncryptBlob(reqBytes, nodePublicKey)
	if err != nil {
		return err
	}

	reqPacket := &relay.HandshakeRequestPacket{
		PublicKey: i.Ident.PublicKey,
		Nonce:     nonce,
		Payload:   encryptedReqBytes,
	}

	reqPacketBytes, err := reqPacket.MarshalBinary()
	if err != nil {
		return err
	}

	conn.SetReadDeadline(time.Now().Add(2 * time.Second))
	_, err = conn.Write(reqPacketBytes)
	if err != nil {
		return err
	}

	buffer := make([]byte, 96)
	left := len(buffer)
	for left > 0 {
		read, readErr := conn.Read(buffer[len(buffer)-left:])
		if readErr != nil {
			return readErr
		}
		left -= read
	}

	res := relay.HandshakeResponsePacket{}
	err = res.UnmarshalBinary(buffer)
	if err != nil {
		return err
	}

	decryptedBytes, err := i.Ident.DecryptBlob(res.Payload, nodePublicKey, res.Nonce)
	if err != nil {
		return err
	}

	resPacket := &relay.HandshakePayload{}
	err = resPacket.UnmarshalBinary(decryptedBytes)
	if err != nil {
		return err
	}

	return relayConn.EndHandshake(resPacket)
}

func (i *instance) getNodes(node *toxNode) (*ping.Ping, error) {
	nodePublicKey := new([crypto.PublicKeySize]byte)
	decPublicKey, err := hex.DecodeString(node.PublicKey)
	if err != nil {
		return nil, err
	}
	copy(nodePublicKey[:], decPublicKey)

	p, err := ping.NewPing(nodePublicKey)
	if err != nil {
		return nil, err
	}

	packet := &dht.GetNodesPacket{
		PublicKey: i.Ident.PublicKey,
		PingID:    p.ID,
	}

	dhtPacket, err := i.Ident.EncryptPacket(transport.Packet(packet), nodePublicKey)
	if err != nil {
		return nil, err
	}

	payload, err := dhtPacket.MarshalBinary()
	if err != nil {
		return nil, err
	}

	if err := i.sendToUDP(payload, node); err != nil {
		return nil, err
	}

	return p, nil
}

func (i *instance) getBootstrapInfo(node *toxNode) error {
	packet, err := bootstrap.ConstructPacket(&bootstrap.InfoRequestPacket{})
	if err != nil {
		return err
	}

	payload, err := packet.MarshalBinary()
	if err != nil {
		return err
	}

	return i.sendToUDP(payload, node)
}

func (i *instance) sendToUDP(data []byte, node *toxNode) error {
	ip, err := getNodeIP(node)
	if err != nil {
		return err
	}

	return i.UDPTransport.Send(
		&transport.Message{
			Data: data,
			Addr: &net.UDPAddr{
				IP:   ip,
				Port: node.Port,
			},
		},
	)
}

func getNodeIP(node *toxNode) (net.IP, error) {
	if node.ip4 != nil {
		return node.ip4, nil
	} else if enableIpv6 && node.ip6 != nil {
		return node.ip6, nil
	}

	return nil, fmt.Errorf("no valid ip found for %s", node.Maintainer)
}

func connectTCP(node *toxNode, port int) (*net.TCPConn, error) {
	ip, err := getNodeIP(node)
	if err != nil {
		return nil, err
	}

	dialer := net.Dialer{}
	dialer.Deadline = time.Now().Add(2 * time.Second)

	tempConn, err := dialer.Dial("tcp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		return nil, err
	}

	conn, ok := tempConn.(*net.TCPConn)
	if !ok {
		return nil, errors.New("not a tcp conn")
	}

	return conn, nil
}

func (i *instance) handleSendNodesPacket(msg *transport.Message) error {
	dhtPacket := &dht.Packet{}
	err := dhtPacket.UnmarshalBinary(msg.Data)
	if err != nil {
		return err
	}

	decryptedPacket, err := i.Ident.DecryptPacket(dhtPacket)
	if err != nil {
		return err
	}

	packet, ok := decryptedPacket.(*dht.SendNodesPacket)
	if !ok {
		return nil
	}

	i.PingsMutex.Lock()
	nodesMutex.Lock()
	if i.Pings.Find(dhtPacket.SenderPublicKey, packet.PingID, true) != nil {
		for _, node := range nodes {
			publicKey, err := hex.DecodeString(node.PublicKey)
			if err != nil {
				continue
			}

			if bytes.Equal(publicKey, dhtPacket.SenderPublicKey[:]) {
				node.UDPStatus = true
				node.LastPing = time.Now().Unix()
				break
			}
		}
	}
	sort.Stable(nodeSlice(nodes))
	i.PingsMutex.Unlock()
	nodesMutex.Unlock()

	return nil
}

func handleBootstrapInfoPacket(msg *transport.Message) error {
	bootstrapPacket := &bootstrap.Packet{}
	err := bootstrapPacket.UnmarshalBinary(msg.Data)
	if err != nil {
		return err
	}

	transPacket, err := bootstrap.DestructPacket(bootstrapPacket)
	if err != nil {
		return err
	}

	packet, ok := transPacket.(*bootstrap.InfoResponsePacket)
	if !ok {
		return errors.New("wtf")
	}

	nodesMutex.Lock()
	for _, node := range nodes {
		if (node.ip4 != nil && node.ip4.Equal(msg.Addr.IP)) ||
			(node.ip6 != nil && node.ip6.Equal(msg.Addr.IP)) {
			node.MOTD = packet.MOTD
			node.Version = fmt.Sprintf("%d", packet.Version)
			break
		}
	}
	nodesMutex.Unlock()

	return nil
}
<commit_msg>Remove redundant assignment<commit_after>package main

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"sync"
	"time"

	"github.com/Impyy/tox4go/bootstrap"
	"github.com/Impyy/tox4go/crypto"
	"github.com/Impyy/tox4go/dht"
	"github.com/Impyy/tox4go/dht/ping"
	"github.com/Impyy/tox4go/relay"
	"github.com/Impyy/tox4go/transport"
	"github.com/didip/tollbooth"
)

const (
	enableIpv6  = true
	probeRate   = 1 * time.Minute
	refreshRate = 5 * time.Minute
)

var (
	lastScan    int64
	lastRefresh int64
	nodes       = []*toxNode{}
	nodesMutex  = sync.Mutex{}
	tcpPorts    = []int{443, 3389, 33445}
)

func main() {
	if parseFlags() {
		return
	}

	// load state if available
	state, err := loadState()
	if err != nil {
		log.Fatalf("error loading state: %s", err.Error())
	}
	lastScan = state.LastScan
	lastRefresh = state.LastRefresh
	nodes = state.Nodes

	if err := loadCountries(); err != nil {
		log.Fatalf("error loading countries.json: %s", err.Error())
	}

	inst, err := NewInstance(":33450")
	if err != nil {
		log.Fatalf("fatal: %s", err.Error())
	}
	inst.UDPTransport.Handle(dht.PacketIDSendNodes, inst.handleSendNodesPacket)
	inst.UDPTransport.Handle(bootstrap.PacketIDBootstrapInfo, handleBootstrapInfoPacket)

	//handle stop signal
	interruptChan := make(chan os.Signal)
	signal.Notify(interruptChan, os.Interrupt)

	//setup http server
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", httpListenPort))
	if err != nil {
		log.Fatalf("error in net.Listen: %s", err.Error())
	}
	limiter := tollbooth.NewLimiter(1, 2*time.Second)
	limiter.Methods = []string{"POST"}
	limiter.IPLookups = []string{"X-Forwarded-For", "RemoteAddr", "X-Real-IP"}
	serveMux := http.NewServeMux()
	serveMux.HandleFunc("/", handleHTTPRequest)
	serveMux.Handle("/test", tollbooth.LimitFuncHandler(limiter, handleHTTPRequest))
	serveMux.HandleFunc("/json", handleJSONRequest)
	go func() {
		err := http.Serve(listener, serveMux)
		if err != nil {
			log.Printf("http server error: %s\n", err.Error())
			interruptChan <- os.Interrupt
		}
	}()

	//listen for tox packets
	go func() {
		err := inst.UDPTransport.Listen()
		if err != nil {
			log.Printf("udp transport error: %s\n", err.Error())
			interruptChan <- os.Interrupt
		}
	}()
	//go tcpTransport.Listen()

	err = refreshNodes()
	if err != nil {
		log.Fatal(err.Error())
	}
	inst.probeNodes()

	probeTicker := time.NewTicker(probeRate)
	refreshTicker := time.NewTicker(refreshRate)
	updateTicker := time.NewTicker(30 * time.Second)
	run := true

	for run {
		select {
		case <-interruptChan:
			fmt.Printf("killing routines\n")
			probeTicker.Stop()
			refreshTicker.Stop()
			updateTicker.Stop()
			inst.UDPTransport.Stop()
			//tcpTransport.Stop()
			listener.Close()
			run = false
		case <-probeTicker.C:
			// we want an empty ping list at the start of every probe
			inst.PingsMutex.Lock()
			inst.Pings.Clear(false)
			inst.PingsMutex.Unlock()

			nodesMutex.Lock()
			err := inst.probeNodes()
			nodesMutex.Unlock()
			if err != nil {
				log.Printf("error while trying to probe nodes: %s", err.Error())
			}
		case <-refreshTicker.C:
			err := refreshNodes()
			if err != nil {
				log.Printf("error while trying to refresh nodes: %s", err.Error())
			}
		case <-updateTicker.C:
			inst.PingsMutex.Lock()
			inst.Pings.Clear(true)
			inst.PingsMutex.Unlock()

			nodesMutex.Lock()
			for _, node := range nodes {
				if time.Now().Sub(time.Unix(node.LastPing, 0)) > time.Minute*2 {
					node.UDPStatus = false
				}
			}
			sort.Stable(nodeSlice(nodes))

			state := getState()
			saveState(state)
			nodesMutex.Unlock()
		}
	}
}

func refreshNodes() error {
	parsedNodes, err := parseNodes()
	if err != nil {
		return err
	}

	nodesMutex.Lock()
	for _, freshNode := range parsedNodes {
		for _, node := range nodes {
			if freshNode.PublicKey == node.PublicKey {
				freshNode.LastPing = node.LastPing
				freshNode.UDPStatus = node.UDPStatus
				freshNode.TCPStatus = node.TCPStatus
				freshNode.TCPPorts = node.TCPPorts
				freshNode.MOTD = node.MOTD
				freshNode.Version = node.Version
				break
			}
		}
	}
	nodes = parsedNodes
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()

	lastRefresh = time.Now().Unix()
	return nil
}

func (i *instance) probeNodes() error {
	for _, node := range nodes {
		err := i.getBootstrapInfo(node)
		if err != nil {
			fmt.Println(err.Error())
		}

		p, err := i.getNodes(node)
		if err != nil {
			fmt.Println(err.Error())
		} else {
			i.PingsMutex.Lock()
			err = i.Pings.Add(p)
			if err != nil {
				fmt.Println(err.Error())
			}
			i.PingsMutex.Unlock()
		}

		ports := tcpPorts
		exists := false
		for _, i := range ports {
			if i == node.Port {
				exists = true
			}
		}
		if !exists {
			ports = append(ports, node.Port)
		}

		go i.probeNodeTCPPorts(node, ports)
	}

	lastScan = time.Now().Unix()
	return nil
}

func (i *instance) probeNodeTCPPorts(node *toxNode, ports []int) {
	c := make(chan int)
	for _, port := range ports {
		go func(p int) {
			conn, err := connectTCP(node, p)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
				c <- -1
				return
			}

			err = i.tcpHandshake(node, conn)
			if err != nil {
				fmt.Printf("%s\n", err.Error())
				c <- -1
			} else {
				c <- p
			}
			conn.Close()
		}(port)
	}

	nodesMutex.Lock()
	node.TCPPorts = []int{}

	for i := 0; i < len(ports); i++ {
		port := <-c
		if port != -1 {
			fmt.Printf("tcp port for %s: %d\n", node.Maintainer, port)
			node.TCPPorts = append(node.TCPPorts, port)
		}
	}
	if len(node.TCPPorts) > 0 {
		node.LastPing = time.Now().Unix()
	}
	node.TCPStatus = len(node.TCPPorts) > 0
	sort.Stable(nodeSlice(nodes))
	nodesMutex.Unlock()
}

func (i *instance) tcpHandshake(node *toxNode, conn *net.TCPConn) error {
	nodePublicKey := new([crypto.PublicKeySize]byte)
	decPublicKey, err := hex.DecodeString(node.PublicKey)
	if err != nil {
		return err
	}
	copy(nodePublicKey[:], decPublicKey)

	relayConn, err := relay.NewConnection()
	if err != nil {
		return err
	}

	req, err := relayConn.StartHandshake()
	if err != nil {
		return err
	}

	reqBytes, err := req.MarshalBinary()
	if err != nil {
		return err
	}

	encryptedReqBytes, nonce, err := i.Ident.EncryptBlob(reqBytes, nodePublicKey)
	if err != nil {
		return err
	}

	reqPacket := &relay.HandshakeRequestPacket{
		PublicKey: i.Ident.PublicKey,
		Nonce:     nonce,
		Payload:   encryptedReqBytes,
	}

	reqPacketBytes, err := reqPacket.MarshalBinary()
	if err != nil {
		return err
	}

	conn.SetReadDeadline(time.Now().Add(2 * time.Second))
	_, err = conn.Write(reqPacketBytes)
	if err != nil {
		return err
	}

	buffer := make([]byte, 96)
	left := len(buffer)
	for left > 0 {
		read, readErr := conn.Read(buffer[len(buffer)-left:])
		if readErr != nil {
			return readErr
		}
		left -= read
	}

	res := relay.HandshakeResponsePacket{}
	err = res.UnmarshalBinary(buffer)
	if err != nil {
		return err
	}

	decryptedBytes, err := i.Ident.DecryptBlob(res.Payload, nodePublicKey, res.Nonce)
	if err != nil {
		return err
	}

	resPacket := &relay.HandshakePayload{}
	err = resPacket.UnmarshalBinary(decryptedBytes)
	if err != nil {
		return err
	}

	return relayConn.EndHandshake(resPacket)
}

func (i *instance) getNodes(node *toxNode) (*ping.Ping, error) {
	nodePublicKey := new([crypto.PublicKeySize]byte)
	decPublicKey, err := hex.DecodeString(node.PublicKey)
	if err != nil {
		return nil, err
	}
	copy(nodePublicKey[:], decPublicKey)

	p, err := ping.NewPing(nodePublicKey)
	if err != nil {
		return nil, err
	}

	packet := &dht.GetNodesPacket{
		PublicKey: i.Ident.PublicKey,
		PingID:    p.ID,
	}

	dhtPacket, err := i.Ident.EncryptPacket(transport.Packet(packet), nodePublicKey)
	if err != nil {
		return nil, err
	}

	payload, err := dhtPacket.MarshalBinary()
	if err != nil {
		return nil, err
	}

	if err := i.sendToUDP(payload, node); err != nil {
		return nil, err
	}

	return p, nil
}

func (i *instance) getBootstrapInfo(node *toxNode) error {
	packet, err := bootstrap.ConstructPacket(&bootstrap.InfoRequestPacket{})
	if err != nil {
		return err
	}

	payload, err := packet.MarshalBinary()
	if err != nil {
		return err
	}

	return i.sendToUDP(payload, node)
}

func (i *instance) sendToUDP(data []byte, node *toxNode) error {
	ip, err := getNodeIP(node)
	if err != nil {
		return err
	}

	return i.UDPTransport.Send(
		&transport.Message{
			Data: data,
			Addr: &net.UDPAddr{
				IP:   ip,
				Port: node.Port,
			},
		},
	)
}

func getNodeIP(node *toxNode) (net.IP, error) {
	if node.ip4 != nil {
		return node.ip4, nil
	} else if enableIpv6 && node.ip6 != nil {
		return node.ip6, nil
	}

	return nil, fmt.Errorf("no valid ip found for %s", node.Maintainer)
}

func connectTCP(node *toxNode, port int) (*net.TCPConn, error) {
	ip, err := getNodeIP(node)
	if err != nil {
		return nil, err
	}

	dialer := net.Dialer{}
	dialer.Deadline = time.Now().Add(2 * time.Second)

	tempConn, err := dialer.Dial("tcp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		return nil, err
	}

	conn, ok := tempConn.(*net.TCPConn)
	if !ok {
		return nil, errors.New("not a tcp conn")
	}

	return conn, nil
}

func (i *instance) handleSendNodesPacket(msg *transport.Message) error {
	dhtPacket := &dht.Packet{}
	err := dhtPacket.UnmarshalBinary(msg.Data)
	if err != nil {
		return err
	}

	decryptedPacket, err := i.Ident.DecryptPacket(dhtPacket)
	if err != nil {
		return err
	}

	packet, ok := decryptedPacket.(*dht.SendNodesPacket)
	if !ok {
		return nil
	}

	i.PingsMutex.Lock()
	nodesMutex.Lock()
	if i.Pings.Find(dhtPacket.SenderPublicKey, packet.PingID, true) != nil {
		for _, node := range nodes {
			publicKey, err := hex.DecodeString(node.PublicKey)
			if err != nil {
				continue
			}

			if bytes.Equal(publicKey, dhtPacket.SenderPublicKey[:]) {
				node.UDPStatus = true
				node.LastPing = time.Now().Unix()
				break
			}
		}
	}
	sort.Stable(nodeSlice(nodes))
	i.PingsMutex.Unlock()
	nodesMutex.Unlock()

	return nil
}

func handleBootstrapInfoPacket(msg *transport.Message) error {
	bootstrapPacket := &bootstrap.Packet{}
	err := bootstrapPacket.UnmarshalBinary(msg.Data)
	if err != nil {
		return err
	}

	transPacket, err := bootstrap.DestructPacket(bootstrapPacket)
	if err != nil {
		return err
	}

	packet, ok := transPacket.(*bootstrap.InfoResponsePacket)
	if !ok {
		return errors.New("wtf")
	}

	nodesMutex.Lock()
	for _, node := range nodes {
		if (node.ip4 != nil && node.ip4.Equal(msg.Addr.IP)) ||
			(node.ip6 != nil && node.ip6.Equal(msg.Addr.IP)) {
			node.MOTD = packet.MOTD
			node.Version = fmt.Sprintf("%d", packet.Version)
			break
		}
	}
	nodesMutex.Unlock()

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 trivago GmbH
//
// 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 router

import (
	"github.com/trivago/gollum/core"
	"sync"
	"sync/atomic"
)

// RoundRobin router plugin
//
// Messages will be sent to one of the producers attached to this router.
// Producers will be switched one-by-one.
//
// The "RoundRobin" router relays each message sent to the stream [Stream] to
// exactly one of the producers connected to [Stream]. The producer is selected
// by rotating the connected producers in sequence, one producer per received
// message.
//
// Configuration example:
//
// # Generate junk
// JunkGenerator:
//   Type: "consumer.Profiler"
//   Message: "%20s"
//   Streams: "junkstream"
//   Characters: "abcdefghijklmZ"
//   KeepRunning: true
//   Runs: 10000
//   Batches: 3000000
//   DelayMs: 500
// # Spread messages to connected producers in round-robin
// JunkRouterRoundRob:
//   Type: "router.RoundRobin"
//   Stream: "junkstream"
// # Produce messages to stdout
// JunkPrinter00:
//   Type: "producer.Console"
//   Streams: "junkstream"
//   Modulators:
//     - "format.Envelope":
//         Prefix: "[junk_00] "
// # Produce messages to stdout
// JunkPrinter01:
//   Type: "producer.Console"
//   Streams: "junkstream"
//   Modulators:
//     - "format.Envelope":
//         Prefix: "[junk_01] "
//
type RoundRobin struct {
	core.SimpleRouter `gollumdoc:"embed_type"`
	index             int32
	indexByStream     map[core.MessageStreamID]*int32
	mapInitLock       *sync.Mutex
}

func init() {
	core.TypeRegistry.Register(RoundRobin{})
}

// Configure initializes this distributor with values from a plugin config.
func (router *RoundRobin) Configure(conf core.PluginConfigReader) {
	router.index = 0
	router.indexByStream = make(map[core.MessageStreamID]*int32)
	router.mapInitLock = new(sync.Mutex)
}

// Start the router
func (router *RoundRobin) Start() error {
	return nil
}

// Enqueue enques a message to the router
func (router *RoundRobin) Enqueue(msg *core.Message) error {
	producers := router.GetProducers()
	if len(producers) == 0 {
		return core.NewModulateResultError("No producers configured for stream %s", router.GetID())
	}
	index := atomic.AddInt32(&router.index, 1) % int32(len(producers))
	producers[index].Enqueue(msg, router.GetTimeout())
	return nil
}
<commit_msg>update plugin docs for router.RoundRobin<commit_after>// Copyright 2015-2017 trivago GmbH
//
// 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 router

import (
	"github.com/trivago/gollum/core"
	"sync"
	"sync/atomic"
)

// RoundRobin router
//
// This router implements round robin routing. Messages will always be routed to
// only and exactly one of the producers registered to the given stream. The
// producer is switched in a round robin fashin after each message.
// This producer can be useful for load balancing, e.g. when the target service
// does not support sharding by itself.
//
// Examples
//
//  loadBalancer:
//    Type: router.RoundRobin
//    Stream: logs
type RoundRobin struct {
	core.SimpleRouter `gollumdoc:"embed_type"`
	index             int32
	indexByStream     map[core.MessageStreamID]*int32
	mapInitLock       *sync.Mutex
}

func init() {
	core.TypeRegistry.Register(RoundRobin{})
}

// Configure initializes this distributor with values from a plugin config.
func (router *RoundRobin) Configure(conf core.PluginConfigReader) {
	router.index = 0
	router.indexByStream = make(map[core.MessageStreamID]*int32)
	router.mapInitLock = new(sync.Mutex)
}

// Start the router
func (router *RoundRobin) Start() error {
	return nil
}

// Enqueue enques a message to the router
func (router *RoundRobin) Enqueue(msg *core.Message) error {
	producers := router.GetProducers()
	if len(producers) == 0 {
		return core.NewModulateResultError("No producers configured for stream %s", router.GetID())
	}
	index := atomic.AddInt32(&router.index, 1) % int32(len(producers))
	producers[index].Enqueue(msg, router.GetTimeout())
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Copyright 2014 Ninja Blocks Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package json2

import (
	"encoding/json"

	"github.com/ninjasphere/go-ninja/rpc3"
)

// ----------------------------------------------------------------------------
// Request and Response
// ----------------------------------------------------------------------------

// clientRequest represents a JSON-RPC request sent by a client.
type clientRequest struct {
	// JSON-RPC protocol.
	Version string `json:"jsonrpc"`

	// A String containing the name of the method to be invoked.
	Method string `json:"method"`

	// Object to pass as request parameter to the method.
	Params interface{} `json:"params"`

	// The request id. This can be of any type. It is used to match the
	// response with the request that it is replying to.
	Id uint32 `json:"id"`
}

// clientResponse represents a JSON-RPC response returned to a client.
type clientResponse struct {
	Version string           `json:"jsonrpc"`
	Result  *json.RawMessage `json:"response"`
	Error   *json.RawMessage `json:"error"`
	Id      uint32           `json:"id"`
}

func NewClientCodec() *ClientCodec {
	return &ClientCodec{}
}

type ClientCodec struct {
}

// EncodeClientRequest encodes parameters for a JSON-RPC client request.
func (c *ClientCodec) EncodeClientRequest(call *rpc.Call) ([]byte, error) {
	req := &clientRequest{
		Version: "2.0",
		Method:  call.ServiceMethod,
		Params:  call.Args,
		Id:      call.Id,
	}
	return json.Marshal(req)
}

func (c *ClientCodec) DecodeIdAndError(msg []byte) (*uint32, error) {
	res := &clientResponse{}

	if err := json.Unmarshal(msg, res); err != nil {
		return nil, err
	}

	if res.Error != nil {
		jsonErr := &Error{}
		if err := json.Unmarshal(*res.Error, jsonErr); err != nil {
			return &res.Id, &Error{
				Code:    E_SERVER,
				Message: string(*res.Error),
			}
		}
		return &res.Id, jsonErr
	}

	return &res.Id, nil

}

// DecodeClientResponse decodes the response body of a client request into
// the interface reply.
func (c *ClientCodec) DecodeClientResponse(msg []byte, reply interface{}) error {
	var res clientResponse
	if err := json.Unmarshal(msg, &res); err != nil {
		return err
	}
	return json.Unmarshal(*res.Result, reply)
}
<commit_msg>Force nil args in roc client to send as an empty array<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Copyright 2014 Ninja Blocks Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package json2

import (
	"encoding/json"

	"github.com/ninjasphere/go-ninja/rpc3"
)

// ----------------------------------------------------------------------------
// Request and Response
// ----------------------------------------------------------------------------

// clientRequest represents a JSON-RPC request sent by a client.
type clientRequest struct {
	// JSON-RPC protocol.
	Version string `json:"jsonrpc"`

	// A String containing the name of the method to be invoked.
	Method string `json:"method"`

	// Object to pass as request parameter to the method.
	Params interface{} `json:"params"`

	// The request id. This can be of any type. It is used to match the
	// response with the request that it is replying to.
	Id uint32 `json:"id"`
}

// clientResponse represents a JSON-RPC response returned to a client.
type clientResponse struct {
	Version string           `json:"jsonrpc"`
	Result  *json.RawMessage `json:"response"`
	Error   *json.RawMessage `json:"error"`
	Id      uint32           `json:"id"`
}

func NewClientCodec() *ClientCodec {
	return &ClientCodec{}
}

type ClientCodec struct {
}

// EncodeClientRequest encodes parameters for a JSON-RPC client request.
func (c *ClientCodec) EncodeClientRequest(call *rpc.Call) ([]byte, error) {
	req := &clientRequest{
		Version: "2.0",
		Method:  call.ServiceMethod,
		Params:  []interface{}{},
		Id:      call.Id,
	}

	if call.Args != nil {
		req.Params = call.Args
	}

	return json.Marshal(req)
}

func (c *ClientCodec) DecodeIdAndError(msg []byte) (*uint32, error) {
	res := &clientResponse{}

	if err := json.Unmarshal(msg, res); err != nil {
		return nil, err
	}

	if res.Error != nil {
		jsonErr := &Error{}
		if err := json.Unmarshal(*res.Error, jsonErr); err != nil {
			return &res.Id, &Error{
				Code:    E_SERVER,
				Message: string(*res.Error),
			}
		}
		return &res.Id, jsonErr
	}

	return &res.Id, nil

}

// DecodeClientResponse decodes the response body of a client request into
// the interface reply.
func (c *ClientCodec) DecodeClientResponse(msg []byte, reply interface{}) error {
	var res clientResponse
	if err := json.Unmarshal(msg, &res); err != nil {
		return err
	}
	return json.Unmarshal(*res.Result, reply)
}
<|endoftext|>


<commit_before>package server

import (
	"testing"

	"github.com/pachyderm/pachyderm/v2/src/internal/require"
	tu "github.com/pachyderm/pachyderm/v2/src/internal/testutil"
	"github.com/pachyderm/pachyderm/v2/src/pps"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func wrap(t testing.TB, ps *pps.ParallelismSpec) *pps.PipelineInfo {
	return &pps.PipelineInfo{
		Pipeline: &pps.Pipeline{
			Name: t.Name() + "-pipeline",
		},
		ParallelismSpec: ps,
	}
}

func TestGetExpectedNumWorkers(t *testing.T) {
	kubeClient := tu.GetKubeClient(t)

	// An empty parallelism spec should default to 1 worker
	workers, err := getExpectedNumWorkers(kubeClient, wrap(t,
		&pps.ParallelismSpec{}))
	require.NoError(t, err)
	require.Equal(t, 1, workers)

	// A constant should literally be returned
	workers, err = getExpectedNumWorkers(kubeClient, wrap(t,
		&pps.ParallelismSpec{
			Constant: 1,
		}))
	require.NoError(t, err)
	require.Equal(t, 1, workers)
	workers, err = getExpectedNumWorkers(kubeClient, wrap(t,
		&pps.ParallelismSpec{
			Constant: 3,
		}))
	require.NoError(t, err)
	require.Equal(t, 3, workers)

	// Constant and Coefficient cannot both be non-zero
	_, err = getExpectedNumWorkers(kubeClient, wrap(t,
		&pps.ParallelismSpec{
			Constant:    3,
			Coefficient: 0.5,
		}))
	require.YesError(t, err)

	// No parallelism spec should default to 1 worker
	workers, err = getExpectedNumWorkers(kubeClient, wrap(t, nil))
	require.NoError(t, err)
	require.Equal(t, 1, workers)

	nodes, err := kubeClient.CoreV1().Nodes().List(metav1.ListOptions{})
	require.NoError(t, err)
	numNodes := len(nodes.Items)

	// Coefficient == 1
	parellelism, err := getExpectedNumWorkers(kubeClient, wrap(t,
		&pps.ParallelismSpec{
			Coefficient: 1,
		}))
	require.NoError(t, err)
	require.Equal(t, numNodes, parellelism)

	// Coefficient > 1
	parellelism, err = getExpectedNumWorkers(kubeClient, wrap(t,
		&pps.ParallelismSpec{
			Coefficient: 2,
		}))
	require.NoError(t, err)
	require.Equal(t, 2*numNodes, parellelism)

	// Make sure we start at least one worker
	parellelism, err = getExpectedNumWorkers(kubeClient, wrap(t,
		&pps.ParallelismSpec{
			Coefficient: 0.01,
		}))
	require.NoError(t, err)
	require.Equal(t, 1, parellelism)
}
<commit_msg>Remove testing for coefficient.<commit_after>package server

import (
	"testing"

	"github.com/pachyderm/pachyderm/v2/src/internal/require"
	tu "github.com/pachyderm/pachyderm/v2/src/internal/testutil"
	"github.com/pachyderm/pachyderm/v2/src/pps"
)

func wrap(t testing.TB, ps *pps.ParallelismSpec) *pps.PipelineInfo {
	return &pps.PipelineInfo{
		Pipeline: &pps.Pipeline{
			Name: t.Name() + "-pipeline",
		},
		ParallelismSpec: ps,
	}
}

func TestGetExpectedNumWorkers(t *testing.T) {
	kubeClient := tu.GetKubeClient(t)

	// An empty parallelism spec should default to 1 worker
	workers, err := getExpectedNumWorkers(kubeClient, wrap(t,
		&pps.ParallelismSpec{}))
	require.NoError(t, err)
	require.Equal(t, 1, workers)

	// A constant should literally be returned
	workers, err = getExpectedNumWorkers(kubeClient, wrap(t,
		&pps.ParallelismSpec{
			Constant: 1,
		}))
	require.NoError(t, err)
	require.Equal(t, 1, workers)
	workers, err = getExpectedNumWorkers(kubeClient, wrap(t,
		&pps.ParallelismSpec{
			Constant: 3,
		}))
	require.NoError(t, err)
	require.Equal(t, 3, workers)

	// No parallelism spec should default to 1 worker
	workers, err = getExpectedNumWorkers(kubeClient, wrap(t, nil))
	require.NoError(t, err)
	require.Equal(t, 1, workers)
}
<|endoftext|>


<commit_before>package agent

import (
	"fmt"
	"github.com/armon/mdns"
	"io"
	"log"
	"net"
	"time"
)

const (
	mdnsPollInterval  = 60 * time.Second
	mdnsQuietInterval = 100 * time.Millisecond
)

// AgentMDNS is used to advertise ourself using mDNS and to
// attempt to join peers periodically using mDNS queries.
type AgentMDNS struct {
	agent    *Agent
	discover string
	logger   *log.Logger
	seen     map[string]struct{}
	server   *mdns.Server
	replay   bool
	iface    *net.Interface
}

// NewAgentMDNS is used to create a new AgentMDNS
func NewAgentMDNS(agent *Agent, logOutput io.Writer, replay bool,
	node, discover string, iface *net.Interface, bind net.IP, port int) (*AgentMDNS, error) {
	// Create the service
	service := &mdns.MDNSService{
		Instance: node,
		Service:  mdnsName(discover),
		Addr:     bind,
		Port:     port,
		Info:     fmt.Sprintf("Serf '%s' cluster", discover),
	}
	if err := service.Init(); err != nil {
		return nil, err
	}

	// Configure mdns server
	conf := &mdns.Config{
		Zone:  service,
		Iface: iface,
	}

	// Create the server
	server, err := mdns.NewServer(conf)
	if err != nil {
		return nil, err
	}

	// Initialize the AgentMDNS
	m := &AgentMDNS{
		agent:    agent,
		discover: discover,
		logger:   log.New(logOutput, "", log.LstdFlags),
		seen:     make(map[string]struct{}),
		server:   server,
		replay:   replay,
		iface:    iface,
	}

	// Start the background workers
	go m.run()
	return m, nil
}

// run is a long running goroutine that scans for new hosts periodically
func (m *AgentMDNS) run() {
	hosts := make(chan *mdns.ServiceEntry, 32)
	poll := time.After(0)
	var quiet <-chan time.Time
	var join []string

	for {
		select {
		case h := <-hosts:
			// Format the host address
			addr := net.TCPAddr{IP: h.Addr, Port: h.Port}
			addrS := addr.String()

			// Skip if we've handled this host already
			if _, ok := m.seen[addrS]; ok {
				continue
			}

			// Queue for handling
			join = append(join, addrS)
			quiet = time.After(mdnsQuietInterval)

		case <-quiet:
			// Attempt the join
			n, err := m.agent.Join(join, m.replay)
			if err != nil {
				m.logger.Printf("[ERR] agent.mdns: Failed to join: %v", err)
			}
			if n > 0 {
				m.logger.Printf("[INFO] agent.mdns: Joined %d hosts", n)
			}

			// Mark all as seen
			for _, n := range join {
				m.seen[n] = struct{}{}
			}
			join = nil

		case <-poll:
			poll = time.After(mdnsPollInterval)
			go m.poll(hosts)
		}
	}
}

// poll is invoked periodically to check for new hosts
func (m *AgentMDNS) poll(hosts chan *mdns.ServiceEntry) {
	params := mdns.QueryParam{
		Service:   mdnsName(m.discover),
		Interface: m.iface,
		Entries:   hosts,
	}
	if err := mdns.Query(&params); err != nil {
		m.logger.Printf("[ERR] agent.mdns: Failed to poll for new hosts: %v", err)
	}
}

// mdnsName returns the service name to register and to lookup
func mdnsName(discover string) string {
	return fmt.Sprintf("_serf_%s._tcp", discover)
}
<commit_msg>Changing the package import for mdns<commit_after>package agent

import (
	"fmt"
	"io"
	"log"
	"net"
	"time"

	"github.com/hashicorp/mdns"
)

const (
	mdnsPollInterval  = 60 * time.Second
	mdnsQuietInterval = 100 * time.Millisecond
)

// AgentMDNS is used to advertise ourself using mDNS and to
// attempt to join peers periodically using mDNS queries.
type AgentMDNS struct {
	agent    *Agent
	discover string
	logger   *log.Logger
	seen     map[string]struct{}
	server   *mdns.Server
	replay   bool
	iface    *net.Interface
}

// NewAgentMDNS is used to create a new AgentMDNS
func NewAgentMDNS(agent *Agent, logOutput io.Writer, replay bool,
	node, discover string, iface *net.Interface, bind net.IP, port int) (*AgentMDNS, error) {
	// Create the service
	service := &mdns.MDNSService{
		Instance: node,
		Service:  mdnsName(discover),
		Addr:     bind,
		Port:     port,
		Info:     fmt.Sprintf("Serf '%s' cluster", discover),
	}
	if err := service.Init(); err != nil {
		return nil, err
	}

	// Configure mdns server
	conf := &mdns.Config{
		Zone:  service,
		Iface: iface,
	}

	// Create the server
	server, err := mdns.NewServer(conf)
	if err != nil {
		return nil, err
	}

	// Initialize the AgentMDNS
	m := &AgentMDNS{
		agent:    agent,
		discover: discover,
		logger:   log.New(logOutput, "", log.LstdFlags),
		seen:     make(map[string]struct{}),
		server:   server,
		replay:   replay,
		iface:    iface,
	}

	// Start the background workers
	go m.run()
	return m, nil
}

// run is a long running goroutine that scans for new hosts periodically
func (m *AgentMDNS) run() {
	hosts := make(chan *mdns.ServiceEntry, 32)
	poll := time.After(0)
	var quiet <-chan time.Time
	var join []string

	for {
		select {
		case h := <-hosts:
			// Format the host address
			addr := net.TCPAddr{IP: h.Addr, Port: h.Port}
			addrS := addr.String()

			// Skip if we've handled this host already
			if _, ok := m.seen[addrS]; ok {
				continue
			}

			// Queue for handling
			join = append(join, addrS)
			quiet = time.After(mdnsQuietInterval)

		case <-quiet:
			// Attempt the join
			n, err := m.agent.Join(join, m.replay)
			if err != nil {
				m.logger.Printf("[ERR] agent.mdns: Failed to join: %v", err)
			}
			if n > 0 {
				m.logger.Printf("[INFO] agent.mdns: Joined %d hosts", n)
			}

			// Mark all as seen
			for _, n := range join {
				m.seen[n] = struct{}{}
			}
			join = nil

		case <-poll:
			poll = time.After(mdnsPollInterval)
			go m.poll(hosts)
		}
	}
}

// poll is invoked periodically to check for new hosts
func (m *AgentMDNS) poll(hosts chan *mdns.ServiceEntry) {
	params := mdns.QueryParam{
		Service:   mdnsName(m.discover),
		Interface: m.iface,
		Entries:   hosts,
	}
	if err := mdns.Query(&params); err != nil {
		m.logger.Printf("[ERR] agent.mdns: Failed to poll for new hosts: %v", err)
	}
}

// mdnsName returns the service name to register and to lookup
func mdnsName(discover string) string {
	return fmt.Sprintf("_serf_%s._tcp", discover)
}
<|endoftext|>


<commit_before>package cache

import (
	"go.polydawn.net/go-timeless-api"
	"go.polydawn.net/rio/fs"
)

type Cache interface {
	ShelfFor(wareID api.WareID) fs.RelPath
}
<commit_msg>cache: shared behavior is this simple.<commit_after>package cache

import (
	"fmt"

	"go.polydawn.net/go-timeless-api"
	"go.polydawn.net/rio/fs"
	whutil "go.polydawn.net/rio/warehouse/util"
)

func ShelfFor(wareID api.WareID) fs.RelPath {
	chunk1, chunk2, _ := whutil.ChunkifyHash(wareID)
	return fs.MustRelPath(fmt.Sprintf("%s/committed/%s/%s/%s",
		wareID.Type,
		chunk1, chunk2, wareID.Hash,
	))
}
<|endoftext|>


<commit_before>package commands

import (
	"github.com/spf13/cobra"
)

const (
	VERSION = "0.1.0"
)

var Location string

var PunchCardCmd = &cobra.Command{
	Use:   "punchcard",
	Short: "Punchcard is a fun tool to create fake git commits.",
	Long: `Punchcard can create fake git commits in a repo.
The larger purposes is to have fun with contribution graphs, punchcards etc.`,
	Run: nil,
}

func init() {
	PunchCardCmd.PersistentFlags().StringVar(&Location, "location", ".",
		"location where the git repo will be initialized")
}
<commit_msg>Fix typo and use variable for default location.<commit_after>package commands

import (
	"github.com/spf13/cobra"
)

const (
	VERSION          = "0.1.0"
	DEFAULT_LOCATION = "."
)

var Location string

var PunchCardCmd = &cobra.Command{
	Use:   "punchcard",
	Short: "Punchcard is a fun tool to create fake git commits.",
	Long: `Punchcard can create fake git commits in a repo.
The larger purpose is to have fun with contribution graphs, punchcards etc.`,
	Run: nil,
}

func init() {
	PunchCardCmd.PersistentFlags().StringVar(&Location, "location", DEFAULT_LOCATION,
		"location where the git repo will be initialized")
}
<|endoftext|>


<commit_before>package inbound

import (
	"context"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/log"
	v2net "v2ray.com/core/common/net"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/internet/tcp"
	"v2ray.com/core/transport/internet/udp"
)

type worker interface {
	Start() error
	Close()
	Port() v2net.Port
	Proxy() proxy.Inbound
}

type tcpWorker struct {
	address          v2net.Address
	port             v2net.Port
	proxy            proxy.Inbound
	stream           *internet.StreamConfig
	recvOrigDest     bool
	tag              string
	allowPassiveConn bool

	ctx    context.Context
	cancel context.CancelFunc
	hub    *internet.TCPHub
}

func (w *tcpWorker) callback(conn internet.Connection) {
	ctx, cancel := context.WithCancel(w.ctx)
	if w.recvOrigDest {
		dest := tcp.GetOriginalDestination(conn)
		if dest.IsValid() {
			ctx = proxy.ContextWithOriginalDestination(ctx, dest)
		}
	}
	if len(w.tag) > 0 {
		ctx = proxy.ContextWithInboundTag(ctx, w.tag)
	}
	ctx = proxy.ContextWithAllowPassiveConnection(ctx, w.allowPassiveConn)
	ctx = proxy.ContextWithInboundDestination(ctx, v2net.TCPDestination(w.address, w.port))
	w.proxy.Process(ctx, v2net.Network_TCP, conn)
	cancel()
	conn.Close()
}

func (w *tcpWorker) Proxy() proxy.Inbound {
	return w.proxy
}

func (w *tcpWorker) Start() error {
	ctx, cancel := context.WithCancel(context.Background())
	w.ctx = ctx
	w.cancel = cancel
	hub, err := internet.ListenTCP(w.address, w.port, w.callback, w.stream)
	if err != nil {
		return err
	}
	w.hub = hub
	return nil
}

func (w *tcpWorker) Close() {
	log.Debug("Proxyman|TCPWorker: Closed. ", w.port)
	w.hub.Close()
	w.cancel()
}

func (w *tcpWorker) Port() v2net.Port {
	return w.port
}

type udpConn struct {
	cancel           context.CancelFunc
	lastActivityTime int64 // in seconds
	input            chan []byte
	output           func([]byte) (int, error)
	closer           func() error
	remote           net.Addr
	local            net.Addr
}

func (c *udpConn) updateActivity() {
	atomic.StoreInt64(&c.lastActivityTime, time.Now().Unix())
}

func (c *udpConn) Read(buf []byte) (int, error) {
	in, open := <-c.input
	if !open {
		return 0, io.EOF
	}
	c.updateActivity()
	return copy(buf, in), nil
}

func (c *udpConn) Write(buf []byte) (int, error) {
	n, err := c.output(buf)
	if err == nil {
		c.updateActivity()
	}
	return n, err
}

func (c *udpConn) Close() error {
	close(c.input)
	c.cancel()
	return nil
}

func (c *udpConn) RemoteAddr() net.Addr {
	return c.remote
}

func (c *udpConn) LocalAddr() net.Addr {
	return c.remote
}

func (*udpConn) SetDeadline(time.Time) error {
	return nil
}

func (*udpConn) SetReadDeadline(time.Time) error {
	return nil
}

func (*udpConn) SetWriteDeadline(time.Time) error {
	return nil
}

func (*udpConn) Reusable() bool {
	return false
}

func (*udpConn) SetReusable(bool) {}

type udpWorker struct {
	sync.RWMutex

	proxy        proxy.Inbound
	hub          *udp.Hub
	address      v2net.Address
	port         v2net.Port
	recvOrigDest bool
	tag          string

	ctx        context.Context
	cancel     context.CancelFunc
	activeConn map[v2net.Destination]*udpConn
}

func (w *udpWorker) getConnection(src v2net.Destination) (*udpConn, bool) {
	w.Lock()
	defer w.Unlock()

	if conn, found := w.activeConn[src]; found {
		return conn, true
	}

	conn := &udpConn{
		input: make(chan []byte, 32),
		output: func(b []byte) (int, error) {
			return w.hub.WriteTo(b, src)
		},
		closer: func() error {
			w.Lock()
			delete(w.activeConn, src)
			w.Unlock()
			return nil
		},
		remote: &net.UDPAddr{
			IP:   src.Address.IP(),
			Port: int(src.Port),
		},
		local: &net.UDPAddr{
			IP:   w.address.IP(),
			Port: int(w.port),
		},
	}
	w.activeConn[src] = conn

	conn.updateActivity()
	return conn, false
}

func (w *udpWorker) callback(b *buf.Buffer, source v2net.Destination, originalDest v2net.Destination) {
	conn, existing := w.getConnection(source)
	conn.input <- b.Bytes()

	if !existing {
		go func() {
			ctx := w.ctx
			ctx, cancel := context.WithCancel(ctx)
			conn.cancel = cancel
			if originalDest.IsValid() {
				ctx = proxy.ContextWithOriginalDestination(ctx, originalDest)
			}
			if len(w.tag) > 0 {
				ctx = proxy.ContextWithInboundTag(ctx, w.tag)
			}
			ctx = proxy.ContextWithSource(ctx, source)
			ctx = proxy.ContextWithInboundDestination(ctx, v2net.UDPDestination(w.address, w.port))
			w.proxy.Process(ctx, v2net.Network_UDP, conn)
			conn.cancel()
		}()
	}
}

func (w *udpWorker) removeConn(src v2net.Destination) {
	w.Lock()
	delete(w.activeConn, src)
	w.Unlock()
}

func (w *udpWorker) Start() error {
	w.activeConn = make(map[v2net.Destination]*udpConn)
	ctx, cancel := context.WithCancel(context.Background())
	w.ctx = ctx
	w.cancel = cancel
	h, err := udp.ListenUDP(w.address, w.port, udp.ListenOption{
		Callback:            w.callback,
		ReceiveOriginalDest: w.recvOrigDest,
	})
	if err != nil {
		return err
	}
	w.hub = h
	return nil
}

func (w *udpWorker) Close() {
	w.hub.Close()
	w.cancel()
}

func (w *udpWorker) monitor() {
	for {
		select {
		case <-w.ctx.Done():
			return
		case <-time.After(time.Second * 16):
			nowSec := time.Now().Unix()
			w.Lock()
			for addr, conn := range w.activeConn {
				if nowSec-conn.lastActivityTime > 8 {
					w.removeConn(addr)
					conn.Close()
				}
			}
		}
	}
}

func (w *udpWorker) Port() v2net.Port {
	return w.port
}

func (w *udpWorker) Proxy() proxy.Inbound {
	return w.proxy
}
<commit_msg>fix udp worker<commit_after>package inbound

import (
	"context"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/log"
	v2net "v2ray.com/core/common/net"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/internet/tcp"
	"v2ray.com/core/transport/internet/udp"
)

type worker interface {
	Start() error
	Close()
	Port() v2net.Port
	Proxy() proxy.Inbound
}

type tcpWorker struct {
	address          v2net.Address
	port             v2net.Port
	proxy            proxy.Inbound
	stream           *internet.StreamConfig
	recvOrigDest     bool
	tag              string
	allowPassiveConn bool

	ctx    context.Context
	cancel context.CancelFunc
	hub    *internet.TCPHub
}

func (w *tcpWorker) callback(conn internet.Connection) {
	ctx, cancel := context.WithCancel(w.ctx)
	if w.recvOrigDest {
		dest := tcp.GetOriginalDestination(conn)
		if dest.IsValid() {
			ctx = proxy.ContextWithOriginalDestination(ctx, dest)
		}
	}
	if len(w.tag) > 0 {
		ctx = proxy.ContextWithInboundTag(ctx, w.tag)
	}
	ctx = proxy.ContextWithAllowPassiveConnection(ctx, w.allowPassiveConn)
	ctx = proxy.ContextWithInboundDestination(ctx, v2net.TCPDestination(w.address, w.port))
	w.proxy.Process(ctx, v2net.Network_TCP, conn)
	cancel()
	conn.Close()
}

func (w *tcpWorker) Proxy() proxy.Inbound {
	return w.proxy
}

func (w *tcpWorker) Start() error {
	ctx, cancel := context.WithCancel(context.Background())
	w.ctx = ctx
	w.cancel = cancel
	hub, err := internet.ListenTCP(w.address, w.port, w.callback, w.stream)
	if err != nil {
		return err
	}
	w.hub = hub
	return nil
}

func (w *tcpWorker) Close() {
	log.Debug("Proxyman|TCPWorker: Closed. ", w.port)
	w.hub.Close()
	w.cancel()
}

func (w *tcpWorker) Port() v2net.Port {
	return w.port
}

type udpConn struct {
	lastActivityTime int64 // in seconds
	input            chan []byte
	output           func([]byte) (int, error)
	remote           net.Addr
	local            net.Addr
	cancel           context.CancelFunc
}

func (c *udpConn) updateActivity() {
	atomic.StoreInt64(&c.lastActivityTime, time.Now().Unix())
}

func (c *udpConn) Read(buf []byte) (int, error) {
	in, open := <-c.input
	if !open {
		return 0, io.EOF
	}
	c.updateActivity()
	return copy(buf, in), nil
}

func (c *udpConn) Write(buf []byte) (int, error) {
	n, err := c.output(buf)
	if err == nil {
		c.updateActivity()
	}
	return n, err
}

func (c *udpConn) Close() error {
	return nil
}

func (c *udpConn) RemoteAddr() net.Addr {
	return c.remote
}

func (c *udpConn) LocalAddr() net.Addr {
	return c.remote
}

func (*udpConn) SetDeadline(time.Time) error {
	return nil
}

func (*udpConn) SetReadDeadline(time.Time) error {
	return nil
}

func (*udpConn) SetWriteDeadline(time.Time) error {
	return nil
}

func (*udpConn) Reusable() bool {
	return false
}

func (*udpConn) SetReusable(bool) {}

type udpWorker struct {
	sync.RWMutex

	proxy        proxy.Inbound
	hub          *udp.Hub
	address      v2net.Address
	port         v2net.Port
	recvOrigDest bool
	tag          string

	ctx        context.Context
	cancel     context.CancelFunc
	activeConn map[v2net.Destination]*udpConn
}

func (w *udpWorker) getConnection(src v2net.Destination) (*udpConn, bool) {
	w.Lock()
	defer w.Unlock()

	if conn, found := w.activeConn[src]; found {
		return conn, true
	}

	conn := &udpConn{
		input: make(chan []byte, 32),
		output: func(b []byte) (int, error) {
			return w.hub.WriteTo(b, src)
		},
		remote: &net.UDPAddr{
			IP:   src.Address.IP(),
			Port: int(src.Port),
		},
		local: &net.UDPAddr{
			IP:   w.address.IP(),
			Port: int(w.port),
		},
	}
	w.activeConn[src] = conn

	conn.updateActivity()
	return conn, false
}

func (w *udpWorker) callback(b *buf.Buffer, source v2net.Destination, originalDest v2net.Destination) {
	conn, existing := w.getConnection(source)
	conn.input <- b.Bytes()

	if !existing {
		go func() {
			ctx := w.ctx
			ctx, cancel := context.WithCancel(ctx)
			conn.cancel = cancel
			if originalDest.IsValid() {
				ctx = proxy.ContextWithOriginalDestination(ctx, originalDest)
			}
			if len(w.tag) > 0 {
				ctx = proxy.ContextWithInboundTag(ctx, w.tag)
			}
			ctx = proxy.ContextWithSource(ctx, source)
			ctx = proxy.ContextWithInboundDestination(ctx, v2net.UDPDestination(w.address, w.port))
			w.proxy.Process(ctx, v2net.Network_UDP, conn)
			w.removeConn(source)
			cancel()
		}()
	}
}

func (w *udpWorker) removeConn(src v2net.Destination) {
	w.Lock()
	delete(w.activeConn, src)
	w.Unlock()
}

func (w *udpWorker) Start() error {
	w.activeConn = make(map[v2net.Destination]*udpConn)
	ctx, cancel := context.WithCancel(context.Background())
	w.ctx = ctx
	w.cancel = cancel
	h, err := udp.ListenUDP(w.address, w.port, udp.ListenOption{
		Callback:            w.callback,
		ReceiveOriginalDest: w.recvOrigDest,
	})
	if err != nil {
		return err
	}
	go w.monitor()
	w.hub = h
	return nil
}

func (w *udpWorker) Close() {
	w.hub.Close()
	w.cancel()
}

func (w *udpWorker) monitor() {
	for {
		select {
		case <-w.ctx.Done():
			return
		case <-time.After(time.Second * 16):
			nowSec := time.Now().Unix()
			w.Lock()
			for addr, conn := range w.activeConn {
				if nowSec-conn.lastActivityTime > 8 {
					delete(w.activeConn, addr)
					conn.cancel()
				}
			}
			w.Unlock()
		}
	}
}

func (w *udpWorker) Port() v2net.Port {
	return w.port
}

func (w *udpWorker) Proxy() proxy.Inbound {
	return w.proxy
}
<|endoftext|>


<commit_before>package common

import (
	"fmt"
	"strconv"
	"strings"

	. "firempq/api"
	"io"
)

type CallFuncType func([]string) IResponse

type ItemsResponse struct {
	items []IItem
}

type DictResponse struct {
	dict map[string]interface{}
}

func NewDictResponse(dict map[string]interface{}) *DictResponse {
	return &DictResponse{dict}
}

func (self *DictResponse) GetDict() map[string]interface{} {
	return self.dict
}

func (self *DictResponse) getResponseChunks() []string {
	data := make([]string, 0, 3+9*len(self.dict))
	data = append(data, "+DATA %")
	data = append(data, strconv.Itoa(len(self.dict)))
	for k, v := range self.dict {
		data = append(data, "\n")
		data = append(data, k)
		data = append(data, " ")
		switch t := v.(type) {
		case string:
			data = append(data, t)
		case int:
			data = append(data, ":")
			data = append(data, strconv.Itoa(t))
		case int64:
			data = append(data, ":")
			data = append(data, strconv.Itoa(int(t)))
		case bool:
			if t {
				data = append(data, "?t")
			} else {
				data = append(data, "?f")
			}
		}
	}
	return data
}

func (self *DictResponse) GetResponse() string {
	return strings.Join(self.getResponseChunks(), "")
}

func (self *DictResponse) WriteResponse(buff io.Writer) error {
	var err error
	for _, s := range self.getResponseChunks() {
		_, err = buff.Write(UnsafeStringToBytes(s))
		if err != nil {
			return err
		}
	}
	return nil
}

func (self *DictResponse) IsError() bool { return false }

func NewItemsResponse(items []IItem) *ItemsResponse {
	return &ItemsResponse{items}
}

func (self *ItemsResponse) GetItems() []IItem {
	return self.items
}

func (self *ItemsResponse) getResponseChunks() []string {
	data := make([]string, 0, 3+9*len(self.items))
	data = append(data, "+DATA *")
	data = append(data, strconv.Itoa(len(self.items)))
	for _, item := range self.items {
		data = append(data, " ")
		data = append(data, "%2 ID ")
		data = append(data, EncodeRespStringTo(data, item.GetId())...)
		data = append(data, " PL ")
		data = append(data, EncodeRespStringTo(data, item.GetPayload())...)
	}
	return data
}

func (self *ItemsResponse) GetResponse() string {
	return strings.Join(self.getResponseChunks(), "")
}

func (self *ItemsResponse) WriteResponse(buff io.Writer) error {
	var err error
	for _, s := range self.getResponseChunks() {
		_, err = buff.Write(UnsafeStringToBytes(s))
		if err != nil {
			return err
		}
	}
	return nil
}

func (self *ItemsResponse) IsError() bool {
	return false
}

// ErrorResponse is an error response.
type ErrorResponse struct {
	ErrorText string
	ErrorCode int64
}

func (e *ErrorResponse) Error() string {
	return e.ErrorText
}

func (e *ErrorResponse) GetResponse() string {
	return fmt.Sprintf("-ERR %s %s",
		EncodeRespInt64(e.ErrorCode),
		EncodeRespString(e.ErrorText))
}

func (e *ErrorResponse) WriteResponse(buff io.Writer) error {
	_, err := buff.Write(UnsafeStringToBytes(e.GetResponse()))
	return err
}

func (e *ErrorResponse) IsError() bool {
	return true
}
<commit_msg>Data encoder that uses built in encoder from messages.<commit_after>package common

import (
	"fmt"
	"io"
	"strconv"
	"strings"

	. "firempq/api"
	. "firempq/common/response_encoder"
)

type CallFuncType func([]string) IResponse

type DictResponse struct {
	dict map[string]interface{}
}

func NewDictResponse(dict map[string]interface{}) *DictResponse {
	return &DictResponse{
		dict: dict,
	}
}

func (self *DictResponse) GetDict() map[string]interface{} {
	return self.dict
}

func (self *DictResponse) getResponseChunks() []string {
	data := make([]string, 0, 3+9*len(self.dict))
	data = append(data, "+DATA %")
	data = append(data, strconv.Itoa(len(self.dict)))
	for k, v := range self.dict {
		data = append(data, "\n")
		data = append(data, k)
		data = append(data, " ")
		switch t := v.(type) {
		case string:
			data = append(data, t)
		case int:
			data = append(data, ":")
			data = append(data, strconv.Itoa(t))
		case int64:
			data = append(data, ":")
			data = append(data, strconv.Itoa(int(t)))
		case bool:
			if t {
				data = append(data, "?t")
			} else {
				data = append(data, "?f")
			}
		}
	}
	return data
}

func (self *DictResponse) GetResponse() string {
	return strings.Join(self.getResponseChunks(), "")
}

func (self *DictResponse) WriteResponse(buff io.Writer) error {
	var err error
	for _, s := range self.getResponseChunks() {
		_, err = buff.Write(UnsafeStringToBytes(s))
		if err != nil {
			return err
		}
	}
	return nil
}

func (self *DictResponse) IsError() bool { return false }

type ItemsResponse struct {
	items []IResponseItem
}

func NewItemsResponse(items []IResponseItem) *ItemsResponse {
	return &ItemsResponse{
		items: items,
	}
}

func (self *ItemsResponse) GetItems() []IResponseItem {
	return self.items
}

func (self *ItemsResponse) getResponseChunks() []string {
	data := make([]string, 0, 3+9*len(self.items))
	data = append(data, "+DATA")
	data = append(data, EncodeArraySize(len(self.items)))
	for _, item := range self.items {
		data = append(data, item.Encode())
	}
	return data
}

func (self *ItemsResponse) GetResponse() string {
	return strings.Join(self.getResponseChunks(), "")
}

func (self *ItemsResponse) WriteResponse(buff io.Writer) error {
	var err error
	for _, s := range self.getResponseChunks() {
		_, err = buff.Write(UnsafeStringToBytes(s))
		if err != nil {
			return err
		}
	}
	return nil
}

func (self *ItemsResponse) IsError() bool {
	return false
}

// ErrorResponse is an error response.
type ErrorResponse struct {
	ErrorText string
	ErrorCode int64
}

func (e *ErrorResponse) Error() string {
	return e.ErrorText
}

func (e *ErrorResponse) GetResponse() string {
	return fmt.Sprintf("-ERR %s %s",
		EncodeRespInt64(e.ErrorCode),
		EncodeRespString(e.ErrorText))
}

func (e *ErrorResponse) WriteResponse(buff io.Writer) error {
	_, err := buff.Write(UnsafeStringToBytes(e.GetResponse()))
	return err
}

func (e *ErrorResponse) IsError() bool {
	return true
}
<|endoftext|>