<commit_before>/* plumber: a deployment pipeline template
 * Copyright (C) 2014 Recruit Technologies Co., Ltd. and contributors
 * (see CONTRIBUTORS.md)
 *
 * 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 stages

import (
	"container/list"
	"fmt"
)

type Stage interface {
	AddChildStage(Stage)
	GetChildStages() list.List
	GetStageName() string
	SetStageName(string)
	GetInputCh() *chan Mediator
	SetInputCh(*chan Mediator)
	GetOutputCh() *chan Mediator
	SetOutputCh(*chan Mediator)
}

type Runner interface {
	Run() bool
}

type Mediator struct {
	States map[string]string
}

func InitStage(stageType string) Stage {
	switch stageType {
	case "command":
		return new(CommandStage)
	}
	return nil
}

func ExecuteStage(stage Stage, monitorCh *chan Mediator) {
	mediator := <-*stage.GetInputCh()

	fmt.Println("mediator received: %v", mediator)
	fmt.Println("execute as parent: %v", stage)
	fmt.Println("execute as parent name %v", stage.GetStageName())

	result := stage.(Runner).Run()
	fmt.Println("execute as parent result %v", result)

	mediator.States[stage.GetStageName()] = fmt.Sprintf("%v", result)

	setChildStatus(&stage, &mediator, "waiting")

	if childStages := stage.GetChildStages(); childStages.Len() > 0 {
		fmt.Println("execute childstage: %v", childStages)

		for childStage := childStages.Front(); childStage != nil; childStage = childStage.Next() {
			fmt.Printf("child name %+v\n", childStage.Value.(Stage).GetStageName())
			childInputCh := make(chan Mediator)
			childOutputCh := make(chan Mediator)
			childStage.Value.(Stage).SetInputCh(&childInputCh)
			childStage.Value.(Stage).SetOutputCh(&childOutputCh)

			name := childStage.Value.(Stage).GetStageName()
			mediator.States[name] = fmt.Sprintf("%v", "waiting")

			go func(m Mediator) {
				childInputCh <- m
				close(childInputCh)
			}(mediator)

			go func(stage Stage) {
				ExecuteStage(stage, monitorCh)
			}(childStage.Value.(Stage))
		}

		for childStage := childStages.Front(); childStage != nil; childStage = childStage.Next() {
			go func(stage Stage) {
				childReceived := <-*stage.GetOutputCh()
				name := stage.GetStageName()
				fmt.Printf("child state %+v\n", childReceived.States[name])
				mediator.States[name] = fmt.Sprintf("%v", childReceived.States[name])
			}(childStage.Value.(Stage))
		}
	}

	go func() {
		*stage.GetOutputCh() <- mediator
		close(*stage.GetOutputCh())
	}()
	*monitorCh <- mediator

	closeAfterExecute(&mediator, monitorCh)
}

func setChildStatus(stage *Stage, mediator *Mediator, status string) {
	if childStages := (*stage).GetChildStages(); childStages.Len() > 0 {
		for childStage := childStages.Front(); childStage != nil; childStage = childStage.Next() {
			name := childStage.Value.(Stage).GetStageName()
			mediator.States[name] = fmt.Sprintf("%v", status)
		}
	}
}

func closeAfterExecute(mediator *Mediator, monitCh *chan Mediator) {
	allDone := true
	for _, v := range mediator.States {
		if v == "waiting" {
			allDone = false
		}
	}

	if allDone {
		fmt.Printf("closing monitor channel.. %v\n", mediator)
		close(*monitCh)
	}
}

func Execute(stage Stage, mediator Mediator) Mediator {
	inputCh := make(chan Mediator)
	outputCh := make(chan Mediator)
	monitorCh := make(chan Mediator)
	stage.SetInputCh(&inputCh)
	stage.SetOutputCh(&outputCh)
	name := stage.GetStageName()
	fmt.Printf("----- Execute %v start ------\n", name)

	mediator.States[name] = fmt.Sprintf("%v", "waiting")

	go func(mediator Mediator) {
		inputCh <- mediator
		close(inputCh)
	}(mediator)

	var lastReceive Mediator

	go ExecuteStage(stage, &monitorCh)

	for {
		receive, ok := <-monitorCh
		if !ok {
			fmt.Println("monitorCh closed")
			fmt.Printf("----- Execute %v done ------\n\n", name)
			return lastReceive
		}
		fmt.Printf("monitorCh received  %+v\n", receive)
		lastReceive = receive
	}
}
<commit_msg>initialize channels of stage<commit_after>/* plumber: a deployment pipeline template
 * Copyright (C) 2014 Recruit Technologies Co., Ltd. and contributors
 * (see CONTRIBUTORS.md)
 *
 * 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 stages

import (
	"container/list"
	"fmt"
)

type Stage interface {
	AddChildStage(Stage)
	GetChildStages() list.List
	GetStageName() string
	SetStageName(string)
	GetInputCh() *chan Mediator
	SetInputCh(*chan Mediator)
	GetOutputCh() *chan Mediator
	SetOutputCh(*chan Mediator)
}

type Runner interface {
	Run() bool
}

type Mediator struct {
	States map[string]string
}

func InitStage(stageType string) Stage {
	var stage Stage
	switch stageType {
	case "command":
		stage = new(CommandStage)
	}

	prepareCh(&stage)
	return stage
}

func prepareCh(stage *Stage) {
	in := make(chan Mediator)
	out := make(chan Mediator)
	(*stage).SetInputCh(&in)
	(*stage).SetOutputCh(&out)
}

func ExecuteStage(stage Stage, monitorCh *chan Mediator) {
	mediator := <-*stage.GetInputCh()

	fmt.Println("mediator received: %v", mediator)
	fmt.Println("execute as parent: %v", stage)
	fmt.Println("execute as parent name %v", stage.GetStageName())

	result := stage.(Runner).Run()
	fmt.Println("execute as parent result %v", result)

	mediator.States[stage.GetStageName()] = fmt.Sprintf("%v", result)

	setChildStatus(&stage, &mediator, "waiting")

	if childStages := stage.GetChildStages(); childStages.Len() > 0 {
		fmt.Println("execute childstage: %v", childStages)

		for childStage := childStages.Front(); childStage != nil; childStage = childStage.Next() {
			fmt.Printf("child name %+v\n", childStage.Value.(Stage).GetStageName())
			childInputCh := *childStage.Value.(Stage).GetInputCh()

			name := childStage.Value.(Stage).GetStageName()
			mediator.States[name] = fmt.Sprintf("%v", "waiting")

			go func(m Mediator) {
				childInputCh <- m
				close(childInputCh)
			}(mediator)

			go func(stage Stage) {
				ExecuteStage(stage, monitorCh)
			}(childStage.Value.(Stage))
		}

		for childStage := childStages.Front(); childStage != nil; childStage = childStage.Next() {
			go func(stage Stage) {
				childReceived := <-*stage.GetOutputCh()
				name := stage.GetStageName()
				fmt.Printf("child state %+v\n", childReceived.States[name])
				mediator.States[name] = fmt.Sprintf("%v", childReceived.States[name])
			}(childStage.Value.(Stage))
		}
	}

	go func() {
		*stage.GetOutputCh() <- mediator
		close(*stage.GetOutputCh())
	}()
	*monitorCh <- mediator

	closeAfterExecute(&mediator, monitorCh)
}

func setChildStatus(stage *Stage, mediator *Mediator, status string) {
	if childStages := (*stage).GetChildStages(); childStages.Len() > 0 {
		for childStage := childStages.Front(); childStage != nil; childStage = childStage.Next() {
			name := childStage.Value.(Stage).GetStageName()
			mediator.States[name] = fmt.Sprintf("%v", status)
		}
	}
}

func closeAfterExecute(mediator *Mediator, monitCh *chan Mediator) {
	allDone := true
	for _, v := range mediator.States {
		if v == "waiting" {
			allDone = false
		}
	}

	if allDone {
		fmt.Printf("closing monitor channel.. %v\n", mediator)
		close(*monitCh)
	}
}

func Execute(stage Stage, mediator Mediator) Mediator {
	monitorCh := make(chan Mediator)
	name := stage.GetStageName()
	fmt.Printf("----- Execute %v start ------\n", name)

	mediator.States[name] = fmt.Sprintf("%v", "waiting")

	go func(mediator Mediator) {
		*stage.GetInputCh() <- mediator
		close(*stage.GetInputCh())
	}(mediator)

	var lastReceive Mediator

	go ExecuteStage(stage, &monitorCh)

	for {
		receive, ok := <-monitorCh
		if !ok {
			fmt.Println("monitorCh closed")
			fmt.Printf("----- Execute %v done ------\n\n", name)
			return lastReceive
		}
		fmt.Printf("monitorCh received  %+v\n", receive)
		lastReceive = receive
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"bytes"
	"math"
	"time"
	"io/ioutil"
	"net/http"

	"os"
	"image"
	"image/png"
	"image/draw"
    "image/color"
)

// http://stackoverflow.com/questions/35964656/golang-how-to-concatenate-append-images-to-one-another

// Create a struct to deal with pixel
type Pixel struct {
    Point image.Point
    Color color.Color
}

// Decode image.Image's pixel data into []*Pixel
func DecodePixelsFromImage(img image.Image, offsetX, offsetY int) []*Pixel {
    pixels := []*Pixel{}
    for y := 0; y <= img.Bounds().Max.Y; y++ {
        for x := 0; x <= img.Bounds().Max.X; x++ {
            p := &Pixel{
                Point: image.Point{x + offsetX, y + offsetY},
                Color: img.At(x, y),
            }
            pixels = append(pixels, p)
        }
    }
    return pixels
}

var client = &http.Client{
	Timeout: time.Second * 5,
}

var tiles_map map[int][]image.Image

func init() {
	tiles_map = make(map[int][]image.Image)
}

// degTorad converts degree to radians.
func degTorad(deg float64) float64 {
	return deg * math.Pi / 180;
}

// deg2num converts latlng to tile number
func deg2num(lat_deg float64, lon_deg float64, zoom int) (int, int) {
    lat_rad := degTorad(lat_deg)
    n := math.Pow(2.0, float64(zoom))
    xtile := int((lon_deg + 180.0) / 360.0 * n)
    ytile := int((1.0 - math.Log(math.Tan(lat_rad) + (1 / math.Cos(lat_rad))) / math.Pi) / 2.0 * n)
    return xtile, ytile
}

// xyz
type xyz struct {
	x int
	y int
	z int
}

// getTileNames
func getTileNames(minlat, maxlat, minlng, maxlng float64, z int) []xyz {
	tiles := []xyz{}
	// // upper right
	// ur_tile_x, ur_tile_y := deg2num(float64(70), float64(16), z)
	// // lower left
	// ll_tile_x, ll_tile_y := deg2num(float64(35), float64(0), z)
	// upper right
	ur_tile_x, ur_tile_y := deg2num(maxlat, maxlng, z)
	// lower left
	ll_tile_x, ll_tile_y := deg2num(minlat, minlng, z)

	for x := ll_tile_x-1; x < ur_tile_x+1; x++ {
		if x < 0 { x++ }
		for y := ur_tile_y-1; y < ll_tile_y+1; y++ {
			if y < 0 { y++ }
			tiles = append(tiles, xyz{x,y,z})
		}
	}
	return tiles
}

// getTilePngFromUrl
func getTilePngBytesFromUrl(tile_url string) []byte {
	fmt.Println("GET", tile_url)

	// Just a simple GET request to the image URL
    // We get back a *Response, and an error
	res, err := client.Get(tile_url)
	if err != nil {
        fmt.Printf("Error http.Get -> %v\n", err)
		return []byte("")
    }

	// We read all the bytes of the image
    // Types: data []byte
    data, err := ioutil.ReadAll(res.Body)

	// You have to manually close the body, check docs
    // This is required if you want to use things like
    // Keep-Alive and other HTTP sorcery.
    defer res.Body.Close()

    if err != nil {
        fmt.Printf("Error ioutil.ReadAll -> %v\n", err)
		return []byte("")
    }

	// You can now save it to disk or whatever...
	// ioutil.WriteFile("TMP_TILE.png", data, 0644)
	// ioutil.WriteFile("TMP_TILE.png", data, 0666)

	return data
}

// BytesToPngImage
func BytesToPngImage(b []byte) image.Image {
	img, err := png.Decode(bytes.NewReader(b))
	if nil != err {
		panic(err)
	}
	return img
}

func mergePngTiles() {
	// Get bounds for new image
	size := 256
	cols := 0
	rows := 0
	for i := range tiles_map {
		cols += size
		rows = len(tiles_map[i]) * size
	}

	// collect pixel data from each image.
	// each image has a x-offset and Y-offset from the first.
	var pixelSum []*Pixel
	x := 0
	for i := range tiles_map {
		y := 0
		for j := range tiles_map[i] {
			pixels := DecodePixelsFromImage(tiles_map[i][j], x, y)
			pixelSum = append(pixelSum, pixels...)
			y += size
			fmt.Println(x,y)
		}
		x += size
	}

	// Set a new size for the new image equal to the max width
	// of bigger image and max height of two images combined
	newRect := image.Rectangle{
		Min: image.Point{X: 0, Y: 0},
		Max: image.Point{X: cols, Y: rows},
	}

	finImage := image.NewRGBA(newRect)
	// This is the cool part, all you have to do is loop through
	// each Pixel and set the image's color on the go
	for _, px := range pixelSum {
			finImage.Set(
				px.Point.X,
				px.Point.Y,
				px.Color,
			)
	}
	draw.Draw(finImage, finImage.Bounds(), finImage, image.Point{0, 0}, draw.Src)

	// Create a new file and write to it
	out, err := os.Create("./output.png")
	if err != nil {
		panic(err)
		os.Exit(1)
	}
	err = png.Encode(out, finImage)
	if err != nil {
		panic(err)
		os.Exit(1)
	}
}

func main() {
	base_url := "http://localhost:8080/tms/1.0/population"
	// base_url := "http://localhost:8080/tms/1.0/osm"

	minlat := float64(35)
	maxlat := float64(70)
	minlng := float64(0)
	maxlng := float64(16)
	zoom := 7

	tiles := getTileNames(minlat, maxlat, minlng, maxlng, zoom)

	for _, v := range tiles {
		tile_url := fmt.Sprintf("/%v/%v/%v.png", v.z, v.x, v.y)
		data := getTilePngBytesFromUrl(base_url+tile_url)

		img := BytesToPngImage(data)
		tiles_map[v.x] = append(tiles_map[v.x], img)

	}

	mergePngTiles()
}

/*

export GOPATH="`pwd`"

*/
<commit_msg>minor changes to stitch tiles<commit_after>package main

import (
	"fmt"
	"bytes"
	"math"
	"time"
	"io/ioutil"
	"net/http"

	"os"
	"image"
	"image/png"
	"image/draw"
    "image/color"
)

// http://stackoverflow.com/questions/35964656/golang-how-to-concatenate-append-images-to-one-another

// Create a struct to deal with pixel
type Pixel struct {
    Point image.Point
    Color color.Color
}

// Decode image.Image's pixel data into []*Pixel
func DecodePixelsFromImage(img image.Image, offsetX, offsetY int) []*Pixel {
    pixels := []*Pixel{}
    for y := 0; y <= img.Bounds().Max.Y; y++ {
        for x := 0; x <= img.Bounds().Max.X; x++ {
            p := &Pixel{
                Point: image.Point{x + offsetX, y + offsetY},
                Color: img.At(x, y),
            }
            pixels = append(pixels, p)
        }
    }
    return pixels
}

var client = &http.Client{
	Timeout: time.Second * 5,
}

var tiles_map map[int][]image.Image

func init() {
	tiles_map = make(map[int][]image.Image)
}

// degTorad converts degree to radians.
func degTorad(deg float64) float64 {
	return deg * math.Pi / 180;
}

// deg2num converts latlng to tile number
func deg2num(lat_deg float64, lon_deg float64, zoom int) (int, int) {
    lat_rad := degTorad(lat_deg)
    n := math.Pow(2.0, float64(zoom))
    xtile := int((lon_deg + 180.0) / 360.0 * n)
    ytile := int((1.0 - math.Log(math.Tan(lat_rad) + (1 / math.Cos(lat_rad))) / math.Pi) / 2.0 * n)
    return xtile, ytile
}

// xyz
type xyz struct {
	x int
	y int
	z int
}

// getTileNames
func getTileNames(minlat, maxlat, minlng, maxlng float64, z int) []xyz {
	tiles := []xyz{}

	// upper right
	ur_tile_x, ur_tile_y := deg2num(maxlat, maxlng, z)
	// lower left
	ll_tile_x, ll_tile_y := deg2num(minlat, minlng, z)

	for x := ll_tile_x-1; x < ur_tile_x+1; x++ {
		if x < 0 { x++ }
		for y := ur_tile_y-1; y < ll_tile_y+1; y++ {
			if y < 0 { y++ }
			tiles = append(tiles, xyz{x,y,z})
		}
	}
	return tiles
}

// getTilePngFromUrl
func getTilePngBytesFromUrl(tile_url string) []byte {
	fmt.Println("GET", tile_url)

	// Just a simple GET request to the image URL
    // We get back a *Response, and an error
	res, err := client.Get(tile_url)
	if err != nil {
        fmt.Printf("Error http.Get -> %v\n", err)
		return []byte("")
    }

	// We read all the bytes of the image
    // Types: data []byte
    data, err := ioutil.ReadAll(res.Body)

	// You have to manually close the body, check docs
    // This is required if you want to use things like
    // Keep-Alive and other HTTP sorcery.
    defer res.Body.Close()

    if err != nil {
        fmt.Printf("Error ioutil.ReadAll -> %v\n", err)
		return []byte("")
    }

	// You can now save it to disk or whatever...
	// ioutil.WriteFile("TMP_TILE.png", data, 0644)
	// ioutil.WriteFile("TMP_TILE.png", data, 0666)

	return data
}

// BytesToPngImage
func BytesToPngImage(b []byte) image.Image {
	img, err := png.Decode(bytes.NewReader(b))
	if nil != err {
		panic(err)
	}
	return img
}

func mergePngTiles() {
	// Get bounds for new image
	size := 256
	cols := 0
	rows := 0
	for i := range tiles_map {
		cols += size
		rows = len(tiles_map[i]) * size
	}

	// collect pixel data from each image.
	// each image has a x-offset and Y-offset from the first.
	var pixelSum []*Pixel
	x := 0
	for i := range tiles_map {
		y := 0
		for j := range tiles_map[i] {
			pixels := DecodePixelsFromImage(tiles_map[i][j], x, y)
			pixelSum = append(pixelSum, pixels...)
			y += size
			fmt.Println(x,y)
		}
		x += size
	}

	// Set a new size for the new image equal to the max width
	// of bigger image and max height of two images combined
	newRect := image.Rectangle{
		Min: image.Point{X: 0, Y: 0},
		Max: image.Point{X: cols, Y: rows},
	}

	// Create new image
	finImage := image.NewRGBA(newRect)

	// This is the cool part, all you have to do is loop through
	// each Pixel and set the image's color on the go
	for _, px := range pixelSum {
			finImage.Set(
				px.Point.X,
				px.Point.Y,
				px.Color,
			)
	}
	draw.Draw(finImage, finImage.Bounds(), finImage, image.Point{0, 0}, draw.Src)

	// Create a new file and write to it
	out, err := os.Create("./output.png")
	if err != nil {
		panic(err)
		os.Exit(1)
	}
	err = png.Encode(out, finImage)
	if err != nil {
		panic(err)
		os.Exit(1)
	}
}

func main() {
	base_url := "http://localhost:8080/tms/1.0/population"
	// base_url := "http://localhost:8080/tms/1.0/osm"

	minlat := float64(35)
	maxlat := float64(70)
	minlng := float64(0)
	maxlng := float64(16)
	zoom := 7

	tiles := getTileNames(minlat, maxlat, minlng, maxlng, zoom)

	for _, v := range tiles {
		tile_url := fmt.Sprintf("/%v/%v/%v.png", v.z, v.x, v.y)
		data := getTilePngBytesFromUrl(base_url+tile_url)

		img := BytesToPngImage(data)
		tiles_map[v.x] = append(tiles_map[v.x], img)

	}

	mergePngTiles()
}

/*

export GOPATH="`pwd`"

*/
<|endoftext|>


<commit_before>// Copyright (c) 2019 Shivaram Lingamneni <slingamn@cs.stanford.edu>
// released under the MIT license

package irc

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/oragono/oragono/irc/history"
	"github.com/oragono/oragono/irc/utils"
)

const (
	// #829, also see "Case 2" in the "three cases" below:
	zncPlaybackCommandExpiration = time.Second * 30
)

type zncCommandHandler func(client *Client, command string, params []string, rb *ResponseBuffer)

var zncHandlers = map[string]zncCommandHandler{
	"*playback": zncPlaybackHandler,
}

func zncPrivmsgHandler(client *Client, command string, privmsg string, rb *ResponseBuffer) {
	zncModuleHandler(client, command, strings.Fields(privmsg), rb)
}

func zncModuleHandler(client *Client, command string, params []string, rb *ResponseBuffer) {
	command = strings.ToLower(command)
	if subHandler, ok := zncHandlers[command]; ok {
		subHandler(client, command, params, rb)
	} else {
		nick := rb.target.Nick()
		rb.Add(nil, client.server.name, "NOTICE", nick, fmt.Sprintf(client.t("Oragono does not emulate the ZNC module %s"), command))
		rb.Add(nil, "*status!znc@znc.in", "NOTICE", nick, fmt.Sprintf(client.t("No such module [%s]"), command))
	}
}

// "number of seconds (floating point for millisecond precision) elapsed since January 1, 1970"
func zncWireTimeToTime(str string) (result time.Time) {
	var secondsPortion, fracPortion string
	dot := strings.IndexByte(str, '.')
	if dot == -1 {
		secondsPortion = str
	} else {
		secondsPortion = str[:dot]
		fracPortion = str[dot:]
	}
	seconds, _ := strconv.ParseInt(secondsPortion, 10, 64)
	fraction, _ := strconv.ParseFloat(fracPortion, 64)
	return time.Unix(seconds, int64(fraction*1000000000)).UTC()
}

func timeToZncWireTime(t time.Time) (result string) {
	secs := t.Unix()
	nano := t.UnixNano() - (secs * 1000000000)
	return fmt.Sprintf("%d.%d", secs, nano)
}

type zncPlaybackTimes struct {
	start   time.Time
	end     time.Time
	targets utils.StringSet // nil for "*" (everything), otherwise the channel names
	setAt   time.Time
}

func (z *zncPlaybackTimes) ValidFor(target string) bool {
	if z == nil {
		return false
	}

	if time.Now().Sub(z.setAt) > zncPlaybackCommandExpiration {
		return false
	}

	if z.targets == nil {
		return true
	}

	return z.targets.Has(target)
}

// https://wiki.znc.in/Playback
func zncPlaybackHandler(client *Client, command string, params []string, rb *ResponseBuffer) {
	if len(params) == 0 {
		return
	}
	switch strings.ToLower(params[0]) {
	case "play":
		zncPlaybackPlayHandler(client, command, params, rb)
	case "list":
		zncPlaybackListHandler(client, command, params, rb)
	default:
		return
	}
}

// PRIVMSG *playback :play <target> [lower_bound] [upper_bound]
// e.g., PRIVMSG *playback :play * 1558374442
func zncPlaybackPlayHandler(client *Client, command string, params []string, rb *ResponseBuffer) {
	if len(params) < 2 || len(params) > 4 {
		return
	}
	targetString := params[1]

	now := time.Now().UTC()
	var start, end time.Time
	switch len(params) {
	case 2:
		// #1205: this should have the same semantics as `LATEST *`
	case 3:
		// #831: this should have the same semantics as `LATEST timestamp=qux`,
		// or equivalently `BETWEEN timestamp=$now timestamp=qux`, as opposed to
		// `AFTER timestamp=qux` (this matters in the case where there are
		// more than znc-maxmessages available)
		start = now
		end = zncWireTimeToTime(params[2])
	case 4:
		start = zncWireTimeToTime(params[2])
		end = zncWireTimeToTime(params[3])
	}

	var targets utils.StringSet
	var nickTargets []string

	// three cases:
	// 1. the user's PMs get played back immediately upon receiving this
	// 2. if this is a new connection (from the server's POV), save the information
	// and use it to process subsequent joins. (This is the Textual behavior:
	// first send the playback PRIVMSG, then send the JOIN lines.)
	// 3. if this is a reattach (from the server's POV), immediately play back
	// history for channels that the client is already joined to. In this scenario,
	// there are three total attempts to play the history:
	//     3.1. During the initial reattach (no-op because the *playback privmsg
	//          hasn't been received yet, but they negotiated the znc.in/playback
	//          cap so we know we're going to receive it later)
	//     3.2  Upon receiving the *playback privmsg, i.e., now: we should play
	//          the relevant history lines
	//     3.3  When the client sends a subsequent redundant JOIN line for those
	//          channels; redundant JOIN is a complete no-op so we won't replay twice

	playPrivmsgs := false
	if params[1] == "*" {
		playPrivmsgs = true // XXX nil `targets` means "every channel"
	} else {
		targets = make(utils.StringSet)
		for _, targetName := range strings.Split(targetString, ",") {
			if targetName == "*self" {
				playPrivmsgs = true
			} else if strings.HasPrefix(targetName, "#") {
				if cfTarget, err := CasefoldChannel(targetName); err == nil {
					targets.Add(cfTarget)
				}
			} else {
				if cfNick, err := CasefoldName(targetName); err == nil {
					nickTargets = append(nickTargets, cfNick)
				}
			}
		}
	}

	if playPrivmsgs {
		zncPlayPrivmsgs(client, rb, "*", start, end)
	}

	rb.session.zncPlaybackTimes = &zncPlaybackTimes{
		start:   start,
		end:     end,
		targets: targets,
		setAt:   time.Now().UTC(),
	}

	for _, channel := range client.Channels() {
		if targets == nil || targets.Has(channel.NameCasefolded()) {
			channel.autoReplayHistory(client, rb, "")
			rb.Flush(true)
		}
	}

	for _, cfNick := range nickTargets {
		zncPlayPrivmsgs(client, rb, cfNick, start, end)
		rb.Flush(true)
	}
}

func zncPlayPrivmsgs(client *Client, rb *ResponseBuffer, target string, after, before time.Time) {
	_, sequence, _ := client.server.GetHistorySequence(nil, client, target)
	if sequence == nil {
		return
	}
	zncMax := client.server.Config().History.ZNCMax
	items, _, err := sequence.Between(history.Selector{Time: after}, history.Selector{Time: before}, zncMax)
	if err == nil && len(items) != 0 {
		client.replayPrivmsgHistory(rb, items, "", true)
	}
}

// PRIVMSG *playback :list
func zncPlaybackListHandler(client *Client, command string, params []string, rb *ResponseBuffer) {
	nick := client.Nick()
	for _, channel := range client.Channels() {
		_, sequence, err := client.server.GetHistorySequence(channel, client, "")
		if err != nil {
			client.server.logger.Error("internal", "couldn't get history sequence for ZNC list", err.Error())
			continue
		}
		items, _, err := sequence.Between(history.Selector{}, history.Selector{}, 1) // i.e., LATEST * 1
		if err != nil {
			client.server.logger.Error("internal", "couldn't query history for ZNC list", err.Error())
		} else if len(items) != 0 {
			stamp := timeToZncWireTime(items[0].Message.Time)
			rb.Add(nil, "*playback!znc@znc.in", "PRIVMSG", nick, fmt.Sprintf("%s 0 %s", channel.Name(), stamp))
		}
	}
}
<commit_msg>fix #1552<commit_after>// Copyright (c) 2019 Shivaram Lingamneni <slingamn@cs.stanford.edu>
// released under the MIT license

package irc

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/oragono/oragono/irc/history"
	"github.com/oragono/oragono/irc/utils"
)

const (
	// #829, also see "Case 2" in the "three cases" below:
	zncPlaybackCommandExpiration = time.Second * 30
)

type zncCommandHandler func(client *Client, command string, params []string, rb *ResponseBuffer)

var zncHandlers = map[string]zncCommandHandler{
	"*playback": zncPlaybackHandler,
}

func zncPrivmsgHandler(client *Client, command string, privmsg string, rb *ResponseBuffer) {
	zncModuleHandler(client, command, strings.Fields(privmsg), rb)
}

func zncModuleHandler(client *Client, command string, params []string, rb *ResponseBuffer) {
	command = strings.ToLower(command)
	if subHandler, ok := zncHandlers[command]; ok {
		subHandler(client, command, params, rb)
	} else {
		nick := rb.target.Nick()
		rb.Add(nil, client.server.name, "NOTICE", nick, fmt.Sprintf(client.t("Oragono does not emulate the ZNC module %s"), command))
		rb.Add(nil, "*status!znc@znc.in", "NOTICE", nick, fmt.Sprintf(client.t("No such module [%s]"), command))
	}
}

// "number of seconds (floating point for millisecond precision) elapsed since January 1, 1970"
func zncWireTimeToTime(str string) (result time.Time) {
	var secondsPortion, fracPortion string
	dot := strings.IndexByte(str, '.')
	if dot == -1 {
		secondsPortion = str
	} else {
		secondsPortion = str[:dot]
		fracPortion = str[dot:]
	}
	seconds, _ := strconv.ParseInt(secondsPortion, 10, 64)
	fraction, _ := strconv.ParseFloat(fracPortion, 64)
	return time.Unix(seconds, int64(fraction*1000000000)).UTC()
}

func timeToZncWireTime(t time.Time) (result string) {
	secs := t.Unix()
	nano := t.UnixNano() - (secs * 1000000000)
	return fmt.Sprintf("%d.%d", secs, nano)
}

type zncPlaybackTimes struct {
	start   time.Time
	end     time.Time
	targets utils.StringSet // nil for "*" (everything), otherwise the channel names
	setAt   time.Time
}

func (z *zncPlaybackTimes) ValidFor(target string) bool {
	if z == nil {
		return false
	}

	if time.Now().Sub(z.setAt) > zncPlaybackCommandExpiration {
		return false
	}

	if z.targets == nil {
		return true
	}

	return z.targets.Has(target)
}

// https://wiki.znc.in/Playback
func zncPlaybackHandler(client *Client, command string, params []string, rb *ResponseBuffer) {
	if len(params) == 0 {
		return
	}
	switch strings.ToLower(params[0]) {
	case "play":
		zncPlaybackPlayHandler(client, command, params, rb)
	case "list":
		zncPlaybackListHandler(client, command, params, rb)
	default:
		return
	}
}

// PRIVMSG *playback :play <target> [lower_bound] [upper_bound]
// e.g., PRIVMSG *playback :play * 1558374442
func zncPlaybackPlayHandler(client *Client, command string, params []string, rb *ResponseBuffer) {
	if len(params) < 2 || len(params) > 4 {
		return
	}
	targetString := params[1]

	now := time.Now().UTC()
	var start, end time.Time
	switch len(params) {
	case 2:
		// #1205: this should have the same semantics as `LATEST *`
	case 3:
		// #831: this should have the same semantics as `LATEST timestamp=qux`,
		// or equivalently `BETWEEN timestamp=$now timestamp=qux`, as opposed to
		// `AFTER timestamp=qux` (this matters in the case where there are
		// more than znc-maxmessages available)
		start = now
		end = zncWireTimeToTime(params[2])
	case 4:
		start = zncWireTimeToTime(params[2])
		end = zncWireTimeToTime(params[3])
	}

	var targets utils.StringSet
	var nickTargets []string

	// three cases:
	// 1. the user's PMs get played back immediately upon receiving this
	// 2. if this is a new connection (from the server's POV), save the information
	// and use it to process subsequent joins. (This is the Textual behavior:
	// first send the playback PRIVMSG, then send the JOIN lines.)
	// 3. if this is a reattach (from the server's POV), immediately play back
	// history for channels that the client is already joined to. In this scenario,
	// there are three total attempts to play the history:
	//     3.1. During the initial reattach (no-op because the *playback privmsg
	//          hasn't been received yet, but they negotiated the znc.in/playback
	//          cap so we know we're going to receive it later)
	//     3.2  Upon receiving the *playback privmsg, i.e., now: we should play
	//          the relevant history lines
	//     3.3  When the client sends a subsequent redundant JOIN line for those
	//          channels; redundant JOIN is a complete no-op so we won't replay twice

	playPrivmsgs := false
	if params[1] == "*" {
		playPrivmsgs = true // XXX nil `targets` means "every channel"
	} else {
		targets = make(utils.StringSet)
		for _, targetName := range strings.Split(targetString, ",") {
			if targetName == "*self" {
				playPrivmsgs = true
			} else if strings.HasPrefix(targetName, "#") {
				if cfTarget, err := CasefoldChannel(targetName); err == nil {
					targets.Add(cfTarget)
				}
			} else {
				if cfNick, err := CasefoldName(targetName); err == nil {
					nickTargets = append(nickTargets, cfNick)
				}
			}
		}
	}

	if playPrivmsgs {
		zncPlayPrivmsgs(client, rb, "*", start, end)
	}

	rb.session.zncPlaybackTimes = &zncPlaybackTimes{
		start:   start,
		end:     end,
		targets: targets,
		setAt:   time.Now().UTC(),
	}

	for _, channel := range client.Channels() {
		if targets == nil || targets.Has(channel.NameCasefolded()) {
			channel.autoReplayHistory(client, rb, "")
			rb.Flush(true)
		}
	}

	for _, cfNick := range nickTargets {
		zncPlayPrivmsgs(client, rb, cfNick, start, end)
		rb.Flush(true)
	}
}

func zncPlayPrivmsgs(client *Client, rb *ResponseBuffer, target string, after, before time.Time) {
	_, sequence, _ := client.server.GetHistorySequence(nil, client, target)
	if sequence == nil {
		return
	}
	zncMax := client.server.Config().History.ZNCMax
	items, _, err := sequence.Between(history.Selector{Time: after}, history.Selector{Time: before}, zncMax)
	if err == nil && len(items) != 0 {
		client.replayPrivmsgHistory(rb, items, "", true)
	}
}

// PRIVMSG *playback :list
func zncPlaybackListHandler(client *Client, command string, params []string, rb *ResponseBuffer) {
	nick := client.Nick()
	for _, channel := range client.Channels() {
		_, sequence, err := client.server.GetHistorySequence(channel, client, "")
		if sequence == nil {
			continue
		} else if err != nil {
			client.server.logger.Error("internal", "couldn't get history sequence for ZNC list", err.Error())
			continue
		}
		items, _, err := sequence.Between(history.Selector{}, history.Selector{}, 1) // i.e., LATEST * 1
		if err != nil {
			client.server.logger.Error("internal", "couldn't query history for ZNC list", err.Error())
		} else if len(items) != 0 {
			stamp := timeToZncWireTime(items[0].Message.Time)
			rb.Add(nil, "*playback!znc@znc.in", "PRIVMSG", nick, fmt.Sprintf("%s 0 %s", channel.Name(), stamp))
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 Volker Dobler.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// html.go contains checks on a HTML body.

package ht

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

	"github.com/andybalholm/cascadia"
	"golang.org/x/net/html"
)

func init() {
	RegisterCheck(&HTMLContains{})
	RegisterCheck(&HTMLContainsText{})
	RegisterCheck(ValidHTML{})
	RegisterCheck(W3CValidHTML{})
	RegisterCheck(&Links{})
}

// ----------------------------------------------------------------------------
// ValidHTML and W3CValidHTML

// ValidHTML checks for valid HTML 5. Kinda: It never fails. TODO: make it useful.
type ValidHTML struct{}

func (c ValidHTML) Execute(t *Test) error {
	if t.Response.BodyErr != nil {
		return BadBody
	}
	_, err := html.Parse(t.Response.Body())
	if err != nil {
		return fmt.Errorf("Invalid HTML: %s", err.Error())
	}

	return nil
}

func (_ ValidHTML) Prepare() error { return nil }

// W3CValidHTML checks for valid HTML but checking the response body via
// the online checker from W3C which is very strict.
type W3CValidHTML struct {
	// AllowedErrors is the number of allowed errors (after ignoring errors).
	AllowedErrors int `json:",omitempty"`

	// IgnoredErrros is a list of error messages to be ignored completely.
	IgnoredErrors []Condition `json:",omitempty"`
}

func (w W3CValidHTML) Execute(t *Test) error {
	file := "@file:@sample.html:" + string(t.Response.BodyBytes)
	test := &Test{
		Name: "W3CValidHTML",
		Request: Request{
			Method: "POST",
			URL:    "http://validator.w3.org/check",
			Params: URLValues{
				"charset":       {"(detect automatically)"},
				"fbc":           {"1"},
				"doctype":       {"Inline"},
				"fbd":           {"1"},
				"group":         {"0"},
				"uploaded_file": {file},
			},
			ParamsAs: "multipart",
			Header: http.Header{
				"Accept":     {"text/html,application/xhtml+xml"},
				"User-Agent": {"Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/41.0.2272.101 Safari/537.36"},
			},
		},
		Checks: CheckList{
			StatusCode{Expect: 200},
		},
		Timeout: Duration(20 * time.Second),
	}

	// TODO: properly limit gloabl rate at which we fire to W3C validator
	time.Sleep(100 * time.Millisecond)

	err := test.Run(nil)
	if err != nil {
		return CantCheck{err}
	}
	if test.Status != Pass {
		return CantCheck{test.Error}
	}

	// Interprete response from validator
	valStat := test.Response.Response.Header.Get("X-W3C-Validator-Status")
	if valStat == "Abort" {
		return CantCheck{fmt.Errorf("validator service sent Abort")}
	} else if valStat == "Valid" {
		return nil
	}

	doc, err := html.Parse(test.Response.Body())
	if err != nil {
		return CantCheck{err}
	}

	// Errors
	validationErrors := []ValidationIssue{}
	noIgnored := 0
outer:
	for _, ve := range w3cValidatorErrSel.MatchAll(doc) {
		vi := extractValidationIssue(ve)
		for _, c := range w.IgnoredErrors {
			if c.Fullfilled(vi.Message) == nil {
				noIgnored++
				continue outer
			}
		}
		validationErrors = append(validationErrors, vi)
	}
	// TODO: maybe warnings too?

	if len(validationErrors) > w.AllowedErrors {
		errmsg := fmt.Sprintf("Ignored %d errors", noIgnored)
		for i, e := range validationErrors {
			errmsg = fmt.Sprintf("%s\nError %d:\n  %s\n  %s\n  %s",
				errmsg, i, e.Position, e.Message, e.Input)
		}
		return fmt.Errorf("%s", errmsg)
	}

	return nil
}

func (_ W3CValidHTML) Prepare() error { return nil }

// ValidationIssue contains extracted information from the output of
// a W3C validator run.
type ValidationIssue struct {
	Position string
	Message  string
	Input    string
}

var (
	w3cValidatorErrSel   = cascadia.MustCompile("li.msg_err")
	w3cValidatorWarnSel  = cascadia.MustCompile("li.msg_warn")
	w3cValidatorEmSel    = cascadia.MustCompile("em")
	w3cValidatorMsgSel   = cascadia.MustCompile("span.msg")
	w3cValidatorInputSel = cascadia.MustCompile("code.input")
)

func extractValidationIssue(node *html.Node) ValidationIssue {
	p := textContent(w3cValidatorEmSel.MatchFirst(node))
	p = strings.Replace(p, "\n     ", "", -1)
	return ValidationIssue{
		Position: p,
		Message:  textContent(w3cValidatorMsgSel.MatchFirst(node)),
		Input:    textContent(w3cValidatorInputSel.MatchFirst(node)),
	}
}

// ----------------------------------------------------------------------------
// HTMLContains

// HTMLContains checks for the existens of HTML elements selected by CSS selectors.
type HTMLContains struct {
	// Selector is the CSS selector of the HTML elements.
	Selector string

	// Count determines the number of occurences to check for:
	//     < 0: no occurence
	//    == 0: one ore more occurences
	//     > 0: exactly that many occurences
	Count int `json:",omitempty"`

	sel cascadia.Selector
}

func (c *HTMLContains) Execute(t *Test) error {
	if c.sel == nil {
		if err := c.Prepare(); err != nil {
			return err
		}
	}
	if t.Response.BodyErr != nil {
		return BadBody
	}

	doc, err := html.Parse(t.Response.Body())
	if err != nil {
		return CantCheck{err}
	}

	matches := c.sel.MatchAll(doc)

	switch {
	case c.Count < 0 && len(matches) > 0:
		return FoundForbidden
	case c.Count == 0 && len(matches) == 0:
		return NotFound
	case c.Count > 0:
		if len(matches) != c.Count {
			return WrongCount{Got: len(matches), Want: c.Count}
		}
	}

	return nil
}

func (c *HTMLContains) Prepare() (err error) {
	c.sel, err = cascadia.Compile(c.Selector)
	if err != nil {
		c.sel = nil
		return MalformedCheck{Err: err}
	}
	return nil
}

// ----------------------------------------------------------------------------
// HTMLContainsText

// HTMLContainsText check the text content off HTML elements selected by a CSS rule.
type HTMLContainsText struct {
	// Selector is the CSS selector of the HTML elements.
	Selector string

	// The plain text content of each selected element.
	Text []string `json:",omitempty"`

	// If true: Text contains the all matches of Selector.
	Complete bool `json:",omitempty"`

	sel cascadia.Selector
}

func (c *HTMLContainsText) Execute(t *Test) error {
	if c.sel == nil {
		if err := c.Prepare(); err != nil {
			return err
		}
	}
	if t.Response.BodyErr != nil {
		return BadBody
	}
	doc, err := html.Parse(t.Response.Body())
	if err != nil {
		return CantCheck{err}
	}

	matches := c.sel.MatchAll(doc)

	for i, want := range c.Text {
		if i == len(matches) {
			return WrongCount{Got: len(matches), Want: len(c.Text)}
		}

		got := textContent(matches[i])
		if want != got {
			return fmt.Errorf("found %q, want %q", got, want)
		}
	}

	if c.Complete && len(c.Text) != len(matches) {
		return WrongCount{Got: len(matches), Want: len(c.Text)}

	}
	return nil
}

func (c *HTMLContainsText) Prepare() (err error) {
	c.sel, err = cascadia.Compile(c.Selector)
	if err != nil {
		c.sel = nil
		return MalformedCheck{Err: err}
	}
	return nil
}

// textContent returns the full text content of n.
func textContent(n *html.Node) string {
	if n == nil {
		return ""
	}
	switch n.Type {
	case html.TextNode:
		return n.Data
	case html.ElementNode:
		s := ""
		for child := n.FirstChild; child != nil; child = child.NextSibling {
			cs := textContent(child)
			if cs != "" {
				if s != "" && s[len(s)-1] != ' ' {
					s += " "
				}
				s += cs
			}
		}
		return s
	}
	return ""
}

// ----------------------------------------------------------------------------
// Links

// Links checks links and references in HTML pages for availability
type Links struct {
	// Head triggers HEAD request instead of GET requests.
	Head bool

	// Which links to test; a combination of "a", "img", "link" and "script".
	// E.g. use "a img" to check the href of all a-tags and src of all img-tags.
	Which string

	// Concurrency determines how many of the found links are checked
	// concurrently. A zero value indicats sequential checking.
	Concurrency int `json:",omitempty"`

	// Timeout if different from main test.
	Timeout Duration `json:",omitempty"`

	IgnoredLinks []Condition `json:",omitempty"`

	tags []string
}

func (c *Links) collectURLs(t *Test) (map[string]struct{}, error) {
	if t.Response.BodyErr != nil {
		return nil, BadBody
	}
	doc, err := html.Parse(t.Response.Body())
	if err != nil {
		return nil, CantCheck{err}
	}

	// Collect all non-ignored URL as map keys (for automatic deduplication).
	refs := make(map[string]struct{})
outer:
	for _, tag := range c.tags {
		for _, u := range c.linkURL(doc, tag, t.Request.Request.URL) {
			for _, cond := range c.IgnoredLinks {
				if cond.Fullfilled(u) == nil {
					continue outer
				}
			}
			refs[u] = struct{}{}
		}
	}

	return refs, nil
}

func (c *Links) Execute(t *Test) error {
	refs, err := c.collectURLs(t)
	if err != nil {
		return err
	}
	suite := &Suite{}
	method := "GET"
	if c.Head {
		method = "HEAD"
	}
	timeout := t.Timeout
	if c.Timeout > 0 {
		timeout = c.Timeout
	}

	for r, _ := range refs {
		test := &Test{
			Name: r,
			Request: Request{
				Method:          method,
				URL:             r,
				FollowRedirects: true,
			},
			ClientPool: t.ClientPool,
			Checks: CheckList{
				StatusCode{Expect: 200},
			},
			Verbosity: t.Verbosity - 1,
			Timeout:   timeout,
		}
		suite.Tests = append(suite.Tests, test)
	}

	err = suite.Prepare()
	if err != nil {
		return CantCheck{fmt.Errorf("Constructed meta test are bad: %s", err)}
	}
	conc := 1
	if c.Concurrency > 1 {
		conc = c.Concurrency
	}
	suite.ExecuteConcurrent(conc)
	if suite.Status != Pass {
		broken := []string{}
		for _, test := range suite.Tests {
			if test.Status == Error || test.Status == Bogus {
				broken = append(broken, fmt.Sprintf("%s  -->  %s",
					test.Request.URL, test.Error))
			} else if test.Status == Fail {
				broken = append(broken, fmt.Sprintf("%s  -->  %d",
					test.Request.URL,
					test.Response.Response.StatusCode))
			}
		}
		return fmt.Errorf("%s", strings.Join(broken, "\n"))
	}
	return nil
}

func (_ Links) linkURL(doc *html.Node, tag string, reqURL *url.URL) []string {
	attr := map[string]string{
		"a":      "href",
		"img":    "src",
		"script": "src",
		"link":   "href",
	}
	href := cascadia.MustCompile(tag)
	matches := href.MatchAll(doc)
	ak := attr[tag]
	refs := []string{}
	for _, m := range matches {
		for _, a := range m.Attr {
			if a.Key != ak || a.Val == "#" ||
				strings.HasPrefix(a.Val, "mailto:") {
				continue
			}
			u, err := reqURL.Parse(a.Val)
			if err != nil {
				refs = append(refs, "Error: "+err.Error())
			} else {
				u.Fragment = "" // easier to clear here
				refs = append(refs, u.String())
			}
		}
	}
	return refs
}

func (c *Links) Prepare() (err error) {
	// TODO: compile IgnoredLinks
	c.tags = nil
	for _, tag := range strings.Split(c.Which, " ") {
		tag = strings.TrimSpace(tag)
		switch tag {
		case "": // ignored
		case "a", "img", "link", "script":
			c.tags = append(c.tags, tag)
		default:
			fmt.Println("Bad", tag)
			return fmt.Errorf("Unknown link tag %q", tag)
		}
	}

	if len(c.tags) == 0 {
		return fmt.Errorf("Bad or missing value for Which: %q", c.Which)
	}
	return nil
}
<commit_msg>ht: Update to new interface of W3C Validator<commit_after>// Copyright 2014 Volker Dobler.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// html.go contains checks on a HTML body.

package ht

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

	"github.com/andybalholm/cascadia"
	"golang.org/x/net/html"
)

func init() {
	RegisterCheck(&HTMLContains{})
	RegisterCheck(&HTMLContainsText{})
	RegisterCheck(ValidHTML{})
	RegisterCheck(W3CValidHTML{})
	RegisterCheck(&Links{})
}

// ----------------------------------------------------------------------------
// ValidHTML and W3CValidHTML

// ValidHTML checks for valid HTML 5. Kinda: It never fails. TODO: make it useful.
type ValidHTML struct{}

func (c ValidHTML) Execute(t *Test) error {
	if t.Response.BodyErr != nil {
		return BadBody
	}
	_, err := html.Parse(t.Response.Body())
	if err != nil {
		return fmt.Errorf("Invalid HTML: %s", err.Error())
	}

	return nil
}

func (_ ValidHTML) Prepare() error { return nil }

// W3CValidHTML checks for valid HTML but checking the response body via
// the online checker from W3C which is very strict.
type W3CValidHTML struct {
	// AllowedErrors is the number of allowed errors (after ignoring errors).
	AllowedErrors int `json:",omitempty"`

	// IgnoredErrros is a list of error messages to be ignored completely.
	IgnoredErrors []Condition `json:",omitempty"`
}

func (w W3CValidHTML) Execute(t *Test) error {
	file := "@file:@sample.html:" + string(t.Response.BodyBytes)
	test := &Test{
		Name: "W3CValidHTML",
		Request: Request{
			Method: "POST",
			URL:    "http://validator.w3.org/nu/",
			Params: URLValues{
				"file": {file},
			},
			ParamsAs: "multipart",
			Header: http.Header{
				"Accept":     {"text/html,application/xhtml+xml"},
				"User-Agent": {"Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/41.0.2272.101 Safari/537.36"},
			},
		},
		Checks: CheckList{
			StatusCode{Expect: 200},
		},
		Timeout: Duration(20 * time.Second),
	}

	// TODO: properly limit gloabl rate at which we fire to W3C validator
	time.Sleep(100 * time.Millisecond)

	err := test.Run(nil)
	if err != nil {
		return CantCheck{err}
	}
	if test.Status != Pass {
		return CantCheck{test.Error}
	}

	// Interprete response from validator
	valStat := test.Response.Response.Header.Get("X-W3C-Validator-Status")
	if valStat == "Abort" {
		return CantCheck{fmt.Errorf("validator service sent Abort")}
	} else if valStat == "Valid" {
		return nil
	}

	doc, err := html.Parse(test.Response.Body())
	if err != nil {
		return CantCheck{err}
	}

	// Errors
	validationErrors := []ValidationIssue{}
	noIgnored := 0
outer:
	for _, ve := range w3cValidatorErrSel.MatchAll(doc) {
		vi := extractValidationIssue(ve)
		for _, c := range w.IgnoredErrors {
			if c.Fullfilled(vi.Message) == nil {
				noIgnored++
				continue outer
			}
		}
		validationErrors = append(validationErrors, vi)
	}
	// TODO: maybe warnings too?

	if len(validationErrors) > w.AllowedErrors {
		errmsg := fmt.Sprintf("Ignored %d errors", noIgnored)
		for i, e := range validationErrors {
			errmsg = fmt.Sprintf("%s\nError %d:\n  %s\n  %s\n  %s",
				errmsg, i, e.Position, e.Message, e.Input)
		}
		return fmt.Errorf("%s", errmsg)
	}

	return nil
}

func (_ W3CValidHTML) Prepare() error { return nil }

// ValidationIssue contains extracted information from the output of
// a W3C validator run.
type ValidationIssue struct {
	Position string
	Message  string
	Input    string
}

var (
	w3cValidatorErrSel   = cascadia.MustCompile("li.error")
	w3cValidatorWarnSel  = cascadia.MustCompile("li.msg_warn")
	w3cValidatorLocSel   = cascadia.MustCompile("p.location")
	w3cValidatorMsgSel   = cascadia.MustCompile("p span")
	w3cValidatorInputSel = cascadia.MustCompile("p.extract")
)

func extractValidationIssue(node *html.Node) ValidationIssue {
	p := textContent(w3cValidatorLocSel.MatchFirst(node))
	p = strings.Replace(p, "\n     ", "", -1)
	return ValidationIssue{
		Position: p,
		Message:  textContent(w3cValidatorMsgSel.MatchFirst(node)),
		Input:    textContent(w3cValidatorInputSel.MatchFirst(node)),
	}
}

// ----------------------------------------------------------------------------
// HTMLContains

// HTMLContains checks for the existens of HTML elements selected by CSS selectors.
type HTMLContains struct {
	// Selector is the CSS selector of the HTML elements.
	Selector string

	// Count determines the number of occurences to check for:
	//     < 0: no occurence
	//    == 0: one ore more occurences
	//     > 0: exactly that many occurences
	Count int `json:",omitempty"`

	sel cascadia.Selector
}

func (c *HTMLContains) Execute(t *Test) error {
	if c.sel == nil {
		if err := c.Prepare(); err != nil {
			return err
		}
	}
	if t.Response.BodyErr != nil {
		return BadBody
	}

	doc, err := html.Parse(t.Response.Body())
	if err != nil {
		return CantCheck{err}
	}

	matches := c.sel.MatchAll(doc)

	switch {
	case c.Count < 0 && len(matches) > 0:
		return FoundForbidden
	case c.Count == 0 && len(matches) == 0:
		return NotFound
	case c.Count > 0:
		if len(matches) != c.Count {
			return WrongCount{Got: len(matches), Want: c.Count}
		}
	}

	return nil
}

func (c *HTMLContains) Prepare() (err error) {
	c.sel, err = cascadia.Compile(c.Selector)
	if err != nil {
		c.sel = nil
		return MalformedCheck{Err: err}
	}
	return nil
}

// ----------------------------------------------------------------------------
// HTMLContainsText

// HTMLContainsText check the text content off HTML elements selected by a CSS rule.
type HTMLContainsText struct {
	// Selector is the CSS selector of the HTML elements.
	Selector string

	// The plain text content of each selected element.
	Text []string `json:",omitempty"`

	// If true: Text contains the all matches of Selector.
	Complete bool `json:",omitempty"`

	sel cascadia.Selector
}

func (c *HTMLContainsText) Execute(t *Test) error {
	if c.sel == nil {
		if err := c.Prepare(); err != nil {
			return err
		}
	}
	if t.Response.BodyErr != nil {
		return BadBody
	}
	doc, err := html.Parse(t.Response.Body())
	if err != nil {
		return CantCheck{err}
	}

	matches := c.sel.MatchAll(doc)

	for i, want := range c.Text {
		if i == len(matches) {
			return WrongCount{Got: len(matches), Want: len(c.Text)}
		}

		got := textContent(matches[i])
		if want != got {
			return fmt.Errorf("found %q, want %q", got, want)
		}
	}

	if c.Complete && len(c.Text) != len(matches) {
		return WrongCount{Got: len(matches), Want: len(c.Text)}

	}
	return nil
}

func (c *HTMLContainsText) Prepare() (err error) {
	c.sel, err = cascadia.Compile(c.Selector)
	if err != nil {
		c.sel = nil
		return MalformedCheck{Err: err}
	}
	return nil
}

// textContent returns the full text content of n.
func textContent(n *html.Node) string {
	if n == nil {
		return ""
	}
	switch n.Type {
	case html.TextNode:
		return n.Data
	case html.ElementNode:
		s := ""
		for child := n.FirstChild; child != nil; child = child.NextSibling {
			cs := textContent(child)
			if cs != "" {
				if s != "" && s[len(s)-1] != ' ' {
					s += " "
				}
				s += cs
			}
		}
		return s
	}
	return ""
}

// ----------------------------------------------------------------------------
// Links

// Links checks links and references in HTML pages for availability
type Links struct {
	// Head triggers HEAD request instead of GET requests.
	Head bool

	// Which links to test; a combination of "a", "img", "link" and "script".
	// E.g. use "a img" to check the href of all a-tags and src of all img-tags.
	Which string

	// Concurrency determines how many of the found links are checked
	// concurrently. A zero value indicats sequential checking.
	Concurrency int `json:",omitempty"`

	// Timeout if different from main test.
	Timeout Duration `json:",omitempty"`

	IgnoredLinks []Condition `json:",omitempty"`

	tags []string
}

func (c *Links) collectURLs(t *Test) (map[string]struct{}, error) {
	if t.Response.BodyErr != nil {
		return nil, BadBody
	}
	doc, err := html.Parse(t.Response.Body())
	if err != nil {
		return nil, CantCheck{err}
	}

	// Collect all non-ignored URL as map keys (for automatic deduplication).
	refs := make(map[string]struct{})
outer:
	for _, tag := range c.tags {
		for _, u := range c.linkURL(doc, tag, t.Request.Request.URL) {
			for _, cond := range c.IgnoredLinks {
				if cond.Fullfilled(u) == nil {
					continue outer
				}
			}
			refs[u] = struct{}{}
		}
	}

	return refs, nil
}

func (c *Links) Execute(t *Test) error {
	refs, err := c.collectURLs(t)
	if err != nil {
		return err
	}
	suite := &Suite{}
	method := "GET"
	if c.Head {
		method = "HEAD"
	}
	timeout := t.Timeout
	if c.Timeout > 0 {
		timeout = c.Timeout
	}

	for r, _ := range refs {
		test := &Test{
			Name: r,
			Request: Request{
				Method:          method,
				URL:             r,
				FollowRedirects: true,
			},
			ClientPool: t.ClientPool,
			Checks: CheckList{
				StatusCode{Expect: 200},
			},
			Verbosity: t.Verbosity - 1,
			Timeout:   timeout,
		}
		suite.Tests = append(suite.Tests, test)
	}

	err = suite.Prepare()
	if err != nil {
		return CantCheck{fmt.Errorf("Constructed meta test are bad: %s", err)}
	}
	conc := 1
	if c.Concurrency > 1 {
		conc = c.Concurrency
	}
	suite.ExecuteConcurrent(conc)
	if suite.Status != Pass {
		broken := []string{}
		for _, test := range suite.Tests {
			if test.Status == Error || test.Status == Bogus {
				broken = append(broken, fmt.Sprintf("%s  -->  %s",
					test.Request.URL, test.Error))
			} else if test.Status == Fail {
				broken = append(broken, fmt.Sprintf("%s  -->  %d",
					test.Request.URL,
					test.Response.Response.StatusCode))
			}
		}
		return fmt.Errorf("%s", strings.Join(broken, "\n"))
	}
	return nil
}

func (_ Links) linkURL(doc *html.Node, tag string, reqURL *url.URL) []string {
	attr := map[string]string{
		"a":      "href",
		"img":    "src",
		"script": "src",
		"link":   "href",
	}
	href := cascadia.MustCompile(tag)
	matches := href.MatchAll(doc)
	ak := attr[tag]
	refs := []string{}
	for _, m := range matches {
		for _, a := range m.Attr {
			if a.Key != ak || a.Val == "#" ||
				strings.HasPrefix(a.Val, "mailto:") {
				continue
			}
			u, err := reqURL.Parse(a.Val)
			if err != nil {
				refs = append(refs, "Error: "+err.Error())
			} else {
				u.Fragment = "" // easier to clear here
				refs = append(refs, u.String())
			}
		}
	}
	return refs
}

func (c *Links) Prepare() (err error) {
	// TODO: compile IgnoredLinks
	c.tags = nil
	for _, tag := range strings.Split(c.Which, " ") {
		tag = strings.TrimSpace(tag)
		switch tag {
		case "": // ignored
		case "a", "img", "link", "script":
			c.tags = append(c.tags, tag)
		default:
			fmt.Println("Bad", tag)
			return fmt.Errorf("Unknown link tag %q", tag)
		}
	}

	if len(c.tags) == 0 {
		return fmt.Errorf("Bad or missing value for Which: %q", c.Which)
	}
	return nil
}
<|endoftext|>


<commit_before>package vfs

import (
	"bytes"
	"crypto/md5" // #nosec
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"path"
	"time"

	"github.com/cozy/cozy-stack/couchdb"
	"github.com/spf13/afero"
)

type fileAttributes struct {
	Name       string    `json:"name"`
	CreatedAt  time.Time `json:"created_at"`
	UpdatedAt  time.Time `json:"updated_at"`
	Size       int64     `json:"size,string"`
	Tags       []string  `json:"tags"`
	MD5Sum     []byte    `json:"md5sum"`
	Executable bool      `json:"executable"`
	Class      string    `json:"class"`
	Mime       string    `json:"mime"`
}

// FileDoc is a struct containing all the informations about a file.
// It implements the couchdb.Doc and jsonapi.JSONApier interfaces.
type FileDoc struct {
	// Qualified file identifier
	FID string `json:"_id,omitempty"`
	// File revision
	FRev string `json:"_rev,omitempty"`
	// File attributes
	Attrs *fileAttributes `json:"attributes"`
	// Parent folder identifier
	FolderID string `json:"folderID"`
	// File path on VFS
	Path string `json:"path"`
}

// ID returns the file qualified identifier (part of couchdb.Doc
// interface)
func (f *FileDoc) ID() string {
	return f.FID
}

// Rev returns the file revision (part of couchdb.Doc interface)
func (f *FileDoc) Rev() string {
	return f.FRev
}

// DocType returns the file document type (part of couchdb.Doc
// interface)
func (f *FileDoc) DocType() string {
	return string(FileDocType)
}

// SetID is used to change the file qualified identifier (part of
// couchdb.Doc interface)
func (f *FileDoc) SetID(id string) {
	f.FID = id
}

// SetRev is used to change the file revision (part of couchdb.Doc
// interface)
func (f *FileDoc) SetRev(rev string) {
	f.FRev = rev
}

// ToJSONApi implements temporary interface JSONApier to serialize
// the file document
func (f *FileDoc) ToJSONApi() ([]byte, error) {
	id := f.FID
	data := map[string]interface{}{
		"id":         id,
		"type":       f.DocType(),
		"rev":        f.Rev(),
		"attributes": f.Attrs,
	}
	m := map[string]interface{}{
		"data": data,
	}
	return json.Marshal(m)
}

// GetFileDoc is used to fetch file document information form our
// database.
func GetFileDoc(fileID, dbPrefix string) (doc *FileDoc, err error) {
	doc = &FileDoc{}
	err = couchdb.GetDoc(dbPrefix, string(FileDocType), fileID, doc)
	return
}

// ServeFileContent replies to a http request using the content of a
// file given its FileDoc.
//
// It uses internally http.ServeContent and benefits from it by
// offering support to Range, If-Modified-Since and If-None-Match
// requests. It uses the revision of the file as the Etag value for
// non-ranged requests
//
// The content disposition is inlined.
func ServeFileContent(fileDoc *FileDoc, req *http.Request, w http.ResponseWriter, fs afero.Fs) (err error) {
	attrs := fileDoc.Attrs
	header := w.Header()
	header.Set("Content-Type", attrs.Mime)
	header.Set("Content-Disposition", "inline; filename="+attrs.Name)

	if header.Get("Range") == "" {
		eTag := base64.StdEncoding.EncodeToString(fileDoc.Attrs.MD5Sum)
		header.Set("Etag", eTag)
	}

	serveContent(req, w, fs, fileDoc.Path, attrs.Name, attrs.UpdatedAt)
	return
}

// ServeFileContentByPath replies to a http request using the content
// of a file identified by its full path on the VFS. Unlike
// ServeFileContent, this method does not require the full file
// document but only its path.
//
// It also uses internally http.ServeContent but does not provide an
// Etag.
//
// The content disposition is attached
func ServeFileContentByPath(pth string, req *http.Request, w http.ResponseWriter, fs afero.Fs) (err error) {
	fileInfo, err := fs.Stat(pth)
	if err != nil {
		return
	}

	name := path.Base(pth)
	w.Header().Set("Content-Disposition", "attachment; filename="+name)

	serveContent(req, w, fs, pth, name, fileInfo.ModTime())
	return
}

func serveContent(req *http.Request, w http.ResponseWriter, fs afero.Fs, pth, name string, modtime time.Time) (err error) {
	content, err := fs.Open(pth)
	if err != nil {
		return
	}

	defer content.Close()
	http.ServeContent(w, req, name, modtime, content)
	return
}

// CreateFileAndUpload is the method for uploading a file onto the filesystem.
func CreateFileAndUpload(m *DocAttributes, fs afero.Fs, dbPrefix string, body io.ReadCloser) (doc *FileDoc, err error) {
	if m.docType != FileDocType {
		err = ErrDocTypeInvalid
		return
	}

	pth, _, err := createNewFilePath(m, fs, dbPrefix)
	if err != nil {
		return
	}

	createDate := time.Now()
	attrs := &fileAttributes{
		Name:       m.name,
		CreatedAt:  createDate,
		UpdatedAt:  createDate,
		Size:       m.size,
		Tags:       m.tags,
		MD5Sum:     m.givenMD5,
		Executable: m.executable,
		Class:      m.class,
		Mime:       m.mime,
	}

	doc = &FileDoc{
		Attrs:    attrs,
		FolderID: m.folderID,
		Path:     pth,
	}

	// Error handling to make sure the steps of uploading the file and
	// creating the corresponding are both rollbacked in case of an
	// error. This should preserve our VFS coherency a little.
	defer func() {
		if err != nil {
			fs.Remove(pth)
		}
	}()

	var written int64
	if written, err = copyOnFsAndCheckIntegrity(m, fs, pth, body); err != nil {
		return
	}

	if attrs.Size < 0 {
		attrs.Size = written
	}

	if attrs.Size != written {
		err = ErrContentLengthMismatch
		return
	}

	if err = couchdb.CreateDoc(dbPrefix, doc); err != nil {
		return
	}

	return
}

func copyOnFsAndCheckIntegrity(m *DocAttributes, fs afero.Fs, pth string, r io.ReadCloser) (written int64, err error) {
	f, err := fs.Create(pth)
	if err != nil {
		return
	}

	defer f.Close()
	defer r.Close()

	md5H := md5.New() // #nosec
	written, err = io.Copy(f, io.TeeReader(r, md5H))
	if err != nil {
		return
	}

	calcMD5 := md5H.Sum(nil)
	if !bytes.Equal(m.givenMD5, calcMD5) {
		err = ErrInvalidHash
		return
	}

	return
}
<commit_msg>Allow to not pass a Content-MD5 for upload<commit_after>package vfs

import (
	"bytes"
	"crypto/md5" // #nosec
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"path"
	"time"

	"github.com/cozy/cozy-stack/couchdb"
	"github.com/spf13/afero"
)

type fileAttributes struct {
	Name       string    `json:"name"`
	CreatedAt  time.Time `json:"created_at"`
	UpdatedAt  time.Time `json:"updated_at"`
	Size       int64     `json:"size,string"`
	Tags       []string  `json:"tags"`
	MD5Sum     []byte    `json:"md5sum"`
	Executable bool      `json:"executable"`
	Class      string    `json:"class"`
	Mime       string    `json:"mime"`
}

// FileDoc is a struct containing all the informations about a file.
// It implements the couchdb.Doc and jsonapi.JSONApier interfaces.
type FileDoc struct {
	// Qualified file identifier
	FID string `json:"_id,omitempty"`
	// File revision
	FRev string `json:"_rev,omitempty"`
	// File attributes
	Attrs *fileAttributes `json:"attributes"`
	// Parent folder identifier
	FolderID string `json:"folderID"`
	// File path on VFS
	Path string `json:"path"`
}

// ID returns the file qualified identifier (part of couchdb.Doc
// interface)
func (f *FileDoc) ID() string {
	return f.FID
}

// Rev returns the file revision (part of couchdb.Doc interface)
func (f *FileDoc) Rev() string {
	return f.FRev
}

// DocType returns the file document type (part of couchdb.Doc
// interface)
func (f *FileDoc) DocType() string {
	return string(FileDocType)
}

// SetID is used to change the file qualified identifier (part of
// couchdb.Doc interface)
func (f *FileDoc) SetID(id string) {
	f.FID = id
}

// SetRev is used to change the file revision (part of couchdb.Doc
// interface)
func (f *FileDoc) SetRev(rev string) {
	f.FRev = rev
}

// ToJSONApi implements temporary interface JSONApier to serialize
// the file document
func (f *FileDoc) ToJSONApi() ([]byte, error) {
	id := f.FID
	data := map[string]interface{}{
		"id":         id,
		"type":       f.DocType(),
		"rev":        f.Rev(),
		"attributes": f.Attrs,
	}
	m := map[string]interface{}{
		"data": data,
	}
	return json.Marshal(m)
}

// GetFileDoc is used to fetch file document information form our
// database.
func GetFileDoc(fileID, dbPrefix string) (doc *FileDoc, err error) {
	doc = &FileDoc{}
	err = couchdb.GetDoc(dbPrefix, string(FileDocType), fileID, doc)
	return
}

// ServeFileContent replies to a http request using the content of a
// file given its FileDoc.
//
// It uses internally http.ServeContent and benefits from it by
// offering support to Range, If-Modified-Since and If-None-Match
// requests. It uses the revision of the file as the Etag value for
// non-ranged requests
//
// The content disposition is inlined.
func ServeFileContent(fileDoc *FileDoc, req *http.Request, w http.ResponseWriter, fs afero.Fs) (err error) {
	attrs := fileDoc.Attrs
	header := w.Header()
	header.Set("Content-Type", attrs.Mime)
	header.Set("Content-Disposition", "inline; filename="+attrs.Name)

	if header.Get("Range") == "" {
		eTag := base64.StdEncoding.EncodeToString(fileDoc.Attrs.MD5Sum)
		header.Set("Etag", eTag)
	}

	serveContent(req, w, fs, fileDoc.Path, attrs.Name, attrs.UpdatedAt)
	return
}

// ServeFileContentByPath replies to a http request using the content
// of a file identified by its full path on the VFS. Unlike
// ServeFileContent, this method does not require the full file
// document but only its path.
//
// It also uses internally http.ServeContent but does not provide an
// Etag.
//
// The content disposition is attached
func ServeFileContentByPath(pth string, req *http.Request, w http.ResponseWriter, fs afero.Fs) (err error) {
	fileInfo, err := fs.Stat(pth)
	if err != nil {
		return
	}

	name := path.Base(pth)
	w.Header().Set("Content-Disposition", "attachment; filename="+name)

	serveContent(req, w, fs, pth, name, fileInfo.ModTime())
	return
}

func serveContent(req *http.Request, w http.ResponseWriter, fs afero.Fs, pth, name string, modtime time.Time) (err error) {
	content, err := fs.Open(pth)
	if err != nil {
		return
	}

	defer content.Close()
	http.ServeContent(w, req, name, modtime, content)
	return
}

// CreateFileAndUpload is the method for uploading a file onto the filesystem.
func CreateFileAndUpload(m *DocAttributes, fs afero.Fs, dbPrefix string, body io.ReadCloser) (doc *FileDoc, err error) {
	if m.docType != FileDocType {
		err = ErrDocTypeInvalid
		return
	}

	pth, _, err := createNewFilePath(m, fs, dbPrefix)
	if err != nil {
		return
	}

	createDate := time.Now()
	attrs := &fileAttributes{
		Name:       m.name,
		CreatedAt:  createDate,
		UpdatedAt:  createDate,
		Size:       m.size,
		Tags:       m.tags,
		MD5Sum:     m.givenMD5,
		Executable: m.executable,
		Class:      m.class,
		Mime:       m.mime,
	}

	doc = &FileDoc{
		Attrs:    attrs,
		FolderID: m.folderID,
		Path:     pth,
	}

	// Error handling to make sure the steps of uploading the file and
	// creating the corresponding are both rollbacked in case of an
	// error. This should preserve our VFS coherency a little.
	defer func() {
		if err != nil {
			fs.Remove(pth)
		}
	}()

	var written int64
	var md5Sum []byte
	if written, md5Sum, err = copyOnFsAndCheckIntegrity(m, fs, pth, body); err != nil {
		return
	}

	if attrs.Size < 0 {
		attrs.Size = written
	}

	if attrs.MD5Sum == nil {
		attrs.MD5Sum = md5Sum
	}

	if attrs.Size != written {
		err = ErrContentLengthMismatch
		return
	}

	if err = couchdb.CreateDoc(dbPrefix, doc); err != nil {
		return
	}

	return
}

func copyOnFsAndCheckIntegrity(m *DocAttributes, fs afero.Fs, pth string, r io.ReadCloser) (written int64, md5Sum []byte, err error) {
	f, err := fs.Create(pth)
	if err != nil {
		return
	}

	defer f.Close()
	defer r.Close()

	md5H := md5.New() // #nosec

	written, err = io.Copy(f, io.TeeReader(r, md5H))
	if err != nil {
		return
	}

	doCheck := m.givenMD5 != nil
	md5Sum = md5H.Sum(nil)
	if doCheck && !bytes.Equal(m.givenMD5, md5Sum) {
		err = ErrInvalidHash
		return
	}

	return
}
<|endoftext|>


<commit_before>package hub

import (
	"code.google.com/p/go.net/websocket"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"sync"
	"time"
)

const (
	defaultBufferSize     = 1024
	defaultHeartbeat      = time.Second * 5
	defaultSessionTimeout = time.Second * 30
	idLength              = 16
	bufLength             = 4096
	defaultLoglevel       = 1
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

type Authorizer func(uri, token string, write bool) (authorized bool, err error)

func prettify(obj interface{}) string {
	b, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		panic(err)
	}
	return string(b)
}

type Server struct {
	heartbeat      time.Duration
	sessionTimeout time.Duration
	sessions       map[string]*Session
	subscriptions  map[string]map[string]*Session
	subscribers    map[string]map[string]bool
	lock           *sync.RWMutex
	loglevel       int
	bufferSize     int
	logger         *log.Logger
	authorizer     Authorizer
}

func NewServer() *Server {
	return &Server{
		heartbeat:      defaultHeartbeat,
		bufferSize:     defaultBufferSize,
		sessionTimeout: defaultSessionTimeout,
		loglevel:       defaultLoglevel,
		sessions:       map[string]*Session{},
		subscriptions:  map[string]map[string]*Session{},
		subscribers:    map[string]map[string]bool{},
		lock:           &sync.RWMutex{},
		logger:         log.New(os.Stdout, "pusher: ", 0),
	}
}

func (self *Server) Loglevel(i int) *Server {
	self.loglevel = i
	return self
}

func (self *Server) Logger(l *log.Logger) *Server {
	self.logger = l
	return self
}

func (self *Server) Fatalf(fmt string, i ...interface{}) {
	self.logger.Printf(fmt, i...)
}

func (self *Server) Errorf(fmt string, i ...interface{}) {
	if self.loglevel > 0 {
		self.logger.Printf(fmt, i...)
	}
}

func (self *Server) Infof(fmt string, i ...interface{}) {
	if self.loglevel > 1 {
		self.logger.Printf(fmt, i...)
	}
}

func (self *Server) Debugf(fmt string, i ...interface{}) {
	if self.loglevel > 2 {
		self.logger.Printf(fmt, i...)
	}
}

func (self *Server) Authorizer(f Authorizer) *Server {
	self.authorizer = f
	return self
}

func (self *Server) addSubscription(sess *Session, uri string) {
	self.lock.Lock()
	defer self.lock.Unlock()

	if _, found := self.subscriptions[uri]; !found {
		self.subscriptions[uri] = map[string]*Session{}
	}
	self.subscriptions[uri][sess.id] = sess

	if _, found := self.subscribers[sess.id]; !found {
		self.subscribers[sess.id] = map[string]bool{}
	}
	self.subscribers[sess.id][uri] = true

	self.Infof("%v\t-\t%v\t%v\t%v\t[subscribe]", time.Now(), uri, sess.RemoteAddr, sess.id)
}

func (self *Server) Emit(message Message) {
	self.lock.Lock()
	defer self.lock.Unlock()

	sent := 0
	for _, sess := range self.subscriptions[message.URI] {
		sess.send(message)
		sent++
	}
	self.Debugf("%v\t%v\t[emitted] x %v", time.Now(), message.URI, sent)
}

func (self *Server) removeSubscription(id, uri string, withLocking bool) {
	if withLocking {
		self.lock.Lock()
		defer self.lock.Unlock()
	}

	delete(self.subscriptions[uri], id)
	if len(self.subscriptions[uri]) == 0 {
		delete(self.subscriptions, uri)
	}

	delete(self.subscribers[id], uri)
	if len(self.subscribers[id]) == 0 {
		delete(self.subscribers, id)
	}
	self.Infof("%v\t-\t%v\t%v\t-\t[unsubscribe]", time.Now(), uri, id)
}

func (self *Server) randomId() string {
	buf := make([]byte, idLength)
	for index, _ := range buf {
		buf[index] = byte(rand.Int31())
	}
	return (base64.URLEncoding.EncodeToString(buf))
}

func (self *Server) removeSession(id string) {
	self.lock.Lock()
	defer self.lock.Unlock()

	delete(self.sessions, id)

	for uri, _ := range self.subscribers[id] {
		self.removeSubscription(id, uri, false)
	}
	self.Infof("%v\t-\t[cleanup]\t%v", time.Now(), id)
}

func (self *Server) GetSession(id string) (result *Session) {
	self.lock.Lock()
	defer self.lock.Unlock()

	// if the id is not found (because the server restarted, or the client gave the initial empty id) we just create a new id and insert a new session
	if result = self.sessions[id]; result == nil {
		result = &Session{
			output:         make(chan Message, self.bufferSize),
			id:             self.randomId(),
			server:         self,
			authorizations: map[string]bool{},
			lock:           &sync.RWMutex{},
		}
		result.cleanupTimer = &time.Timer{}
		self.sessions[result.id] = result
	}
	return
}

func (self *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	self.Infof("%v\t%v\t%v\t%v", time.Now(), r.Method, r.URL, r.RemoteAddr)
	websocket.Handler(self.GetSession(r.URL.Query().Get("session_id")).handleWS).ServeHTTP(w, r)
}

type MessageType string
type ErrorType string

const (
	TypeError       = "Error"
	TypeHeartbeat   = "Heartbeat"
	TypeWelcome     = "Welcome"
	TypeSubscribe   = "Subscribe"
	TypeUnsubscribe = "Unsubscribe"
	TypeMessage     = "Message"
	TypeAuthorize   = "Authorize"
	TypeAck         = "Ack"
)

const (
	TypeJSONError          = "JSONError"
	TypeAuthorizationError = "AuthorizationError"
	TypeSyntaxError        = "SyntaxError"
)

type Welcome struct {
	Heartbeat      time.Duration
	SessionTimeout time.Duration
	Id             string
}

type Error struct {
	Message string
	Type    ErrorType
}

type Message struct {
	Type    MessageType
	Id      string      `json:",omitempty"`
	Welcome *Welcome    `json:",omitempty"`
	Error   *Error      `json:",omitempty"`
	Data    interface{} `json:",omitempty"`
	URI     string      `json:",omitempty"`
	Token   string      `json:",omitempty"`
	Write   bool        `json:",omitempty"`
}

type Session struct {
	ws             io.ReadWriteCloser
	id             string
	RemoteAddr     string
	input          chan Message
	output         chan Message
	closing        chan struct{}
	server         *Server
	cleanupTimer   *time.Timer
	authorizations map[string]bool
	lock           *sync.RWMutex
}

func (self *Session) parseMessage(b []byte) (result Message, err error) {
	err = json.Unmarshal(b, &result)
	return
}

func (self *Session) readLoop() {
	defer self.terminate()
	buf := make([]byte, bufLength)
	n, err := self.ws.Read(buf)
	for err == nil {
		buf_copy := make([]byte, n)
		copy(buf_copy, buf[:n])
		if message, err := self.parseMessage(buf_copy); err == nil {
			self.input <- message
			self.server.Debugf("%v\t%v\t%v\t%v\t[received from socket]", time.Now(), message.URI, self.RemoteAddr, self.id)
		} else {
			self.send(Message{
				Type: TypeError,
				Error: &Error{
					Message: err.Error(),
					Type:    TypeJSONError,
				},
				Data: string(buf_copy)})
		}
		n, err = self.ws.Read(buf)
	}
}

func (self *Session) writeLoop() {
	defer self.terminate()
	var message Message
	var err error
	var n int
	var encoded []byte
	for {
		select {
		case message = <-self.output:
			if encoded, err = json.Marshal(message); err == nil {
				if n, err = self.ws.Write(encoded); err != nil {
					self.server.Fatalf("Error sending %s on %+v: %v", encoded, self.ws, err)
					return
				} else if n != len(encoded) {
					self.server.Fatalf("Unable to send all of %s on %+v: only sent %v bytes", encoded, self.ws, n)
					return
				}
			} else {
				self.server.Fatalf("Unable to JSON marshal %+v: %v", message, err)
				return
			}
			self.server.Debugf("%v\t%v\t%v\t%v\t[sent to socket]", time.Now(), message.URI, self.RemoteAddr, self.id)
		case <-self.closing:
			return
		}
	}
}

func (self *Session) heartbeatLoop() {
	for {
		select {
		case <-self.closing:
			return
		case <-time.After(self.server.heartbeat / 2):
			self.send(Message{Type: TypeHeartbeat})
		}
	}
}

func (self *Session) authorized(uri string, wantWrite bool) bool {
	if self.server.authorizer == nil {
		return true
	}
	self.lock.RLock()
	defer self.lock.RUnlock()
	hasWrite, found := self.authorizations[uri]
	if !found {
		return false
	}
	return !wantWrite || hasWrite
}

func (self *Session) send(message Message) {
	select {
	case self.output <- message:
	default:
		self.server.Errorf("Unable to send %+v to %+v, output buffer full", message, self)
	}
}

func (self *Session) handleMessage(message Message) {
	switch message.Type {
	case TypeHeartbeat:
	case TypeMessage:
		if !self.authorized(message.URI, true) {
			self.send(Message{
				Type: TypeError,
				Id:   message.Id,
				Error: &Error{
					Message: fmt.Sprintf("%v not authorized for writing to %v", self.id, message.URI),
					Type:    TypeAuthorizationError,
				},
				Data: message,
			})
			return
		}
		self.server.Emit(message)
	case TypeUnsubscribe:
		self.server.removeSubscription(self.id, message.URI, true)
	case TypeSubscribe:
		if !self.authorized(message.URI, false) {
			self.send(Message{
				Type: TypeError,
				Id:   message.Id,
				Error: &Error{
					Message: fmt.Sprintf("%v not authorized for subscribing to %v", self.id, message.URI),
					Type:    TypeAuthorizationError,
				},
				Data: message,
			})
			return
		}
		self.server.addSubscription(self, message.URI)
	case TypeAuthorize:
		if self.server.authorizer != nil {
			ok, err := self.server.authorizer(message.URI, message.Token, message.Write)
			if err != nil {
				self.send(Message{
					Type: TypeError,
					Id:   message.Id,
					Error: &Error{
						Message: err.Error(),
						Type:    TypeAuthorizationError,
					}, Data: message,
				})
				return
			}
			if !ok {
				self.send(Message{
					Type: TypeError,
					Id:   message.Id,
					Error: &Error{
						Message: fmt.Sprintf("%v does not provide authorization for %v", message.Token, message.URI),
						Type:    TypeAuthorizationError,
					}, Data: message,
				})
				return
			}
		}
		self.lock.Lock()
		defer self.lock.Unlock()
		self.authorizations[message.URI] = (self.authorizations[message.URI] || message.Write)
		self.server.Infof("%v\t-\t[authorize]\t%v\t%v\t%v", time.Now(), self.RemoteAddr, self.id, message.URI)
	default:
		self.send(Message{
			Type: TypeError,
			Id:   message.Id,
			Error: &Error{
				Message: fmt.Sprintf("Unknown message type %#v", message.Type),
				Type:    TypeSyntaxError,
			},
			Data: message,
		})
		return
	}
	if message.Id != "" {
		self.send(Message{
			Type: TypeAck,
			Id:   message.Id,
		})
	}

}

func (self *Session) remove() {
	self.server.removeSession(self.id)
}

func (self *Session) terminate() {
	self.ws.Close()
	select {
	case _ = <-self.closing:
	default:
		close(self.closing)
	}
	self.server.Infof("%v\t-\t[disconnect]\t%v\t%v", time.Now(), self.RemoteAddr, self.id)
	self.cleanupTimer.Stop()
	self.cleanupTimer = time.AfterFunc(self.server.sessionTimeout, self.remove)
}

func (self *Session) handleWS(ws *websocket.Conn) {
	self.RemoteAddr = ws.Request().RemoteAddr
	self.Handle(ws)
}

/*
 * This function will work as the session main-loop listening on events
 * on a websocket or other source that interfaces io.ReadWriteCloser.
 */
func (self *Session) Handle(ws io.ReadWriteCloser) {
	self.server.Infof("%v\t-\t[connect]\t%v\t%v", time.Now(), self.RemoteAddr, self.id)

	defer self.terminate()

	self.ws = ws
	self.cleanupTimer.Stop()
	self.input = make(chan Message)
	self.closing = make(chan struct{})

	go self.readLoop()
	go self.writeLoop()
	go self.heartbeatLoop()

	self.send(Message{
		Type: TypeWelcome,
		Welcome: &Welcome{
			Heartbeat:      self.server.heartbeat / time.Millisecond,
			SessionTimeout: self.server.sessionTimeout / time.Millisecond,
			Id:             self.id,
		},
	})

	var message Message
	for {
		select {
		case _ = <-self.closing:
			return
		case message = <-self.input:
			self.handleMessage(message)
		case <-time.After(self.server.heartbeat):
			return
		}
	}
}
<commit_msg>Seems last fix was only semi-true.<commit_after>package hub

import (
	"code.google.com/p/go.net/websocket"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"sync"
	"time"
)

const (
	defaultBufferSize     = 1024
	defaultHeartbeat      = time.Second * 5
	defaultSessionTimeout = time.Second * 30
	idLength              = 16
	bufLength             = 4096
	defaultLoglevel       = 1
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

type Authorizer func(uri, token string, write bool) (authorized bool, err error)

func prettify(obj interface{}) string {
	b, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		panic(err)
	}
	return string(b)
}

type Server struct {
	heartbeat      time.Duration
	sessionTimeout time.Duration
	sessions       map[string]*Session
	subscriptions  map[string]map[string]*Session
	subscribers    map[string]map[string]bool
	lock           *sync.RWMutex
	loglevel       int
	bufferSize     int
	logger         *log.Logger
	authorizer     Authorizer
}

func NewServer() *Server {
	return &Server{
		heartbeat:      defaultHeartbeat,
		bufferSize:     defaultBufferSize,
		sessionTimeout: defaultSessionTimeout,
		loglevel:       defaultLoglevel,
		sessions:       map[string]*Session{},
		subscriptions:  map[string]map[string]*Session{},
		subscribers:    map[string]map[string]bool{},
		lock:           &sync.RWMutex{},
		logger:         log.New(os.Stdout, "pusher: ", 0),
	}
}

func (self *Server) Loglevel(i int) *Server {
	self.loglevel = i
	return self
}

func (self *Server) Logger(l *log.Logger) *Server {
	self.logger = l
	return self
}

func (self *Server) Fatalf(fmt string, i ...interface{}) {
	self.logger.Printf(fmt, i...)
}

func (self *Server) Errorf(fmt string, i ...interface{}) {
	if self.loglevel > 0 {
		self.logger.Printf(fmt, i...)
	}
}

func (self *Server) Infof(fmt string, i ...interface{}) {
	if self.loglevel > 1 {
		self.logger.Printf(fmt, i...)
	}
}

func (self *Server) Debugf(fmt string, i ...interface{}) {
	if self.loglevel > 2 {
		self.logger.Printf(fmt, i...)
	}
}

func (self *Server) Authorizer(f Authorizer) *Server {
	self.authorizer = f
	return self
}

func (self *Server) addSubscription(sess *Session, uri string) {
	self.lock.Lock()
	defer self.lock.Unlock()

	if _, found := self.subscriptions[uri]; !found {
		self.subscriptions[uri] = map[string]*Session{}
	}
	self.subscriptions[uri][sess.id] = sess

	if _, found := self.subscribers[sess.id]; !found {
		self.subscribers[sess.id] = map[string]bool{}
	}
	self.subscribers[sess.id][uri] = true

	self.Infof("%v\t-\t%v\t%v\t%v\t[subscribe]", time.Now(), uri, sess.RemoteAddr, sess.id)
}

func (self *Server) Emit(message Message) {
	self.lock.Lock()
	defer self.lock.Unlock()

	sent := 0
	for _, sess := range self.subscriptions[message.URI] {
		sess.send(message)
		sent++
	}
	self.Debugf("%v\t%v\t[emitted] x %v", time.Now(), message.URI, sent)
}

func (self *Server) removeSubscription(id, uri string, withLocking bool) {
	if withLocking {
		self.lock.Lock()
		defer self.lock.Unlock()
	}

	delete(self.subscriptions[uri], id)
	if len(self.subscriptions[uri]) == 0 {
		delete(self.subscriptions, uri)
	}

	delete(self.subscribers[id], uri)
	if len(self.subscribers[id]) == 0 {
		delete(self.subscribers, id)
	}
	self.Infof("%v\t-\t%v\t%v\t-\t[unsubscribe]", time.Now(), uri, id)
}

func (self *Server) randomId() string {
	buf := make([]byte, idLength)
	for index, _ := range buf {
		buf[index] = byte(rand.Int31())
	}
	return (base64.URLEncoding.EncodeToString(buf))
}

func (self *Server) removeSession(id string) {
	self.lock.Lock()
	defer self.lock.Unlock()

	delete(self.sessions, id)

	for uri, _ := range self.subscribers[id] {
		self.removeSubscription(id, uri, false)
	}
	self.Infof("%v\t-\t[cleanup]\t%v", time.Now(), id)
}

func (self *Server) GetSession(id string) (result *Session) {
	self.lock.Lock()
	defer self.lock.Unlock()

	// if the id is not found (because the server restarted, or the client gave the initial empty id) we just create a new id and insert a new session
	if result = self.sessions[id]; result == nil {
		result = &Session{
			output:         make(chan Message, self.bufferSize),
			id:             self.randomId(),
			server:         self,
			authorizations: map[string]bool{},
			lock:           &sync.RWMutex{},
		}
		result.cleanupTimer = &time.Timer{}
		self.sessions[result.id] = result
	}
	return
}

func (self *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	self.Infof("%v\t%v\t%v\t%v", time.Now(), r.Method, r.URL, r.RemoteAddr)
	websocket.Handler(self.GetSession(r.URL.Query().Get("session_id")).handleWS).ServeHTTP(w, r)
}

type MessageType string
type ErrorType string

const (
	TypeError       = "Error"
	TypeHeartbeat   = "Heartbeat"
	TypeWelcome     = "Welcome"
	TypeSubscribe   = "Subscribe"
	TypeUnsubscribe = "Unsubscribe"
	TypeMessage     = "Message"
	TypeAuthorize   = "Authorize"
	TypeAck         = "Ack"
)

const (
	TypeJSONError          = "JSONError"
	TypeAuthorizationError = "AuthorizationError"
	TypeSyntaxError        = "SyntaxError"
)

type Welcome struct {
	Heartbeat      time.Duration
	SessionTimeout time.Duration
	Id             string
}

type Error struct {
	Message string
	Type    ErrorType
}

type Message struct {
	Type    MessageType
	Id      string      `json:",omitempty"`
	Welcome *Welcome    `json:",omitempty"`
	Error   *Error      `json:",omitempty"`
	Data    interface{} `json:",omitempty"`
	URI     string      `json:",omitempty"`
	Token   string      `json:",omitempty"`
	Write   bool        `json:",omitempty"`
}

type Session struct {
	ws             io.ReadWriteCloser
	id             string
	RemoteAddr     string
	input          chan Message
	output         chan Message
	closing        chan struct{}
	server         *Server
	cleanupTimer   *time.Timer
	authorizations map[string]bool
	lock           *sync.RWMutex
}

func (self *Session) parseMessage(b []byte) (result Message, err error) {
	err = json.Unmarshal(b, &result)
	return
}

func (self *Session) readLoop() {
	defer self.terminate()
	buf := make([]byte, bufLength)
	n, err := self.ws.Read(buf)
	for err == nil {
		if message, err := self.parseMessage(buf[:n]); err == nil {
			self.input <- message
			self.server.Debugf("%v\t%v\t%v\t%v\t[received from socket]", time.Now(), message.URI, self.RemoteAddr, self.id)
		} else {
			buf_copy := make([]byte, n)
			copy(buf_copy, buf[:n])
			self.send(Message{
				Type: TypeError,
				Error: &Error{
					Message: err.Error(),
					Type:    TypeJSONError,
				},
				Data: string(buf_copy)})
		}
		n, err = self.ws.Read(buf)
	}
}

func (self *Session) writeLoop() {
	defer self.terminate()
	var message Message
	var err error
	var n int
	var encoded []byte
	for {
		select {
		case message = <-self.output:
			if encoded, err = json.Marshal(message); err == nil {
				if n, err = self.ws.Write(encoded); err != nil {
					self.server.Fatalf("Error sending %s on %+v: %v", encoded, self.ws, err)
					return
				} else if n != len(encoded) {
					self.server.Fatalf("Unable to send all of %s on %+v: only sent %v bytes", encoded, self.ws, n)
					return
				}
			} else {
				self.server.Fatalf("Unable to JSON marshal %+v: %v", message, err)
				return
			}
			self.server.Debugf("%v\t%v\t%v\t%v\t[sent to socket]", time.Now(), message.URI, self.RemoteAddr, self.id)
		case <-self.closing:
			return
		}
	}
}

func (self *Session) heartbeatLoop() {
	for {
		select {
		case <-self.closing:
			return
		case <-time.After(self.server.heartbeat / 2):
			self.send(Message{Type: TypeHeartbeat})
		}
	}
}

func (self *Session) authorized(uri string, wantWrite bool) bool {
	if self.server.authorizer == nil {
		return true
	}
	self.lock.RLock()
	defer self.lock.RUnlock()
	hasWrite, found := self.authorizations[uri]
	if !found {
		return false
	}
	return !wantWrite || hasWrite
}

func (self *Session) send(message Message) {
	select {
	case self.output <- message:
	default:
		self.server.Errorf("Unable to send %+v to %+v, output buffer full", message, self)
	}
}

func (self *Session) handleMessage(message Message) {
	switch message.Type {
	case TypeHeartbeat:
	case TypeMessage:
		if !self.authorized(message.URI, true) {
			self.send(Message{
				Type: TypeError,
				Id:   message.Id,
				Error: &Error{
					Message: fmt.Sprintf("%v not authorized for writing to %v", self.id, message.URI),
					Type:    TypeAuthorizationError,
				},
				Data: message,
			})
			return
		}
		self.server.Emit(message)
	case TypeUnsubscribe:
		self.server.removeSubscription(self.id, message.URI, true)
	case TypeSubscribe:
		if !self.authorized(message.URI, false) {
			self.send(Message{
				Type: TypeError,
				Id:   message.Id,
				Error: &Error{
					Message: fmt.Sprintf("%v not authorized for subscribing to %v", self.id, message.URI),
					Type:    TypeAuthorizationError,
				},
				Data: message,
			})
			return
		}
		self.server.addSubscription(self, message.URI)
	case TypeAuthorize:
		if self.server.authorizer != nil {
			ok, err := self.server.authorizer(message.URI, message.Token, message.Write)
			if err != nil {
				self.send(Message{
					Type: TypeError,
					Id:   message.Id,
					Error: &Error{
						Message: err.Error(),
						Type:    TypeAuthorizationError,
					}, Data: message,
				})
				return
			}
			if !ok {
				self.send(Message{
					Type: TypeError,
					Id:   message.Id,
					Error: &Error{
						Message: fmt.Sprintf("%v does not provide authorization for %v", message.Token, message.URI),
						Type:    TypeAuthorizationError,
					}, Data: message,
				})
				return
			}
		}
		self.lock.Lock()
		defer self.lock.Unlock()
		self.authorizations[message.URI] = (self.authorizations[message.URI] || message.Write)
		self.server.Infof("%v\t-\t[authorize]\t%v\t%v\t%v", time.Now(), self.RemoteAddr, self.id, message.URI)
	default:
		self.send(Message{
			Type: TypeError,
			Id:   message.Id,
			Error: &Error{
				Message: fmt.Sprintf("Unknown message type %#v", message.Type),
				Type:    TypeSyntaxError,
			},
			Data: message,
		})
		return
	}
	if message.Id != "" {
		self.send(Message{
			Type: TypeAck,
			Id:   message.Id,
		})
	}

}

func (self *Session) remove() {
	self.server.removeSession(self.id)
}

func (self *Session) terminate() {
	self.ws.Close()
	select {
	case _ = <-self.closing:
	default:
		close(self.closing)
	}
	self.server.Infof("%v\t-\t[disconnect]\t%v\t%v", time.Now(), self.RemoteAddr, self.id)
	self.cleanupTimer.Stop()
	self.cleanupTimer = time.AfterFunc(self.server.sessionTimeout, self.remove)
}

func (self *Session) handleWS(ws *websocket.Conn) {
	self.RemoteAddr = ws.Request().RemoteAddr
	self.Handle(ws)
}

/*
 * This function will work as the session main-loop listening on events
 * on a websocket or other source that interfaces io.ReadWriteCloser.
 */
func (self *Session) Handle(ws io.ReadWriteCloser) {
	self.server.Infof("%v\t-\t[connect]\t%v\t%v", time.Now(), self.RemoteAddr, self.id)

	defer self.terminate()

	self.ws = ws
	self.cleanupTimer.Stop()
	self.input = make(chan Message)
	self.closing = make(chan struct{})

	go self.readLoop()
	go self.writeLoop()
	go self.heartbeatLoop()

	self.send(Message{
		Type: TypeWelcome,
		Welcome: &Welcome{
			Heartbeat:      self.server.heartbeat / time.Millisecond,
			SessionTimeout: self.server.sessionTimeout / time.Millisecond,
			Id:             self.id,
		},
	})

	var message Message
	for {
		select {
		case _ = <-self.closing:
			return
		case message = <-self.input:
			self.handleMessage(message)
		case <-time.After(self.server.heartbeat):
			return
		}
	}
}
<|endoftext|>


<commit_before>package tcp

import (
	"fmt"
	"io"

	"v2ray.com/core/common/net"
)

type Server struct {
	Port         net.Port
	MsgProcessor func(msg []byte) []byte
	SendFirst    []byte
	Listen       net.Address
	accepting    bool
	listener     *net.TCPListener
}

func (server *Server) Start() (net.Destination, error) {
	listenerAddr := server.Listen
	if listenerAddr == nil {
		listenerAddr = net.LocalHostIP
	}
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{
		IP:   listenerAddr.IP(),
		Port: int(server.Port),
		Zone: "",
	})
	if err != nil {
		return net.Destination{}, err
	}
	server.Port = net.Port(listener.Addr().(*net.TCPAddr).Port)
	server.listener = listener
	go server.acceptConnections(listener)
	localAddr := listener.Addr().(*net.TCPAddr)
	return net.TCPDestination(net.IPAddress(localAddr.IP), net.Port(localAddr.Port)), nil
}

func (server *Server) acceptConnections(listener *net.TCPListener) {
	server.accepting = true
	for server.accepting {
		conn, err := listener.Accept()
		if err != nil {
			fmt.Printf("Failed accept TCP connection: %v\n", err)
			continue
		}

		go server.handleConnection(conn)
	}
}

func (server *Server) handleConnection(conn net.Conn) {
	if len(server.SendFirst) > 0 {
		conn.Write(server.SendFirst)
	}
	request := make([]byte, 4096)
	for {
		nBytes, err := conn.Read(request)
		if err != nil {
			if err != io.EOF {
				fmt.Println("Failed to read request:", err)
			}
			break
		}
		response := server.MsgProcessor(request[:nBytes])
		if _, err := conn.Write(response); err != nil {
			fmt.Println("Failed to write response:", err)
			break
		}
	}
	conn.Close()
}

func (v *Server) Close() {
	v.accepting = false
	v.listener.Close()
}
<commit_msg>fix data race in tcp server<commit_after>package tcp

import (
	"fmt"
	"io"

	"v2ray.com/core/common/net"
)

type Server struct {
	Port         net.Port
	MsgProcessor func(msg []byte) []byte
	SendFirst    []byte
	Listen       net.Address
	listener     *net.TCPListener
}

func (server *Server) Start() (net.Destination, error) {
	listenerAddr := server.Listen
	if listenerAddr == nil {
		listenerAddr = net.LocalHostIP
	}
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{
		IP:   listenerAddr.IP(),
		Port: int(server.Port),
		Zone: "",
	})
	if err != nil {
		return net.Destination{}, err
	}
	server.Port = net.Port(listener.Addr().(*net.TCPAddr).Port)
	server.listener = listener
	go server.acceptConnections(listener)
	localAddr := listener.Addr().(*net.TCPAddr)
	return net.TCPDestination(net.IPAddress(localAddr.IP), net.Port(localAddr.Port)), nil
}

func (server *Server) acceptConnections(listener *net.TCPListener) {
	for {
		conn, err := listener.Accept()
		if err != nil {
			fmt.Printf("Failed accept TCP connection: %v\n", err)
			return
		}

		go server.handleConnection(conn)
	}
}

func (server *Server) handleConnection(conn net.Conn) {
	if len(server.SendFirst) > 0 {
		conn.Write(server.SendFirst)
	}
	request := make([]byte, 4096)
	for {
		nBytes, err := conn.Read(request)
		if err != nil {
			if err != io.EOF {
				fmt.Println("Failed to read request:", err)
			}
			break
		}
		response := server.MsgProcessor(request[:nBytes])
		if _, err := conn.Write(response); err != nil {
			fmt.Println("Failed to write response:", err)
			break
		}
	}
	conn.Close()
}

func (server *Server) Close() {
	server.listener.Close()
}
<|endoftext|>


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

package common

import (
	"os"
	"strings"
	"testing"
	"time"

	clientv3 "go.etcd.io/etcd/client/v3"
	"go.etcd.io/etcd/tests/v3/framework/config"
	"go.etcd.io/etcd/tests/v3/framework/testutils"
)

func TestAlarm(t *testing.T) {
	testRunner.BeforeTest(t)
	clus := testRunner.NewCluster(t, config.ClusterConfig{ClusterSize: 3, QuotaBackendBytes: int64(13 * os.Getpagesize())})
	defer clus.Close()
	testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
		// test small put still works
		smallbuf := strings.Repeat("a", 64)
		if err := clus.Client().Put("1st_test", smallbuf, config.PutOptions{}); err != nil {
			t.Fatalf("alarmTest: put kv error (%v)", err)
		}

		// write some chunks to fill up the database
		buf := strings.Repeat("b", os.Getpagesize())
		for {
			if err := clus.Client().Put("2nd_test", buf, config.PutOptions{}); err != nil {
				if !strings.Contains(err.Error(), "etcdserver: mvcc: database space exceeded") {
					t.Fatal(err)
				}
				break
			}
		}

		// quota alarm should now be on
		alarmResp, err := clus.Client().Alarm("list", nil)
		if err != nil {
			t.Fatalf("alarmTest: Alarm error (%v)", err)
		}

		// check that Put is rejected when alarm is on
		if err := clus.Client().Put("3rd_test", smallbuf, config.PutOptions{}); err != nil {
			if !strings.Contains(err.Error(), "etcdserver: mvcc: database space exceeded") {
				t.Fatal(err)
			}
		}

		// get latest revision to compact
		sresp, err := clus.Client().Status()
		if err != nil {
			t.Fatalf("get endpoint status error: %v", err)
		}
		var rvs int64
		for _, resp := range sresp {
			if resp != nil && resp.Header != nil {
				rvs = resp.Header.Revision
				break
			}
		}

		// make some space
		_, err = clus.Client().Compact(rvs, config.CompactOption{Physical: true, Timeout: 10 * time.Second})
		if err != nil {
			t.Fatalf("alarmTest: Compact error (%v)", err)
		}

		if err = clus.Client().Defragment(config.DefragOption{Timeout: 10 * time.Second}); err != nil {
			t.Fatalf("alarmTest: defrag error (%v)", err)
		}

		// turn off alarm
		for _, alarm := range alarmResp.Alarms {
			m := &clientv3.AlarmMember{
				MemberID: alarm.MemberID,
				Alarm:    alarm.Alarm,
			}
			_, err = clus.Client().Alarm("disarm", m)
			if err != nil {
				t.Fatalf("alarmTest: Alarm error (%v)", err)
			}
		}

		// put one more key below quota
		if err := clus.Client().Put("4th_test", smallbuf, config.PutOptions{}); err != nil {
			t.Fatal(err)
		}
	})
}
<commit_msg>fix: set cluster size to 1 to avoid unstable<commit_after>// Copyright 2022 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package common

import (
	"os"
	"strings"
	"testing"
	"time"

	clientv3 "go.etcd.io/etcd/client/v3"
	"go.etcd.io/etcd/tests/v3/framework/config"
	"go.etcd.io/etcd/tests/v3/framework/testutils"
)

func TestAlarm(t *testing.T) {
	testRunner.BeforeTest(t)
	clus := testRunner.NewCluster(t, config.ClusterConfig{ClusterSize: 1, QuotaBackendBytes: int64(13 * os.Getpagesize())})
	defer clus.Close()
	testutils.ExecuteWithTimeout(t, 10*time.Second, func() {
		// test small put still works
		smallbuf := strings.Repeat("a", 64)
		if err := clus.Client().Put("1st_test", smallbuf, config.PutOptions{}); err != nil {
			t.Fatalf("alarmTest: put kv error (%v)", err)
		}

		// write some chunks to fill up the database
		buf := strings.Repeat("b", os.Getpagesize())
		for {
			if err := clus.Client().Put("2nd_test", buf, config.PutOptions{}); err != nil {
				if !strings.Contains(err.Error(), "etcdserver: mvcc: database space exceeded") {
					t.Fatal(err)
				}
				break
			}
		}

		// quota alarm should now be on
		alarmResp, err := clus.Client().Alarm("list", nil)
		if err != nil {
			t.Fatalf("alarmTest: Alarm error (%v)", err)
		}

		// check that Put is rejected when alarm is on
		if err := clus.Client().Put("3rd_test", smallbuf, config.PutOptions{}); err != nil {
			if !strings.Contains(err.Error(), "etcdserver: mvcc: database space exceeded") {
				t.Fatal(err)
			}
		}

		// get latest revision to compact
		sresp, err := clus.Client().Status()
		if err != nil {
			t.Fatalf("get endpoint status error: %v", err)
		}
		var rvs int64
		for _, resp := range sresp {
			if resp != nil && resp.Header != nil {
				rvs = resp.Header.Revision
				break
			}
		}

		// make some space
		_, err = clus.Client().Compact(rvs, config.CompactOption{Physical: true, Timeout: 10 * time.Second})
		if err != nil {
			t.Fatalf("alarmTest: Compact error (%v)", err)
		}

		if err = clus.Client().Defragment(config.DefragOption{Timeout: 10 * time.Second}); err != nil {
			t.Fatalf("alarmTest: defrag error (%v)", err)
		}

		// turn off alarm
		for _, alarm := range alarmResp.Alarms {
			m := &clientv3.AlarmMember{
				MemberID: alarm.MemberID,
				Alarm:    alarm.Alarm,
			}
			_, err = clus.Client().Alarm("disarm", m)
			if err != nil {
				t.Fatalf("alarmTest: Alarm error (%v)", err)
			}
		}

		// put one more key below quota
		if err := clus.Client().Put("4th_test", smallbuf, config.PutOptions{}); err != nil {
			t.Fatal(err)
		}
	})
}
<|endoftext|>