<commit_before>package gosecco

import (
	"os"
	"path"
	"strings"
	"testing"

	"github.com/twtiger/gosecco/asm"
	"golang.org/x/sys/unix"

	. "gopkg.in/check.v1"
)

func Test(t *testing.T) { TestingT(t) }

type SeccompSuite struct{}

var _ = Suite(&SeccompSuite{})

func (s *SeccompSuite) Test_loadingTooBigBpf(c *C) {
	inp := make([]unix.SockFilter, 0xFFFF+1)
	res := Load(inp)
	c.Assert(res, ErrorMatches, "filter program too big: 65536 bpf instructions \\(limit = 65535\\)")
}

func getActualTestFolder() string {
	wd, _ := os.Getwd()
	if strings.HasSuffix(wd, "/parser/test_policies/") {
		return wd
	}
	return path.Join(wd, "parser/test_policies/")
}

func (s *SeccompSuite) Test_parseInvalidFileReturnsErrors(c *C) {
	set := SeccompSettings{}
	f := getActualTestFolder() + "/failing_test_policy"
	_, ee := Prepare(f, set)
	c.Assert(ee, ErrorMatches, ".*parser/test_policies/failing_test_policy:1: unexpected end of line")
}

func (s *SeccompSuite) Test_parseUnificationErrorReturnsError(c *C) {
	set := SeccompSettings{}
	f := getActualTestFolder() + "/missing_variable_policy"
	_, ee := Prepare(f, set)
	c.Assert(ee, ErrorMatches, "Variable not defined")
}

func (s *SeccompSuite) Test_parseValidPolicyFile(c *C) {
	set := SeccompSettings{DefaultPositiveAction: "allow", DefaultNegativeAction: "kill", DefaultPolicyAction: "kill"}
	f := getActualTestFolder() + "/valid_test_policy"
	res, ee := Prepare(f, set)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t04\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t01\t00\t1\n"+
		"jmp\t1\n"+
		"ret_k\t7FFF0000\n"+
		"ret_k\t0\n")
}

func (s *SeccompSuite) Test_parseInvalidTypeReturnsError(c *C) {
	set := SeccompSettings{}
	f := getActualTestFolder() + "/type_checker_error_policy"
	_, ee := Prepare(f, set)
	c.Assert(ee, ErrorMatches, ".*expected boolean expression but found: 42")
}

func (s *SeccompSuite) Test_parseSimplifiesValidExpression(c *C) {
	set := SeccompSettings{DefaultPositiveAction: "allow", DefaultNegativeAction: "kill", DefaultPolicyAction: "kill"}
	f := getActualTestFolder() + "/valid_unsimplified_policy"
	res, ee := Prepare(f, set)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t0E\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t00\t0A\t1\n"+
		"ld_imm\t3\n"+
		"st\t0\n"+
		"ld_abs\t14\n"+
		"ldx_mem\t0\n"+
		"jeq_x\t00\t07\n"+
		"ld_imm\t0\n"+
		"st\t0\n"+
		"ld_abs\t10\n"+
		"ldx_mem\t0\n"+
		"jeq_x\t01\t02\n"+
		"jmp\t1\n"+
		"ret_k\t7FFF0000\n"+
		"ret_k\t0\n")
}

func (s *SeccompSuite) Test_parseSetsDefaultActions(c *C) {
	set := SeccompSettings{DefaultPositiveAction: "kill", DefaultNegativeAction: "trace", DefaultPolicyAction: "trace"}
	f := getActualTestFolder() + "/valid_test_policy"
	res, ee := Prepare(f, set)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t04\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t00\t01\t1\n"+
		"jmp\t1\n"+
		"jmp\t1\n"+
		"ret_k\t0\n"+
		"ret_k\t7FF00000\n")
}

func (s *SeccompSuite) Test_compileWithEnforce(c *C) {
	f := getActualTestFolder() + "/valid_test_policy"
	res, ee := Compile(f, true)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t05\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t00\t01\t1\n"+
		"jmp\t1\n"+
		"jmp\t1\n"+
		"ret_k\t7FFF0000\n"+
		"ret_k\t0\n")
}

func (s *SeccompSuite) Test_compileWithoutEnforce(c *C) {
	f := getActualTestFolder() + "/valid_test_policy"
	res, ee := Compile(f, false)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t05\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t00\t01\t1\n"+
		"jmp\t1\n"+
		"jmp\t2\n"+
		"ret_k\t7FFF0000\n"+
		"ret_k\t0\n"+
		"ret_k\t7FF00000\n")
}

func (s *SeccompSuite) Test_compileBlacklistWithEnforce(c *C) {
	f := getActualTestFolder() + "/valid_test_policy"
	res, ee := CompileBlacklist(f, true)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t05\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t00\t01\t1\n"+
		"jmp\t2\n"+
		"jmp\t0\n"+
		"ret_k\t7FFF0000\n"+
		"ret_k\t0\n")
}

func (s *SeccompSuite) Test_compileBlacklistWithoutEnforce(c *C) {
	f := getActualTestFolder() + "/valid_test_policy"
	res, ee := CompileBlacklist(f, false)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t05\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t00\t01\t1\n"+
		"jmp\t3\n"+
		"jmp\t0\n"+
		"ret_k\t7FFF0000\n"+
		"ret_k\t0\n"+
		"ret_k\t7FF00000\n")
}
<commit_msg>Fix failing tests with new optimizer<commit_after>package gosecco

import (
	"os"
	"path"
	"strings"
	"testing"

	"github.com/twtiger/gosecco/asm"
	"golang.org/x/sys/unix"

	. "gopkg.in/check.v1"
)

func Test(t *testing.T) { TestingT(t) }

type SeccompSuite struct{}

var _ = Suite(&SeccompSuite{})

func (s *SeccompSuite) Test_loadingTooBigBpf(c *C) {
	inp := make([]unix.SockFilter, 0xFFFF+1)
	res := Load(inp)
	c.Assert(res, ErrorMatches, "filter program too big: 65536 bpf instructions \\(limit = 65535\\)")
}

func getActualTestFolder() string {
	wd, _ := os.Getwd()
	if strings.HasSuffix(wd, "/parser/test_policies/") {
		return wd
	}
	return path.Join(wd, "parser/test_policies/")
}

func (s *SeccompSuite) Test_parseInvalidFileReturnsErrors(c *C) {
	set := SeccompSettings{}
	f := getActualTestFolder() + "/failing_test_policy"
	_, ee := Prepare(f, set)
	c.Assert(ee, ErrorMatches, ".*parser/test_policies/failing_test_policy:1: unexpected end of line")
}

func (s *SeccompSuite) Test_parseUnificationErrorReturnsError(c *C) {
	set := SeccompSettings{}
	f := getActualTestFolder() + "/missing_variable_policy"
	_, ee := Prepare(f, set)
	c.Assert(ee, ErrorMatches, "Variable not defined")
}

func (s *SeccompSuite) Test_parseValidPolicyFile(c *C) {
	set := SeccompSettings{DefaultPositiveAction: "allow", DefaultNegativeAction: "kill", DefaultPolicyAction: "kill"}
	f := getActualTestFolder() + "/valid_test_policy"
	res, ee := Prepare(f, set)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t04\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t01\t00\t1\n"+
		"jmp\t1\n"+
		"ret_k\t7FFF0000\n"+
		"ret_k\t0\n")
}

func (s *SeccompSuite) Test_parseInvalidTypeReturnsError(c *C) {
	set := SeccompSettings{}
	f := getActualTestFolder() + "/type_checker_error_policy"
	_, ee := Prepare(f, set)
	c.Assert(ee, ErrorMatches, ".*expected boolean expression but found: 42")
}

func (s *SeccompSuite) Test_parseSimplifiesValidExpression(c *C) {
	set := SeccompSettings{DefaultPositiveAction: "allow", DefaultNegativeAction: "kill", DefaultPolicyAction: "kill"}
	f := getActualTestFolder() + "/valid_unsimplified_policy"
	res, ee := Prepare(f, set)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t0E\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t00\t0A\t1\n"+
		"ld_imm\t3\n"+
		"st\t0\n"+
		"ld_abs\t14\n"+
		"ldx_mem\t0\n"+
		"jeq_x\t00\t07\n"+
		"ld_imm\t0\n"+
		"st\t0\n"+
		"ld_abs\t10\n"+
		"ldx_mem\t0\n"+
		"jeq_x\t01\t02\n"+
		"jmp\t1\n"+
		"ret_k\t7FFF0000\n"+
		"ret_k\t0\n")
}

func (s *SeccompSuite) Test_parseSetsDefaultActions(c *C) {
	set := SeccompSettings{DefaultPositiveAction: "kill", DefaultNegativeAction: "trace", DefaultPolicyAction: "trace"}
	f := getActualTestFolder() + "/valid_test_policy"
	res, ee := Prepare(f, set)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t03\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t01\t00\t1\n"+
		"jmp\t1\n"+
		"ret_k\t0\n"+
		"ret_k\t7FF00000\n")
}

func (s *SeccompSuite) Test_compileWithEnforce(c *C) {
	f := getActualTestFolder() + "/valid_test_policy"
	res, ee := Compile(f, true)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t04\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t01\t00\t1\n"+
		"jmp\t1\n"+
		"ret_k\t7FFF0000\n"+
		"ret_k\t0\n")
}

func (s *SeccompSuite) Test_compileWithoutEnforce(c *C) {
	f := getActualTestFolder() + "/valid_test_policy"
	res, ee := Compile(f, false)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t04\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t01\t00\t1\n"+
		"jmp\t2\n"+
		"ret_k\t7FFF0000\n"+
		"ret_k\t0\n"+
		"ret_k\t7FF00000\n")
}

func (s *SeccompSuite) Test_compileBlacklistWithEnforce(c *C) {
	f := getActualTestFolder() + "/valid_test_policy"
	res, ee := CompileBlacklist(f, true)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t04\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t02\t00\t1\n"+
		"jmp\t0\n"+
		"ret_k\t7FFF0000\n"+
		"ret_k\t0\n")
}

func (s *SeccompSuite) Test_compileBlacklistWithoutEnforce(c *C) {
	f := getActualTestFolder() + "/valid_test_policy"
	res, ee := CompileBlacklist(f, false)

	c.Assert(ee, Equals, nil)

	c.Assert(asm.Dump(res), Equals, ""+
		"ld_abs\t4\n"+
		"jeq_k\t00\t04\tC000003E\n"+
		"ld_abs\t0\n"+
		"jeq_k\t03\t00\t1\n"+
		"jmp\t0\n"+
		"ret_k\t7FFF0000\n"+
		"ret_k\t0\n"+
		"ret_k\t7FF00000\n")
}
<|endoftext|>


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

package generate_test

import (
	. "github.com/jacobsa/ogletest"
	"bytes"
	"flag"
	"github.com/jacobsa/oglemock/generate"
	"image"
	"io"
	"io/ioutil"
	"path"
	"reflect"
	"testing"
)

var dumpNew = flag.Bool("dump_new", false, "Dump new golden files.")

////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////

type GenerateTest struct {
}

func TestOgletest(t *testing.T) { RunTests(t) }
func init() { RegisterTestSuite(&GenerateTest{}) }

func (t *GenerateTest) runGoldenTest(
	caseName string,
	nilPtrs ...interface{}) {
  // Make a slice of interface types to give to GenerateMockSource.
	interfaces := make([]reflect.Type, len(nilPtrs))
	for i, ptr := range nilPtrs {
		interfaces[i] = reflect.TypeOf(ptr).Elem()
	}

	// Create the mock source.
	buf := new(bytes.Buffer)
	err := generate.GenerateMockSource(buf, "some_pkg", interfaces)
	AssertEq(nil, err, "Error from GenerateMockSource: %v", err)

	// Read the golden file.
	goldenPath := path.Join("test_cases", "golden." + caseName + ".go")
	goldenData := readFileOrDie(goldenPath)

	// Compare the two.
	identical := (buf.String() == string(goldenData))
	ExpectTrue(identical, "Output doesn't match for case '%s'.", caseName)

	// Write out a new golden file if requested.
	if !identical && *dumpNew {
		writeContentsToFileOrDie(buf.Bytes(), goldenPath)
	}
}

func writeContentsToFileOrDie(contents []byte, path string) {
	if err := ioutil.WriteFile(path, contents, 0600); err != nil {
		panic("ioutil.WriteFile: " + err.Error())
	}
}

func readFileOrDie(path string) []byte {
	contents, err := ioutil.ReadFile(path)
	if err != nil {
		panic("ioutil.ReadFile: " + err.Error())
	}

	return contents
}

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

func (t *GenerateTest) EmptySet() {
	ExpectEq("TODO", 17)
}

func (t *GenerateTest) InvalidType() {
	ExpectEq("TODO", 17)
}

func (t *GenerateTest) NonInterfaceType() {
	ExpectEq("TODO", 17)
}

func (t *GenerateTest) SomeOfPkgIo() {
	// Mock io.Reader and io.Writer.
	t.runGoldenTest(
		"io_partial",
		(*io.Reader)(nil),
		(*io.Writer)(nil))
}

func (t *GenerateTest) Image() {
	t.runGoldenTest(
		"image",
		(*image.Image)(nil),
		(*image.PalettedImage)(nil))
}
<commit_msg>Implemented several test cases.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package generate_test

import (
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"bytes"
	"flag"
	"github.com/jacobsa/oglemock/generate"
	"image"
	"io"
	"io/ioutil"
	"path"
	"reflect"
	"testing"
)

var dumpNew = flag.Bool("dump_new", false, "Dump new golden files.")

////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////

type GenerateTest struct {
}

func TestOgletest(t *testing.T) { RunTests(t) }
func init() { RegisterTestSuite(&GenerateTest{}) }

func (t *GenerateTest) runGoldenTest(
	caseName string,
	nilPtrs ...interface{}) {
  // Make a slice of interface types to give to GenerateMockSource.
	interfaces := make([]reflect.Type, len(nilPtrs))
	for i, ptr := range nilPtrs {
		interfaces[i] = reflect.TypeOf(ptr).Elem()
	}

	// Create the mock source.
	buf := new(bytes.Buffer)
	err := generate.GenerateMockSource(buf, "some_pkg", interfaces)
	AssertEq(nil, err, "Error from GenerateMockSource: %v", err)

	// Read the golden file.
	goldenPath := path.Join("test_cases", "golden." + caseName + ".go")
	goldenData := readFileOrDie(goldenPath)

	// Compare the two.
	identical := (buf.String() == string(goldenData))
	ExpectTrue(identical, "Output doesn't match for case '%s'.", caseName)

	// Write out a new golden file if requested.
	if !identical && *dumpNew {
		writeContentsToFileOrDie(buf.Bytes(), goldenPath)
	}
}

func writeContentsToFileOrDie(contents []byte, path string) {
	if err := ioutil.WriteFile(path, contents, 0600); err != nil {
		panic("ioutil.WriteFile: " + err.Error())
	}
}

func readFileOrDie(path string) []byte {
	contents, err := ioutil.ReadFile(path)
	if err != nil {
		panic("ioutil.ReadFile: " + err.Error())
	}

	return contents
}

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

func (t *GenerateTest) EmptyPackageName() {
	err := generate.GenerateMockSource(
		new(bytes.Buffer),
		"",
		[]reflect.Type{
			reflect.TypeOf((*io.Reader)(nil)).Elem(),
		})

  ExpectThat(err, Error(HasSubstr("Package name")))
  ExpectThat(err, Error(HasSubstr("non-empty")))
}

func (t *GenerateTest) EmptySetOfInterfaces() {
	err := generate.GenerateMockSource(
		new(bytes.Buffer),
		"foo",
		[]reflect.Type{})

  ExpectThat(err, Error(HasSubstr("at least one")))
  ExpectThat(err, Error(HasSubstr("interface")))
}

func (t *GenerateTest) InvalidType() {
	err := generate.GenerateMockSource(
		new(bytes.Buffer),
		"foo",
		[]reflect.Type{
			reflect.TypeOf((*io.Reader)(nil)).Elem(),
			reflect.TypeOf(nil),
			reflect.TypeOf((*io.Writer)(nil)).Elem(),
		})

  ExpectThat(err, Error(HasSubstr("Invalid type")))
}

func (t *GenerateTest) NonInterfaceType() {
	err := generate.GenerateMockSource(
		new(bytes.Buffer),
		"foo",
		[]reflect.Type{
			reflect.TypeOf((*io.Reader)(nil)).Elem(),
			reflect.TypeOf(17),
			reflect.TypeOf((*io.Writer)(nil)).Elem(),
		})

  ExpectThat(err, Error(HasSubstr("Invalid type")))
}

func (t *GenerateTest) SomeOfPkgIo() {
	// Mock io.Reader and io.Writer.
	t.runGoldenTest(
		"io_partial",
		(*io.Reader)(nil),
		(*io.Writer)(nil))
}

func (t *GenerateTest) Image() {
	t.runGoldenTest(
		"image",
		(*image.Image)(nil),
		(*image.PalettedImage)(nil))
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"strconv"

	"github.com/gorilla/mux"
	"github.com/nytlabs/st-core/core"
)

type Position struct {
	X float64 `json:"x"`
	Y float64 `json:"y"`
}

type ProtoBlock struct {
	Label    string   `json:"label"`
	Parent   int      `json:"parent"`
	Type     string   `json:"type"`
	Position Position `json:"position"`
}

type BlockLedger struct {
	Label       string        `json:"label"`
	Type        string        `json:"type"`
	Id          int           `json:"id"`
	Block       *core.Block   `json:"-"`
	Parent      *Group        `json:"-"`
	Composition int           `json:"composition,omitempty"`
	Inputs      []core.Input  `json:"inputs"`
	Outputs     []core.Output `json:"outputs"`
	Position    Position      `json:"position"`
}

func (bl *BlockLedger) GetID() int {
	return bl.Id
}

func (bl *BlockLedger) GetParent() *Group {
	return bl.Parent
}

func (bl *BlockLedger) SetParent(group *Group) {
	bl.Parent = group
}

func (s *Server) ListBlocks() []BlockLedger {
	blocks := []BlockLedger{}
	for _, b := range s.blocks {
		blocks = append(blocks, *b)
	}
	return blocks
}

func (s *Server) BlockIndexHandler(w http.ResponseWriter, r *http.Request) {
	s.Lock()
	defer s.Unlock()

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(s.ListBlocks()); err != nil {
		panic(err)
	}
}

func (s *Server) BlockHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	b, ok := s.blocks[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find block"})
		return
	}

	w.WriteHeader(http.StatusOK)
	writeJSON(w, b)
	return
}

func (s *Server) BlockModifyPositionHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var p Position
	err = json.Unmarshal(body, &p)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read JSON"})
		return
	}

	s.Lock()
	defer s.Unlock()

	b, ok := s.blocks[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find block"})
		return
	}

	b.Position = p

	s.websocketBroadcast(Update{Action: UPDATE, Type: BLOCK, Data: wsBlock{wsPosition{wsId{id}, p}}})
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) CreateBlock(p ProtoBlock) (*BlockLedger, error) {
	blockSpec, ok := s.library[p.Type]
	if !ok {
		return nil, errors.New("spec not found")
	}

	block := core.NewBlock(blockSpec)

	m := &BlockLedger{
		Label:    p.Label,
		Position: p.Position,
		Type:     p.Type,
		Block:    block,
		Id:       s.GetNextID(),
	}

	if _, ok := s.groups[p.Parent]; !ok {
		return nil, errors.New("invalid group, could not create block")
	}

	go block.Serve()
	m.Inputs = block.GetInputs()
	m.Outputs = block.GetOutputs()
	s.blocks[m.Id] = m

	s.websocketBroadcast(Update{Action: CREATE, Type: BLOCK, Data: wsBlock{*m}})

	err := s.AddChildToGroup(p.Parent, m)
	if err != nil {
		return nil, err

	}

	return m, nil
}

// CreateBlockHandler responds to a POST request to instantiate a new block and add it to the Server.
func (s *Server) BlockCreateHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var m ProtoBlock
	err = json.Unmarshal(body, &m)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	s.Lock()
	defer s.Unlock()

	b, err := s.CreateBlock(m)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusOK)
	writeJSON(w, b)
}

func (s *Server) BlockModifyNameHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	s.Lock()
	defer s.Unlock()

	_, ok := s.blocks[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"block not found"})
		return
	}

	var label string
	err = json.Unmarshal(body, &label)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not unmarshal value"})
		return
	}

	s.blocks[id].Label = label

	s.websocketBroadcast(Update{Action: UPDATE, Type: BLOCK, Data: wsBlock{wsLabel{wsId{id}, label}}})
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) DeleteBlock(id int) error {
	b, ok := s.blocks[id]
	if !ok {
		return errors.New("block not found")
	}

	deleteSet := make(map[int]struct{})

	// build a set of connections that we may need to delete
	// we need to panic here because if any error is thrown we are in huge trouble
	// any panic indicates that our server connection ledger is no longer true
	for _, c := range s.connections {
		if c.Target.Id == id || c.Source.Id == id {
			deleteSet[c.Id] = struct{}{}
		}
	}

	// delete the connections that involve this block
	for k, _ := range deleteSet {
		s.DeleteConnection(k)
	}

	// remove from group
	s.DetachChild(b)

	// stop and delete the block
	b.Block.Stop()
	s.websocketBroadcast(Update{Action: DELETE, Type: BLOCK, Data: wsBlock{wsId{id}}})
	delete(s.blocks, id)
	return nil
}

func (s *Server) BlockDeleteHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.DeleteBlock(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) BlockModifyRouteHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id, err := getIDFromMux(vars)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	routes, ok := vars["index"]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no route index supplied"})
		return
	}

	route, err := strconv.Atoi(routes)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var v *core.InputValue
	err = json.Unmarshal(body, &v)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not unmarshal value"})
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.ModifyBlockRoute(id, route, v)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) ModifyBlockRoute(id int, route int, v *core.InputValue) error {
	b, ok := s.blocks[id]
	if !ok {
		return errors.New("could not find block")
	}

	var value *core.InputValue

	if v.Exists() {
		value = v
	}

	err := b.Block.SetInput(core.RouteIndex(route), value)
	if err != nil {
		return err
	}

	s.blocks[id].Inputs[route].Value = value

	s.websocketBroadcast(Update{Action: UPDATE, Type: ROUTE, Data: wsRouteModify{ConnectionNode{id, route}, value}})
	return nil
}
<commit_msg>blocks now state what kind of source they depend on<commit_after>package server

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"strconv"

	"github.com/gorilla/mux"
	"github.com/nytlabs/st-core/core"
)

type Position struct {
	X float64 `json:"x"`
	Y float64 `json:"y"`
}

type ProtoBlock struct {
	Label    string   `json:"label"`
	Parent   int      `json:"parent"`
	Type     string   `json:"type"`
	Position Position `json:"position"`
}

type BlockLedger struct {
	Label       string          `json:"label"`
	Type        string          `json:"type"`
	Id          int             `json:"id"`
	Block       *core.Block     `json:"-"`
	Parent      *Group          `json:"-"`
	Composition int             `json:"composition,omitempty"`
	Inputs      []core.Input    `json:"inputs"`
	Outputs     []core.Output   `json:"outputs"`
	Source      core.SourceType `json:"source"`
	Position    Position        `json:"position"`
}

func (bl *BlockLedger) GetID() int {
	return bl.Id
}

func (bl *BlockLedger) GetParent() *Group {
	return bl.Parent
}

func (bl *BlockLedger) SetParent(group *Group) {
	bl.Parent = group
}

func (s *Server) ListBlocks() []BlockLedger {
	blocks := []BlockLedger{}
	for _, b := range s.blocks {
		blocks = append(blocks, *b)
	}
	return blocks
}

func (s *Server) BlockIndexHandler(w http.ResponseWriter, r *http.Request) {
	s.Lock()
	defer s.Unlock()

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(s.ListBlocks()); err != nil {
		panic(err)
	}
}

func (s *Server) BlockHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	b, ok := s.blocks[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find block"})
		return
	}

	w.WriteHeader(http.StatusOK)
	writeJSON(w, b)
	return
}

func (s *Server) BlockModifyPositionHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var p Position
	err = json.Unmarshal(body, &p)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read JSON"})
		return
	}

	s.Lock()
	defer s.Unlock()

	b, ok := s.blocks[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find block"})
		return
	}

	b.Position = p

	s.websocketBroadcast(Update{Action: UPDATE, Type: BLOCK, Data: wsBlock{wsPosition{wsId{id}, p}}})
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) CreateBlock(p ProtoBlock) (*BlockLedger, error) {
	blockSpec, ok := s.library[p.Type]
	if !ok {
		return nil, errors.New("spec not found")
	}

	block := core.NewBlock(blockSpec)

	m := &BlockLedger{
		Label:    p.Label,
		Position: p.Position,
		Type:     p.Type,
		Block:    block,
		Source:   blockSpec.Source,
		Id:       s.GetNextID(),
	}

	if _, ok := s.groups[p.Parent]; !ok {
		return nil, errors.New("invalid group, could not create block")
	}

	go block.Serve()
	m.Inputs = block.GetInputs()
	m.Outputs = block.GetOutputs()
	s.blocks[m.Id] = m

	s.websocketBroadcast(Update{Action: CREATE, Type: BLOCK, Data: wsBlock{*m}})

	err := s.AddChildToGroup(p.Parent, m)
	if err != nil {
		return nil, err

	}

	return m, nil
}

// CreateBlockHandler responds to a POST request to instantiate a new block and add it to the Server.
func (s *Server) BlockCreateHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var m ProtoBlock
	err = json.Unmarshal(body, &m)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	s.Lock()
	defer s.Unlock()

	b, err := s.CreateBlock(m)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusOK)
	writeJSON(w, b)
}

func (s *Server) BlockModifyNameHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	s.Lock()
	defer s.Unlock()

	_, ok := s.blocks[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"block not found"})
		return
	}

	var label string
	err = json.Unmarshal(body, &label)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not unmarshal value"})
		return
	}

	s.blocks[id].Label = label

	s.websocketBroadcast(Update{Action: UPDATE, Type: BLOCK, Data: wsBlock{wsLabel{wsId{id}, label}}})
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) DeleteBlock(id int) error {
	b, ok := s.blocks[id]
	if !ok {
		return errors.New("block not found")
	}

	deleteSet := make(map[int]struct{})

	// build a set of connections that we may need to delete
	// we need to panic here because if any error is thrown we are in huge trouble
	// any panic indicates that our server connection ledger is no longer true
	for _, c := range s.connections {
		if c.Target.Id == id || c.Source.Id == id {
			deleteSet[c.Id] = struct{}{}
		}
	}

	// delete the connections that involve this block
	for k, _ := range deleteSet {
		s.DeleteConnection(k)
	}

	// remove from group
	s.DetachChild(b)

	// stop and delete the block
	b.Block.Stop()
	s.websocketBroadcast(Update{Action: DELETE, Type: BLOCK, Data: wsBlock{wsId{id}}})
	delete(s.blocks, id)
	return nil
}

func (s *Server) BlockDeleteHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.DeleteBlock(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) BlockModifyRouteHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id, err := getIDFromMux(vars)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	routes, ok := vars["index"]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no route index supplied"})
		return
	}

	route, err := strconv.Atoi(routes)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var v *core.InputValue
	err = json.Unmarshal(body, &v)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not unmarshal value"})
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.ModifyBlockRoute(id, route, v)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) ModifyBlockRoute(id int, route int, v *core.InputValue) error {
	b, ok := s.blocks[id]
	if !ok {
		return errors.New("could not find block")
	}

	var value *core.InputValue

	if v.Exists() {
		value = v
	}

	err := b.Block.SetInput(core.RouteIndex(route), value)
	if err != nil {
		return err
	}

	s.blocks[id].Inputs[route].Value = value

	s.websocketBroadcast(Update{Action: UPDATE, Type: ROUTE, Data: wsRouteModify{ConnectionNode{id, route}, value}})
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2012-2014 Apcera Inc. All rights reserved.

package server

import (
	"time"
)

const (
	// VERSION is the current version for the server.
	VERSION = "0.5.7"

	// DEFAULT_PORT is the deault port for client connections.
	DEFAULT_PORT = 4222

	// RANDOM_PORT is the value for port that, when supplied, will cause the
	// server to listen on a randomly-chosen available port. The resolved port
	// is available via the Addr() method.
	RANDOM_PORT = -1

	// DEFAULT_HOST defaults to all interfaces.
	DEFAULT_HOST = "0.0.0.0"

	// MAX_CONTROL_LINE_SIZE is the maximum allowed protocol control line size.
	// 1k should be plenty since payloads sans connect string are separate
	MAX_CONTROL_LINE_SIZE = 1024

	// MAX_PAYLOAD_SIZE is the maximum allowed payload size. Should be using
	// something different if > 1MB payloads are needed.
	MAX_PAYLOAD_SIZE = (1024 * 1024)

	// MAX_PENDING_SIZE is the maximum outbound size (in bytes) per client.
	MAX_PENDING_SIZE = (10 * 1024 * 1024)

	// DEFAULT_MAX_CONNECTIONS is the default maximum connections allowed.
	DEFAULT_MAX_CONNECTIONS = (64 * 1024)

	// SSL_TIMEOUT is the TLS/SSL wait time.
	SSL_TIMEOUT = 500 * time.Millisecond

	// AUTH_TIMEOUT is the authorization wait time.
	AUTH_TIMEOUT = 2 * SSL_TIMEOUT

	// DEFAULT_PING_INTERVAL is how often pings are sent to clients and routes.
	DEFAULT_PING_INTERVAL = 2 * time.Minute

	// DEFAULT_PING_MAX_OUT is maximum allowed pings outstanding before disconnect.
	DEFAULT_PING_MAX_OUT = 2

	// CRLF string
	CR_LF = "\r\n"

	// LEN_CR_LF hold onto the computed size.
	LEN_CR_LF = len(CR_LF)

	// DEFAULT_FLUSH_DEADLINE is the write/flush deadlines.
	DEFAULT_FLUSH_DEADLINE = 2 * time.Second

	// DEFAULT_HTTP_PORT is the default monitoring port.
	DEFAULT_HTTP_PORT = 8333

	// ACCEPT_MIN_SLEEP is the minimum acceptable sleep times on temporary errors.
	ACCEPT_MIN_SLEEP = 10 * time.Millisecond

	// ACCEPT_MAX_SLEEP is the maximum acceptable sleep times on temporary errors
	ACCEPT_MAX_SLEEP = 1 * time.Second

	// DEFAULT_ROUTE_CONNECT Route solicitation intervals.
	DEFAULT_ROUTE_CONNECT = 1 * time.Second

	// DEFAULT_ROUTE_RECONNECT Route reconnect intervals.
	DEFAULT_ROUTE_RECONNECT = 1 * time.Second

	// DEFAULT_ROUTE_DIAL Route dial timeout.
	DEFAULT_ROUTE_DIAL = 1 * time.Second

	// PROTO_SNIPPET_SIZE is the default size of proto to print on parse errors.
	PROTO_SNIPPET_SIZE = 32

	// MAX_MSG_ARGS Maximum possible number of arguments from MSG proto.
	MAX_MSG_ARGS = 4

	// MAX_PUB_ARGS Maximum possible number of arguments from PUB proto.
	MAX_PUB_ARGS = 3
)
<commit_msg>Bumped version<commit_after>// Copyright 2012-2014 Apcera Inc. All rights reserved.

package server

import (
	"time"
)

const (
	// VERSION is the current version for the server.
	VERSION = "0.5.8"

	// DEFAULT_PORT is the deault port for client connections.
	DEFAULT_PORT = 4222

	// RANDOM_PORT is the value for port that, when supplied, will cause the
	// server to listen on a randomly-chosen available port. The resolved port
	// is available via the Addr() method.
	RANDOM_PORT = -1

	// DEFAULT_HOST defaults to all interfaces.
	DEFAULT_HOST = "0.0.0.0"

	// MAX_CONTROL_LINE_SIZE is the maximum allowed protocol control line size.
	// 1k should be plenty since payloads sans connect string are separate
	MAX_CONTROL_LINE_SIZE = 1024

	// MAX_PAYLOAD_SIZE is the maximum allowed payload size. Should be using
	// something different if > 1MB payloads are needed.
	MAX_PAYLOAD_SIZE = (1024 * 1024)

	// MAX_PENDING_SIZE is the maximum outbound size (in bytes) per client.
	MAX_PENDING_SIZE = (10 * 1024 * 1024)

	// DEFAULT_MAX_CONNECTIONS is the default maximum connections allowed.
	DEFAULT_MAX_CONNECTIONS = (64 * 1024)

	// SSL_TIMEOUT is the TLS/SSL wait time.
	SSL_TIMEOUT = 500 * time.Millisecond

	// AUTH_TIMEOUT is the authorization wait time.
	AUTH_TIMEOUT = 2 * SSL_TIMEOUT

	// DEFAULT_PING_INTERVAL is how often pings are sent to clients and routes.
	DEFAULT_PING_INTERVAL = 2 * time.Minute

	// DEFAULT_PING_MAX_OUT is maximum allowed pings outstanding before disconnect.
	DEFAULT_PING_MAX_OUT = 2

	// CRLF string
	CR_LF = "\r\n"

	// LEN_CR_LF hold onto the computed size.
	LEN_CR_LF = len(CR_LF)

	// DEFAULT_FLUSH_DEADLINE is the write/flush deadlines.
	DEFAULT_FLUSH_DEADLINE = 2 * time.Second

	// DEFAULT_HTTP_PORT is the default monitoring port.
	DEFAULT_HTTP_PORT = 8333

	// ACCEPT_MIN_SLEEP is the minimum acceptable sleep times on temporary errors.
	ACCEPT_MIN_SLEEP = 10 * time.Millisecond

	// ACCEPT_MAX_SLEEP is the maximum acceptable sleep times on temporary errors
	ACCEPT_MAX_SLEEP = 1 * time.Second

	// DEFAULT_ROUTE_CONNECT Route solicitation intervals.
	DEFAULT_ROUTE_CONNECT = 1 * time.Second

	// DEFAULT_ROUTE_RECONNECT Route reconnect intervals.
	DEFAULT_ROUTE_RECONNECT = 1 * time.Second

	// DEFAULT_ROUTE_DIAL Route dial timeout.
	DEFAULT_ROUTE_DIAL = 1 * time.Second

	// PROTO_SNIPPET_SIZE is the default size of proto to print on parse errors.
	PROTO_SNIPPET_SIZE = 32

	// MAX_MSG_ARGS Maximum possible number of arguments from MSG proto.
	MAX_MSG_ARGS = 4

	// MAX_PUB_ARGS Maximum possible number of arguments from PUB proto.
	MAX_PUB_ARGS = 3
)
<|endoftext|>


<commit_before><commit_msg>delete unused file<commit_after><|endoftext|>


<commit_before>// common.go
//
// This file implements the configuration part for when you need the API
// key to modify things in the Atlas configuration and manage measurements.

package atlas

import (
    "regexp"
    "errors"
	"github.com/sendgrid/rest"
	"fmt"
	"encoding/json"
)

const (
	apiEndpoint = "https://atlas.ripe.net/api/v2"

	ourVersion = "0.9"
)

// APIKey is the API key
var APIKey string

// ErrInvalidMeasurementType is a new error
var ErrInvalidMeasurementType = errors.New("invalid measurement type")

// ErrInvalidAPIKey is returned when the key is invalid
var ErrInvalidAPIKey = errors.New("invalid API key")

// SetAuth stores the credentials for later use
func SetAuth(key string) {
	APIKey = key
}

// HasAPIKey returns whether an API key is stored
func HasAPIKey() (string, bool) {
	if APIKey == "" {
		return "", false
	}
	return APIKey, true
}

// getPageNum returns the value of the page= parameter
func getPageNum(url string) (page string) {
	re := regexp.MustCompile(`page=(\d+)`)
	if m := re.FindStringSubmatch(url); len(m) >= 1 {
		return m[1]
	}
	return ""
}

// prepareRequest insert all pre-defined stuff
func prepareRequest(what string) (req rest.Request) {
	endPoint := apiEndpoint + fmt.Sprintf("/%s/", what)
	key, ok := HasAPIKey()

	// Add at least one option, the APIkey if present
	hdrs := make(map[string]string)
	opts := make(map[string]string)

	// Insert our sig
	hdrs["User-Agent"] = fmt.Sprintf("ripe-atlas/%s", ourVersion)

	// Insert key
	if ok {
		opts["key"] = key
	}

	req = rest.Request{
		BaseURL: endPoint,
		Headers: hdrs,
		QueryParams: opts,
	}
	return
}

// handleAPIResponse check status code & errors from the API
func handleAPIResponse(r *rest.Response) (err error) {
	if r.StatusCode != 200 || r.StatusCode != 0 {
		var aerr APIError

		err = json.Unmarshal([]byte(r.Body), &aerr)
		err = fmt.Errorf("status: %d code: %d - r:%v\n",
			aerr.Error.Status,
			aerr.Error.Code,
			aerr.Error.Detail)
	}
	return
}
<commit_msg>Fix expression.<commit_after>// common.go
//
// This file implements the configuration part for when you need the API
// key to modify things in the Atlas configuration and manage measurements.

package atlas

import (
    "regexp"
    "errors"
	"github.com/sendgrid/rest"
	"fmt"
	"encoding/json"
)

const (
	apiEndpoint = "https://atlas.ripe.net/api/v2"

	ourVersion = "0.9"
)

// APIKey is the API key
var APIKey string

// ErrInvalidMeasurementType is a new error
var ErrInvalidMeasurementType = errors.New("invalid measurement type")

// ErrInvalidAPIKey is returned when the key is invalid
var ErrInvalidAPIKey = errors.New("invalid API key")

// SetAuth stores the credentials for later use
func SetAuth(key string) {
	APIKey = key
}

// HasAPIKey returns whether an API key is stored
func HasAPIKey() (string, bool) {
	if APIKey == "" {
		return "", false
	}
	return APIKey, true
}

// getPageNum returns the value of the page= parameter
func getPageNum(url string) (page string) {
	re := regexp.MustCompile(`page=(\d+)`)
	if m := re.FindStringSubmatch(url); len(m) >= 1 {
		return m[1]
	}
	return ""
}

// prepareRequest insert all pre-defined stuff
func prepareRequest(what string) (req rest.Request) {
	endPoint := apiEndpoint + fmt.Sprintf("/%s/", what)
	key, ok := HasAPIKey()

	// Add at least one option, the APIkey if present
	hdrs := make(map[string]string)
	opts := make(map[string]string)

	// Insert our sig
	hdrs["User-Agent"] = fmt.Sprintf("ripe-atlas/%s", ourVersion)

	// Insert key
	if ok {
		opts["key"] = key
	}

	req = rest.Request{
		BaseURL: endPoint,
		Headers: hdrs,
		QueryParams: opts,
	}
	return
}

// handleAPIResponse check status code & errors from the API
func handleAPIResponse(r *rest.Response) (err error) {
	if r.StatusCode != 200 {
		if r.StatusCode != 0 {
			var aerr APIError

			err = json.Unmarshal([]byte(r.Body), &aerr)
			err = fmt.Errorf("status: %d code: %d - r:%v\n",
				aerr.Error.Status,
				aerr.Error.Code,
				aerr.Error.Detail)
		}
	}
	return
}
<|endoftext|>


<commit_before>package shakers

import (
	"reflect"
	"time"

	check "gopkg.in/check.v1"
)

// As a commodity, we bring all check.Checker variables into the current namespace to avoid having
// to think about check.X versus checker.X.
var (
	DeepEquals   = check.DeepEquals
	Equals       = check.Equals
	ErrorMatches = check.ErrorMatches
	FitsTypeOf   = check.FitsTypeOf
	HasLen       = check.HasLen
	Implements   = check.Implements
	IsNil        = check.IsNil
	Matches      = check.Matches
	Not          = check.Not
	NotNil       = check.NotNil
	PanicMatches = check.PanicMatches
	Panics       = check.Panics
)

// Equaler is an interface implemented if the type has a Equal method.
// This is used to compare struct using shakers.Equals.
type Equaler interface {
	Equal(Equaler) bool
}

// Equals checker verifies the obtained value is equal to the specified one.
// It's is smart in a wait that it supports several *types* (built-in, Equaler,
// time.Time)
//
//    c.Assert(myStruct, Equals, aStruct, check.Commentf("bouuuhh"))
//    c.Assert(myTime, Equals, aTime, check.Commentf("bouuuhh"))
//
var Equals check.Checker = &equalChecker{
	&check.CheckerInfo{
		Name:   "Equals",
		Params: []string{"obtained", "expected"},
	},
}

type equalChecker struct {
	*check.CheckerInfo
}

func (checker *equalChecker) Check(params []interface{}, names []string) (bool, string) {
	return isEqual(params[0], params[1])
}

func isEqual(obtained, expected interface{}) (bool, string) {
	switch obtained.(type) {
	case time.Time:
		return timeEquals(obtained, expected)
	case Equaler:
		return equalerEquals(obtained, expected)
	default:
		if reflect.TypeOf(obtained) != reflect.TypeOf(expected) {
			return false, "obtained value and expected value have not the same type."
		}
		return obtained == expected, ""
	}
}

func equalerEquals(obtained, expected interface{}) (bool, string) {
	expectedEqualer, ok := expected.(Equaler)
	if !ok {
		return false, "expected value must be an Equaler - implementing Equal(Equaler)."
	}
	obtainedEqualer, ok := obtained.(Equaler)
	if !ok {
		return false, "obtained value must be an Equaler - implementing Equal(Equaler)."
	}
	return obtainedEqualer.Equal(expectedEqualer), ""
}
<commit_msg>Fix multiple Equals..<commit_after>package shakers

import (
	"reflect"
	"time"

	check "gopkg.in/check.v1"
)

// As a commodity, we bring all check.Checker variables into the current namespace to avoid having
// to think about check.X versus checker.X.
var (
	DeepEquals   = check.DeepEquals
	ErrorMatches = check.ErrorMatches
	FitsTypeOf   = check.FitsTypeOf
	HasLen       = check.HasLen
	Implements   = check.Implements
	IsNil        = check.IsNil
	Matches      = check.Matches
	Not          = check.Not
	NotNil       = check.NotNil
	PanicMatches = check.PanicMatches
	Panics       = check.Panics
)

// Equaler is an interface implemented if the type has a Equal method.
// This is used to compare struct using shakers.Equals.
type Equaler interface {
	Equal(Equaler) bool
}

// Equals checker verifies the obtained value is equal to the specified one.
// It's is smart in a wait that it supports several *types* (built-in, Equaler,
// time.Time)
//
//    c.Assert(myStruct, Equals, aStruct, check.Commentf("bouuuhh"))
//    c.Assert(myTime, Equals, aTime, check.Commentf("bouuuhh"))
//
var Equals check.Checker = &equalChecker{
	&check.CheckerInfo{
		Name:   "Equals",
		Params: []string{"obtained", "expected"},
	},
}

type equalChecker struct {
	*check.CheckerInfo
}

func (checker *equalChecker) Check(params []interface{}, names []string) (bool, string) {
	return isEqual(params[0], params[1])
}

func isEqual(obtained, expected interface{}) (bool, string) {
	switch obtained.(type) {
	case time.Time:
		return timeEquals(obtained, expected)
	case Equaler:
		return equalerEquals(obtained, expected)
	default:
		if reflect.TypeOf(obtained) != reflect.TypeOf(expected) {
			return false, "obtained value and expected value have not the same type."
		}
		return obtained == expected, ""
	}
}

func equalerEquals(obtained, expected interface{}) (bool, string) {
	expectedEqualer, ok := expected.(Equaler)
	if !ok {
		return false, "expected value must be an Equaler - implementing Equal(Equaler)."
	}
	obtainedEqualer, ok := obtained.(Equaler)
	if !ok {
		return false, "obtained value must be an Equaler - implementing Equal(Equaler)."
	}
	return obtainedEqualer.Equal(expectedEqualer), ""
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/smira/aptly/debian"
)

func ListPackagesRefList(reflist *debian.PackageRefList) (err error) {
	fmt.Printf("Packages:\n")

	packageCollection := debian.NewPackageCollection(context.database)

	err = reflist.ForEach(func(key []byte) error {
		p, err := packageCollection.ByKey(key)
		if err != nil {
			return err
		}
		fmt.Printf("  %s\n", p)
		return nil
	})
	if err != nil {
		return fmt.Errorf("unable to load packages: %s", err)
	}

	return
}
<commit_msg>Add comment to exported function.<commit_after>package main

import (
	"fmt"
	"github.com/smira/aptly/debian"
)

//ListPackagesRefList shows list of packages in PackageRefList
func ListPackagesRefList(reflist *debian.PackageRefList) (err error) {
	fmt.Printf("Packages:\n")

	packageCollection := debian.NewPackageCollection(context.database)

	err = reflist.ForEach(func(key []byte) error {
		p, err := packageCollection.ByKey(key)
		if err != nil {
			return err
		}
		fmt.Printf("  %s\n", p)
		return nil
	})
	if err != nil {
		return fmt.Errorf("unable to load packages: %s", err)
	}

	return
}
<|endoftext|>


<commit_before>package piazza

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

//---------------------------------------------------------------------------

// An AdminResponse represents the data returned from a call to a service's
// /admin API.
type AdminResponse struct {
	StartTime time.Time             `json:"starttime"`
	Uuidgen   *AdminResponseUuidgen `json:"uuidgen,omitempty"`
	Logger    *AdminResponseLogger  `json:"logger,omitempty"`
}

// AdminResponseUuidgen is the response to pz-uuidgen's /admin call
type AdminResponseUuidgen struct {
	NumRequests int `json:"num_requests"`
	NumUUIDs    int `json:"num_uuids"`
}

// AdminResponseLogger is the response to pz-logger's /admin call
type AdminResponseLogger struct {
	NumMessages int `json:"num_messages"`
}

//---------------------------------------------------------------------------

// ServerLogHandler adds traditional logging support to the http server handlers.
func ServerLogHandler(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		log.Printf("%s %s %s", r.RemoteAddr, r.Method, r.URL)
		handler.ServeHTTP(w, r)
	})
}

// ContentTypeJSON is the http content-type for JSON.
const ContentTypeJSON = "application/json"

// Put is because there is not http.Put.
func Put(url string, contentType string, body io.Reader) (*http.Response, error) {
	req, err := http.NewRequest("PUT", url, body)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", contentType)
	client := &http.Client{}
	return client.Do(req)
}

//---------------------------------------------------------------------------

// LogMessage represents the contents of a messge for the logger service.
// All fields are required.
type LogMessage struct {
	Service  string `json:"service"`
	Address  string `json:"address"`
	Time     string `json:"time"`
	Severity string `json:"severity"`
	Message  string `json:"message"`
}

// ToString returns a LogMessage as a formatted string.
func (mssg *LogMessage) ToString() string {
	s := fmt.Sprintf("[%s, %s, %s, %s, %s]",
		mssg.Service, mssg.Address, mssg.Time, mssg.Severity, mssg.Message)
	return s
}

// SeverityDebug is for log messages that are only used in development.
const SeverityDebug = "Debug"

// SeverityInfo is for log messages that are only informative, no action needed.
const SeverityInfo = "Info"

// SeverityWarning is for log messages that indicate possible problems. Execution continues normally.
const SeverityWarning = "Warning"

// SeverityError is for log messages that indicate something went wrong. The problem is usually handled and execution continues.
const SeverityError = "Error"

// SeverityFatal is for log messages that indicate an internal error and the system is likely now unstable. These should never happen.
const SeverityFatal = "Fatal"

// Validate checks to make sure a LogMessage is properly filled out. If not, a non-nil error is returned.
func (mssg *LogMessage) Validate() error {
	if mssg.Service == "" {
		return errors.New("required field 'service' not set")
	}
	if mssg.Address == "" {
		return errors.New("required field 'address' not set")
	}
	if mssg.Time == "" {
		return errors.New("required field 'time' not set")
	}
	if mssg.Severity == "" {
		return errors.New("required field 'severity' not set")
	}
	if mssg.Message == "" {
		return errors.New("required field 'message' not set")
	}

	ok := false
	for _, code := range [...]string{SeverityDebug, SeverityInfo, SeverityWarning, SeverityError, SeverityFatal} {
		if mssg.Severity == code {
			ok = true
			break
		}
	}
	if !ok {
		return errors.New("invalid 'severity' setting")
	}

	return nil
}

// Log sends a LogMessage to the logger.
func Log(service string, address string, severity string, message string) error {

	address, err := GetServiceAddress("pz-logger")
	if err != nil {
		return err
	}

	mssg := LogMessage{Service: service, Address: address, Severity: severity, Message: message, Time: time.Now().String()}
	data, err := json.Marshal(mssg)
	if err != nil {
		return err
	}

	resp, err := http.Post(address, ContentTypeJSON, bytes.NewBuffer(data))
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusOK {
		return errors.New(resp.Status)
	}

	return nil
}

//---------------------------------------------------------------------------

// ReadFrom is a convenience function that returns the bytes taken from a Reader.
// The reader will be closed if necessary.
func ReadFrom(reader io.Reader) ([]byte, error) {
	switch reader.(type) {
	case io.Closer:
		defer reader.(io.Closer).Close()
	}

	data, err := ioutil.ReadAll(reader)
	if err != nil {
		return nil, err
	}
	return data, err
}

//---------------------------------------------------------------------------

// singelton
var registryURL string

type registryItemData struct {
	Type    string `json:"type"`
	Address string `json:"address"`
}
type registryItem struct {
	Name string           `json:"name"`
	Data registryItemData `json:"data"`
}

// RegistryInit initialies the Discovery service from pz-discovery.
func RegistryInit(url string) {
	registryURL = url + "/api/v1/resources"
}

// RegisterService adds the given service to the discovery system.
func RegisterService(name string, itemtype string, url string) error {

	m := registryItem{Name: name, Data: registryItemData{Type: itemtype, Address: url}}

	data, err := json.Marshal(m)
	if err != nil {
		return err
	}

	resp, err := Put(registryURL, ContentTypeJSON, bytes.NewBuffer(data))
	if err != nil {
		return err
	}
	if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK {
		return errors.New(resp.Status)
	}

	return nil
}

// GetServiceAddress returns the URL of the given service.
// If the service is not found, a non-nil error is returned.
func GetServiceAddress(name string) (string, error) {

	target := fmt.Sprintf("%s/%s", registryURL, name)

	resp, err := http.Get(target)
	if err != nil {
		return "", err
	}

	if resp.StatusCode != http.StatusOK {
		return "", errors.New(resp.Status)
	}

	data, err := ReadFrom(resp.Body)
	if err != nil {
		return "", err
	}

	var m registryItemData
	err = json.Unmarshal(data, &m)
	if err != nil {
		return "", err
	}

	return m.Address, nil
}
<commit_msg>added Delete<commit_after>package piazza

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

//---------------------------------------------------------------------------

// An AdminResponse represents the data returned from a call to a service's
// /admin API.
type AdminResponse struct {
	StartTime time.Time             `json:"starttime"`
	Uuidgen   *AdminResponseUuidgen `json:"uuidgen,omitempty"`
	Logger    *AdminResponseLogger  `json:"logger,omitempty"`
}

// AdminResponseUuidgen is the response to pz-uuidgen's /admin call
type AdminResponseUuidgen struct {
	NumRequests int `json:"num_requests"`
	NumUUIDs    int `json:"num_uuids"`
}

// AdminResponseLogger is the response to pz-logger's /admin call
type AdminResponseLogger struct {
	NumMessages int `json:"num_messages"`
}

//---------------------------------------------------------------------------

// ServerLogHandler adds traditional logging support to the http server handlers.
func ServerLogHandler(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		log.Printf("%s %s %s", r.RemoteAddr, r.Method, r.URL)
		handler.ServeHTTP(w, r)
	})
}

// ContentTypeJSON is the http content-type for JSON.
const ContentTypeJSON = "application/json"

// Put is because there is no http.Put.
func Put(url string, contentType string, body io.Reader) (*http.Response, error) {
	req, err := http.NewRequest("PUT", url, body)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", contentType)
	client := &http.Client{}
	return client.Do(req)
}

// Delete is because there is no http.Delete.
func Delete(url string) (*http.Response, error) {
	req, err := http.NewRequest("DELETE", url, nil)
	if err != nil {
		return nil, err
	}

	client := &http.Client{}
	return client.Do(req)
}

//---------------------------------------------------------------------------

// LogMessage represents the contents of a messge for the logger service.
// All fields are required.
type LogMessage struct {
	Service  string `json:"service"`
	Address  string `json:"address"`
	Time     string `json:"time"`
	Severity string `json:"severity"`
	Message  string `json:"message"`
}

// ToString returns a LogMessage as a formatted string.
func (mssg *LogMessage) ToString() string {
	s := fmt.Sprintf("[%s, %s, %s, %s, %s]",
		mssg.Service, mssg.Address, mssg.Time, mssg.Severity, mssg.Message)
	return s
}

// SeverityDebug is for log messages that are only used in development.
const SeverityDebug = "Debug"

// SeverityInfo is for log messages that are only informative, no action needed.
const SeverityInfo = "Info"

// SeverityWarning is for log messages that indicate possible problems. Execution continues normally.
const SeverityWarning = "Warning"

// SeverityError is for log messages that indicate something went wrong. The problem is usually handled and execution continues.
const SeverityError = "Error"

// SeverityFatal is for log messages that indicate an internal error and the system is likely now unstable. These should never happen.
const SeverityFatal = "Fatal"

// Validate checks to make sure a LogMessage is properly filled out. If not, a non-nil error is returned.
func (mssg *LogMessage) Validate() error {
	if mssg.Service == "" {
		return errors.New("required field 'service' not set")
	}
	if mssg.Address == "" {
		return errors.New("required field 'address' not set")
	}
	if mssg.Time == "" {
		return errors.New("required field 'time' not set")
	}
	if mssg.Severity == "" {
		return errors.New("required field 'severity' not set")
	}
	if mssg.Message == "" {
		return errors.New("required field 'message' not set")
	}

	ok := false
	for _, code := range [...]string{SeverityDebug, SeverityInfo, SeverityWarning, SeverityError, SeverityFatal} {
		if mssg.Severity == code {
			ok = true
			break
		}
	}
	if !ok {
		return errors.New("invalid 'severity' setting")
	}

	return nil
}

// Log sends a LogMessage to the logger.
func Log(service string, address string, severity string, message string) error {

	address, err := GetServiceAddress("pz-logger")
	if err != nil {
		return err
	}

	mssg := LogMessage{Service: service, Address: address, Severity: severity, Message: message, Time: time.Now().String()}
	data, err := json.Marshal(mssg)
	if err != nil {
		return err
	}

	resp, err := http.Post(address, ContentTypeJSON, bytes.NewBuffer(data))
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusOK {
		return errors.New(resp.Status)
	}

	return nil
}

//---------------------------------------------------------------------------

// ReadFrom is a convenience function that returns the bytes taken from a Reader.
// The reader will be closed if necessary.
func ReadFrom(reader io.Reader) ([]byte, error) {
	switch reader.(type) {
	case io.Closer:
		defer reader.(io.Closer).Close()
	}

	data, err := ioutil.ReadAll(reader)
	if err != nil {
		return nil, err
	}
	return data, err
}

//---------------------------------------------------------------------------

// singelton
var registryURL string

type registryItemData struct {
	Type    string `json:"type"`
	Address string `json:"address"`
}
type registryItem struct {
	Name string           `json:"name"`
	Data registryItemData `json:"data"`
}

// RegistryInit initialies the Discovery service from pz-discovery.
func RegistryInit(url string) {
	registryURL = url + "/api/v1/resources"
}

// RegisterService adds the given service to the discovery system.
func RegisterService(name string, itemtype string, url string) error {

	m := registryItem{Name: name, Data: registryItemData{Type: itemtype, Address: url}}

	data, err := json.Marshal(m)
	if err != nil {
		return err
	}

	resp, err := Put(registryURL, ContentTypeJSON, bytes.NewBuffer(data))
	if err != nil {
		return err
	}
	if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK {
		return errors.New(resp.Status)
	}

	return nil
}

// GetServiceAddress returns the URL of the given service.
// If the service is not found, a non-nil error is returned.
func GetServiceAddress(name string) (string, error) {

	target := fmt.Sprintf("%s/%s", registryURL, name)

	resp, err := http.Get(target)
	if err != nil {
		return "", err
	}

	if resp.StatusCode != http.StatusOK {
		return "", errors.New(resp.Status)
	}

	data, err := ReadFrom(resp.Body)
	if err != nil {
		return "", err
	}

	var m registryItemData
	err = json.Unmarshal(data, &m)
	if err != nil {
		return "", err
	}

	return m.Address, nil
}
<|endoftext|>


<commit_before>package mq

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

	"github.com/arschles/gorion"
	"golang.org/x/net/context"
)

// Scheme is "http" or "https"
type Scheme string

// String converts a Scheme to a printable string
func (s Scheme) String() string {
	return string(s)
}

const (
	// SchemeHTTP represents http
	SchemeHTTP = "http"
	// SchemeHTTPS represents https
	SchemeHTTPS     = "https"
	applicationJSON = "application/json"
	oauth           = "OAuth"
)

type httpClient struct {
	scheme     Scheme
	host       string
	port       uint16
	transport  *http.Transport
	client     *http.Client
	oauthToken string
}

// NewHTTPClient returns a Client implementation that can talk to the IronMQ v3
// API documented at http://dev.iron.io/mq/3/reference/api/
func NewHTTPClient(scheme Scheme, host string, port uint16) Client {
	transport := &http.Transport{}
	client := &http.Client{Transport: transport}
	return &httpClient{
		scheme:    scheme,
		host:      host,
		port:      port,
		transport: transport,
		client:    client,
	}
}

// headers sets json and oauth headers on r
func (h *httpClient) newReq(method, token, projID, path string, body io.Reader) (*http.Request, error) {
	urlStr := fmt.Sprintf("%s://%s:%d/3/projects/%s/%s", h.scheme, h.host, h.port, projID, path)
	req, err := http.NewRequest(method, urlStr, body)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "OAuth "+token)
	return req, nil
}

type enqueueReq struct {
	Messages []NewMessage `json:"messages"`
}

// Enqueue posts messages to IronMQ using the API defined at http://dev.iron.io/mq/3/reference/api/#post-messages
func (h *httpClient) Enqueue(ctx context.Context, token, projID, qName string, msgs []NewMessage) (*Enqueued, error) {
	reqBody := &bytes.Buffer{}
	if err := json.NewEncoder(reqBody).Encode(enqueueReq{Messages: msgs}); err != nil {
		return nil, err
	}

	req, err := h.newReq("POST", token, projID, fmt.Sprintf("queues/%s/messages", qName), reqBody)
	if err != nil {
		return nil, err
	}
	ret := new(Enqueued)
	doFunc := func(resp *http.Response, err error) error {
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		if err := json.NewDecoder(resp.Body).Decode(ret); err != nil {
			return err
		}
		return nil
	}
	if err := gorion.HTTPDo(ctx, h.client, h.transport, req, doFunc); err != nil {
		return nil, err
	}
	return ret, nil
}

type dequeueReq struct {
	Num     int  `json:"n"`
	Timeout int  `json:"timeout"`
	Wait    int  `json:"wait"`
	Delete  bool `json:"delete"`
}

type dequeueResp struct {
	Messages []DequeuedMessage `json:"messages"`
}

// Dequeue gets messages from IronMQ using the API defined at http://dev.iron.io/mq/3/reference/api/#reserve-messages
func (h *httpClient) Dequeue(ctx context.Context, token, projID, qName string, num int, timeout Timeout, wait Wait, delete bool) ([]DequeuedMessage, error) {
	if !timeoutInRange(timeout) {
		return nil, ErrTimeoutOutOfRange
	}
	if !waitInRange(wait) {
		return nil, ErrWaitOutOfRange
	}

	body := &bytes.Buffer{}
	if err := json.NewEncoder(body).Encode(dequeueReq{Num: num, Timeout: int(timeout), Wait: int(wait), Delete: delete}); err != nil {
		return nil, err
	}
	req, err := h.newReq("POST", token, projID, fmt.Sprintf("queues/%s/reservations", qName), body)
	if err != nil {
		return nil, err
	}
	ret := new(dequeueResp)
	doFunc := func(resp *http.Response, err error) error {
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		if err := json.NewDecoder(resp.Body).Decode(ret); err != nil {
			return err
		}
		return nil
	}
	if err := gorion.HTTPDo(ctx, h.client, h.transport, req, doFunc); err != nil {
		return nil, err
	}
	return ret.Messages, nil
}

type deleteReservedReq struct {
	ReservationID string `json:"reservation_id"`
}

func (h *httpClient) DeleteReserved(ctx context.Context, token, projID, qName string, messageID int, reservationID string) (*Deleted, error) {
	body := &bytes.Buffer{}
	if err := json.NewEncoder(body).Encode(deleteReservedReq{ReservationID: reservationID}); err != nil {
		return nil, err
	}
	req, err := h.newReq("DELETE", token, projID, fmt.Sprintf("queues/%s/messages/%d", qName, messageID), body)
	if err != nil {
		return nil, err
	}
	ret := new(Deleted)
	doFunc := func(resp *http.Response, err error) error {
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		if err := json.NewDecoder(resp.Body).Decode(ret); err != nil {
			return err
		}
		return nil
	}
	if err := gorion.HTTPDo(ctx, h.client, h.transport, req, doFunc); err != nil {
		return nil, err
	}
	return ret, nil
}
<commit_msg>adding scheme 'constructor'<commit_after>package mq

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

	"github.com/arschles/gorion"
	"golang.org/x/net/context"
)

// Scheme is "http" or "https"
type Scheme string

// SchemeFromString returns a Scheme representation of s. If s is not a supported scheme, returns ErrInvalidScheme
func SchemeFromString(s string) (Scheme, error) {
	switch s {
	case SchemeHTTP:
		return SchemeHTTP, nil
	case SchemeHTTPS:
		return SchemeHTTPS, nil
	default:
		return Scheme(""), ErrInvalidScheme
	}
}

// String converts a Scheme to a printable string
func (s Scheme) String() string {
	return string(s)
}

var (
	ErrInvalidScheme = errors.New("invalid scheme")
)

const (
	// SchemeHTTP represents http
	SchemeHTTP = "http"
	// SchemeHTTPS represents https
	SchemeHTTPS     = "https"
	applicationJSON = "application/json"
	oauth           = "OAuth"
)

type httpClient struct {
	scheme     Scheme
	host       string
	port       uint16
	transport  *http.Transport
	client     *http.Client
	oauthToken string
}

// NewHTTPClient returns a Client implementation that can talk to the IronMQ v3
// API documented at http://dev.iron.io/mq/3/reference/api/
func NewHTTPClient(scheme Scheme, host string, port uint16) Client {
	transport := &http.Transport{}
	client := &http.Client{Transport: transport}
	return &httpClient{
		scheme:    scheme,
		host:      host,
		port:      port,
		transport: transport,
		client:    client,
	}
}

// headers sets json and oauth headers on r
func (h *httpClient) newReq(method, token, projID, path string, body io.Reader) (*http.Request, error) {
	urlStr := fmt.Sprintf("%s://%s:%d/3/projects/%s/%s", h.scheme, h.host, h.port, projID, path)
	req, err := http.NewRequest(method, urlStr, body)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "OAuth "+token)
	return req, nil
}

type enqueueReq struct {
	Messages []NewMessage `json:"messages"`
}

// Enqueue posts messages to IronMQ using the API defined at http://dev.iron.io/mq/3/reference/api/#post-messages
func (h *httpClient) Enqueue(ctx context.Context, token, projID, qName string, msgs []NewMessage) (*Enqueued, error) {
	reqBody := &bytes.Buffer{}
	if err := json.NewEncoder(reqBody).Encode(enqueueReq{Messages: msgs}); err != nil {
		return nil, err
	}

	req, err := h.newReq("POST", token, projID, fmt.Sprintf("queues/%s/messages", qName), reqBody)
	if err != nil {
		return nil, err
	}
	ret := new(Enqueued)
	doFunc := func(resp *http.Response, err error) error {
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		if err := json.NewDecoder(resp.Body).Decode(ret); err != nil {
			return err
		}
		return nil
	}
	if err := gorion.HTTPDo(ctx, h.client, h.transport, req, doFunc); err != nil {
		return nil, err
	}
	return ret, nil
}

type dequeueReq struct {
	Num     int  `json:"n"`
	Timeout int  `json:"timeout"`
	Wait    int  `json:"wait"`
	Delete  bool `json:"delete"`
}

type dequeueResp struct {
	Messages []DequeuedMessage `json:"messages"`
}

// Dequeue gets messages from IronMQ using the API defined at http://dev.iron.io/mq/3/reference/api/#reserve-messages
func (h *httpClient) Dequeue(ctx context.Context, token, projID, qName string, num int, timeout Timeout, wait Wait, delete bool) ([]DequeuedMessage, error) {
	if !timeoutInRange(timeout) {
		return nil, ErrTimeoutOutOfRange
	}
	if !waitInRange(wait) {
		return nil, ErrWaitOutOfRange
	}

	body := &bytes.Buffer{}
	if err := json.NewEncoder(body).Encode(dequeueReq{Num: num, Timeout: int(timeout), Wait: int(wait), Delete: delete}); err != nil {
		return nil, err
	}
	req, err := h.newReq("POST", token, projID, fmt.Sprintf("queues/%s/reservations", qName), body)
	if err != nil {
		return nil, err
	}
	ret := new(dequeueResp)
	doFunc := func(resp *http.Response, err error) error {
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		if err := json.NewDecoder(resp.Body).Decode(ret); err != nil {
			return err
		}
		return nil
	}
	if err := gorion.HTTPDo(ctx, h.client, h.transport, req, doFunc); err != nil {
		return nil, err
	}
	return ret.Messages, nil
}

type deleteReservedReq struct {
	ReservationID string `json:"reservation_id"`
}

func (h *httpClient) DeleteReserved(ctx context.Context, token, projID, qName string, messageID int, reservationID string) (*Deleted, error) {
	body := &bytes.Buffer{}
	if err := json.NewEncoder(body).Encode(deleteReservedReq{ReservationID: reservationID}); err != nil {
		return nil, err
	}
	req, err := h.newReq("DELETE", token, projID, fmt.Sprintf("queues/%s/messages/%d", qName, messageID), body)
	if err != nil {
		return nil, err
	}
	ret := new(Deleted)
	doFunc := func(resp *http.Response, err error) error {
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		if err := json.NewDecoder(resp.Body).Decode(ret); err != nil {
			return err
		}
		return nil
	}
	if err := gorion.HTTPDo(ctx, h.client, h.transport, req, doFunc); err != nil {
		return nil, err
	}
	return ret, nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/log"
)

const (
	namespace = "nginx" // For Prometheus metrics.
)

var (
	listeningAddress = flag.String("telemetry.address", ":9113", "Address on which to expose metrics.")
	metricsEndpoint  = flag.String("telemetry.endpoint", "/metrics", "Path under which to expose metrics.")
	nginxScrapeURI   = flag.String("nginx.scrape_uri", "http://localhost/nginx_status", "URI to nginx stub status page")
	insecure         = flag.Bool("insecure", true, "Ignore server certificate if using https")
)

// Exporter collects nginx stats from the given URI and exports them using
// the prometheus metrics package.
type Exporter struct {
	URI    string
	mutex  sync.RWMutex
	client *http.Client

	scrapeFailures       prometheus.Counter
	processedConnections *prometheus.Desc
	currentConnections   *prometheus.GaugeVec
}

// NewExporter returns an initialized Exporter.
func NewExporter(uri string) *Exporter {
	return &Exporter{
		URI: uri,
		scrapeFailures: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Name:      "exporter_scrape_failures_total",
			Help:      "Number of errors while scraping nginx.",
		}),
		processedConnections: prometheus.NewDesc(
			prometheus.BuildFQName(namespace, "", "connections_processed_total"),
			"Number of connections processed by nginx",
			[]string{"stage"},
			nil,
		),
		currentConnections: prometheus.NewGaugeVec(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "connections_current",
			Help:      "Number of connections currently processed by nginx",
		},
			[]string{"state"},
		),
		client: &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: *insecure},
			},
		},
	}
}

// Describe describes all the metrics ever exported by the nginx exporter. It
// implements prometheus.Collector.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	ch <- e.processedConnections
	e.currentConnections.Describe(ch)
	e.scrapeFailures.Describe(ch)
}

func (e *Exporter) collect(ch chan<- prometheus.Metric) error {
	resp, err := e.client.Get(e.URI)
	if err != nil {
		return fmt.Errorf("Error scraping nginx: %v", err)
	}

	data, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		if err != nil {
			data = []byte(err.Error())
		}
		return fmt.Errorf("Status %s (%d): %s", resp.Status, resp.StatusCode, data)
	}

	// Parsing results
	lines := strings.Split(string(data), "\n")
	if len(lines) != 5 {
		return fmt.Errorf("Unexpected number of lines in status: %v", lines)
	}

	// active connections
	parts := strings.Split(lines[0], ":")
	if len(parts) != 2 {
		return fmt.Errorf("Unexpected first line: %s", lines[0])
	}
	v, err := strconv.Atoi(strings.TrimSpace(parts[1]))
	if err != nil {
		return err
	}
	e.currentConnections.WithLabelValues("active").Set(float64(v))

	// processed connections
	parts = strings.Fields(lines[2])
	if len(parts) != 3 {
		return fmt.Errorf("Unexpected third line: %s", lines[2])
	}
	v, err = strconv.Atoi(strings.TrimSpace(parts[0]))
	if err != nil {
		return err
	}
	ch <- prometheus.MustNewConstMetric(e.processedConnections, prometheus.CounterValue, float64(v), "accepted")
	v, err = strconv.Atoi(strings.TrimSpace(parts[1]))
	if err != nil {
		return err
	}
	ch <- prometheus.MustNewConstMetric(e.processedConnections, prometheus.CounterValue, float64(v), "handled")
	v, err = strconv.Atoi(strings.TrimSpace(parts[2]))
	if err != nil {
		return err
	}
	ch <- prometheus.MustNewConstMetric(e.processedConnections, prometheus.CounterValue, float64(v), "any")

	// current connections
	parts = strings.Fields(lines[3])
	if len(parts) != 6 {
		return fmt.Errorf("Unexpected fourth line: %s", lines[3])
	}
	v, err = strconv.Atoi(strings.TrimSpace(parts[1]))
	if err != nil {
		return err
	}
	e.currentConnections.WithLabelValues("reading").Set(float64(v))
	v, err = strconv.Atoi(strings.TrimSpace(parts[3]))
	if err != nil {
		return err
	}

	e.currentConnections.WithLabelValues("writing").Set(float64(v))
	v, err = strconv.Atoi(strings.TrimSpace(parts[5]))
	if err != nil {
		return err
	}
	e.currentConnections.WithLabelValues("waiting").Set(float64(v))
	return nil
}

// Collect fetches the stats from configured nginx location and delivers them
// as Prometheus metrics. It implements prometheus.Collector.
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	e.mutex.Lock() // To protect metrics from concurrent collects.
	defer e.mutex.Unlock()
	if err := e.collect(ch); err != nil {
		log.Printf("Error scraping nginx: %s", err)
		e.scrapeFailures.Inc()
		e.scrapeFailures.Collect(ch)
	}
	return
}

func main() {
	flag.Parse()

	exporter := NewExporter(*nginxScrapeURI)
	prometheus.MustRegister(exporter)

	log.Printf("Starting Server: %s", *listeningAddress)
	http.Handle(*metricsEndpoint, prometheus.Handler())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
			<title>Nginx Exporter</title>

Nginx Exporter

<a href="` + *metricsEndpoint + `">Metrics

			`))
	})

	log.Fatal(http.ListenAndServe(*listeningAddress, nil))
}
<commit_msg>an earlier merge accidentally removed this line, adding it back in<commit_after>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/log"
)

const (
	namespace = "nginx" // For Prometheus metrics.
)

var (
	listeningAddress = flag.String("telemetry.address", ":9113", "Address on which to expose metrics.")
	metricsEndpoint  = flag.String("telemetry.endpoint", "/metrics", "Path under which to expose metrics.")
	nginxScrapeURI   = flag.String("nginx.scrape_uri", "http://localhost/nginx_status", "URI to nginx stub status page")
	insecure         = flag.Bool("insecure", true, "Ignore server certificate if using https")
)

// Exporter collects nginx stats from the given URI and exports them using
// the prometheus metrics package.
type Exporter struct {
	URI    string
	mutex  sync.RWMutex
	client *http.Client

	scrapeFailures       prometheus.Counter
	processedConnections *prometheus.Desc
	currentConnections   *prometheus.GaugeVec
}

// NewExporter returns an initialized Exporter.
func NewExporter(uri string) *Exporter {
	return &Exporter{
		URI: uri,
		scrapeFailures: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: namespace,
			Name:      "exporter_scrape_failures_total",
			Help:      "Number of errors while scraping nginx.",
		}),
		processedConnections: prometheus.NewDesc(
			prometheus.BuildFQName(namespace, "", "connections_processed_total"),
			"Number of connections processed by nginx",
			[]string{"stage"},
			nil,
		),
		currentConnections: prometheus.NewGaugeVec(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "connections_current",
			Help:      "Number of connections currently processed by nginx",
		},
			[]string{"state"},
		),
		client: &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: *insecure},
			},
		},
	}
}

// Describe describes all the metrics ever exported by the nginx exporter. It
// implements prometheus.Collector.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	ch <- e.processedConnections
	e.currentConnections.Describe(ch)
	e.scrapeFailures.Describe(ch)
}

func (e *Exporter) collect(ch chan<- prometheus.Metric) error {
	resp, err := e.client.Get(e.URI)
	if err != nil {
		return fmt.Errorf("Error scraping nginx: %v", err)
	}

	data, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		if err != nil {
			data = []byte(err.Error())
		}
		return fmt.Errorf("Status %s (%d): %s", resp.Status, resp.StatusCode, data)
	}

	// Parsing results
	lines := strings.Split(string(data), "\n")
	if len(lines) != 5 {
		return fmt.Errorf("Unexpected number of lines in status: %v", lines)
	}

	// active connections
	parts := strings.Split(lines[0], ":")
	if len(parts) != 2 {
		return fmt.Errorf("Unexpected first line: %s", lines[0])
	}
	v, err := strconv.Atoi(strings.TrimSpace(parts[1]))
	if err != nil {
		return err
	}
	e.currentConnections.WithLabelValues("active").Set(float64(v))

	// processed connections
	parts = strings.Fields(lines[2])
	if len(parts) != 3 {
		return fmt.Errorf("Unexpected third line: %s", lines[2])
	}
	v, err = strconv.Atoi(strings.TrimSpace(parts[0]))
	if err != nil {
		return err
	}
	ch <- prometheus.MustNewConstMetric(e.processedConnections, prometheus.CounterValue, float64(v), "accepted")
	v, err = strconv.Atoi(strings.TrimSpace(parts[1]))
	if err != nil {
		return err
	}
	ch <- prometheus.MustNewConstMetric(e.processedConnections, prometheus.CounterValue, float64(v), "handled")
	v, err = strconv.Atoi(strings.TrimSpace(parts[2]))
	if err != nil {
		return err
	}
	ch <- prometheus.MustNewConstMetric(e.processedConnections, prometheus.CounterValue, float64(v), "any")

	// current connections
	parts = strings.Fields(lines[3])
	if len(parts) != 6 {
		return fmt.Errorf("Unexpected fourth line: %s", lines[3])
	}
	v, err = strconv.Atoi(strings.TrimSpace(parts[1]))
	if err != nil {
		return err
	}
	e.currentConnections.WithLabelValues("reading").Set(float64(v))
	v, err = strconv.Atoi(strings.TrimSpace(parts[3]))
	if err != nil {
		return err
	}

	e.currentConnections.WithLabelValues("writing").Set(float64(v))
	v, err = strconv.Atoi(strings.TrimSpace(parts[5]))
	if err != nil {
		return err
	}
	e.currentConnections.WithLabelValues("waiting").Set(float64(v))
	return nil
}

// Collect fetches the stats from configured nginx location and delivers them
// as Prometheus metrics. It implements prometheus.Collector.
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	e.mutex.Lock() // To protect metrics from concurrent collects.
	defer e.mutex.Unlock()
	if err := e.collect(ch); err != nil {
		log.Printf("Error scraping nginx: %s", err)
		e.scrapeFailures.Inc()
		e.scrapeFailures.Collect(ch)
	}
	e.currentConnections.Collect(ch)
	return
}

func main() {
	flag.Parse()

	exporter := NewExporter(*nginxScrapeURI)
	prometheus.MustRegister(exporter)

	log.Printf("Starting Server: %s", *listeningAddress)
	http.Handle(*metricsEndpoint, prometheus.Handler())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
			<title>Nginx Exporter</title>

Nginx Exporter

<a href="` + *metricsEndpoint + `">Metrics

			`))
	})

	log.Fatal(http.ListenAndServe(*listeningAddress, nil))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/freeformz/shh/config"
	"github.com/freeformz/shh/mm"
	"github.com/freeformz/shh/output"
	"github.com/freeformz/shh/pollers"
	"github.com/freeformz/shh/utils"
	"log"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var (
	signalChannel = make(chan os.Signal, 1)
)

func main() {
	measurements := make(chan *mm.Measurement, 100)

	mp := pollers.NewMultiPoller(measurements)

	signal.Notify(signalChannel, syscall.SIGINT)
	signal.Notify(signalChannel, syscall.SIGTERM)

	go func() {
		for sig := range signalChannel {
			mp.Exit()
			log.Fatal(utils.Slog{"signal": sig, "finished": time.Now(), "duration": time.Since(config.Start)})
		}
	}()

	if config.ProfilePort != config.DEFAULT_PROFILE_PORT {
		go func() {
			log.Println(http.ListenAndServe("localhost:"+config.ProfilePort, nil))
		}()
	}

	ctx := utils.Slog{"shh_start": true, "at": config.Start.Format(time.RFC3339Nano), "interval": config.Interval}
	fmt.Println(ctx)

	outputter, err := output.NewOutputter(config.Outputter, measurements)
	if err != nil {
		ctx.FatalError(err, "creating outputter")
	}
	outputter.Start()

	// poll at start
	go mp.Poll(time.Now())

	ticks := time.Tick(config.Interval)
	for {
		select {
		case tick := <-ticks:
			go mp.Poll(tick)
		}
	}
}
<commit_msg>Don't use a goroutine here<commit_after>package main

import (
	"fmt"
	"github.com/freeformz/shh/config"
	"github.com/freeformz/shh/mm"
	"github.com/freeformz/shh/output"
	"github.com/freeformz/shh/pollers"
	"github.com/freeformz/shh/utils"
	"log"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var (
	signalChannel = make(chan os.Signal, 1)
)

func main() {
	measurements := make(chan *mm.Measurement, 100)

	mp := pollers.NewMultiPoller(measurements)

	signal.Notify(signalChannel, syscall.SIGINT)
	signal.Notify(signalChannel, syscall.SIGTERM)

	go func() {
		for sig := range signalChannel {
			mp.Exit()
			log.Fatal(utils.Slog{"signal": sig, "finished": time.Now(), "duration": time.Since(config.Start)})
		}
	}()

	if config.ProfilePort != config.DEFAULT_PROFILE_PORT {
		go func() {
			log.Println(http.ListenAndServe("localhost:"+config.ProfilePort, nil))
		}()
	}

	ctx := utils.Slog{"shh_start": true, "at": config.Start.Format(time.RFC3339Nano), "interval": config.Interval}
	fmt.Println(ctx)

	outputter, err := output.NewOutputter(config.Outputter, measurements)
	if err != nil {
		ctx.FatalError(err, "creating outputter")
	}
	outputter.Start()

	// poll at start
	go mp.Poll(time.Now())

	ticks := time.Tick(config.Interval)
	for {
		select {
		case tick := <-ticks:
			mp.Poll(tick)
		}
	}
}
<|endoftext|>


<commit_before>package websocket

import (
	"fmt"
	"log"
	"sort"
	"strings"
	"sync"
	"time"
)

type SubscriptionRequest struct {
	SubID string `json:"subId"`
	Event string `json:"event"`

	// authenticated
	APIKey      string   `json:"apiKey,omitempty"`
	AuthSig     string   `json:"authSig,omitempty"`
	AuthPayload string   `json:"authPayload,omitempty"`
	AuthNonce   string   `json:"authNonce,omitempty"`
	Filter      []string `json:"filter,omitempty"`

	// unauthenticated
	Channel   string `json:"channel,omitempty"`
	Symbol    string `json:"symbol,omitempty"`
	Precision string `json:"prec,omitempty"`
	Frequency string `json:"freq,omitempty"`
	Key       string `json:"key,omitempty"`
	Len       string `json:"len,omitempty"`
	Pair      string `json:"pair,omitempty"`
}

func (s *SubscriptionRequest) String() string {
	if s.Key == "" && s.Channel != "" && s.Symbol != "" {
		return fmt.Sprintf("%s %s", s.Channel, s.Symbol)
	}
	if s.Channel != "" && s.Symbol != "" && s.Precision != "" && s.Frequency != "" {
		return fmt.Sprintf("%s %s %s %s", s.Channel, s.Symbol, s.Precision, s.Frequency)
	}
	if s.Channel != "" && s.Symbol != "" {
		return fmt.Sprintf("%s %s", s.Channel, s.Symbol)
	}
	return ""
}

type UnsubscribeRequest struct {
	Event  string `json:"event"`
	ChanID int64  `json:"chanId"`
}

type messageFactory func(chanID int64, raw []interface{}) (interface{}, error)

type subscription struct {
	ChanID  int64
	pending bool
	Public  bool

	Request *SubscriptionRequest

	// heartbeat timer
	hbInterval time.Duration
	hb         *time.Timer
	die        chan bool

	parentDisconnect chan error
}

func (s *subscription) activate() {
	s.heartbeat()
}

func (s *subscription) timeout() {
	s.parentDisconnect <- fmt.Errorf("heartbeat timed out on channel %d", s.ChanID)
}

// returns death channel
func (s *subscription) heartbeat() {
	if s.hb != nil {
		s.hb.Stop()
	}
	close(s.die) // terminate previous hb timeout
	s.die = make(chan bool)
	s.hb = time.AfterFunc(s.hbInterval, s.timeout)
	go func() {
		select {
		case <-s.die:
			s.hb.Stop()
			return
		}
	}()
}

func (s *subscription) stopHeartbeatTimeout() {
	s.die <- true
	close(s.die)
}

func isPublic(request *SubscriptionRequest) bool {
	switch request.Channel {
	case ChanBook:
		return true
	case ChanCandles:
		return true
	case ChanTicker:
		return true
	case ChanTrades:
		return true
	}
	return false
}

func newSubscription(request *SubscriptionRequest, interval time.Duration, parentDisconnect chan error) *subscription {
	return &subscription{
		Request:          request,
		pending:          true,
		Public:           isPublic(request),
		die:              make(chan bool), // kill pending heartbeats
		hbInterval:       interval,
		parentDisconnect: parentDisconnect, // disconnect parent in hb timeout
	}
}

func (s subscription) SubID() string {
	return s.Request.SubID
}

func (s subscription) Pending() bool {
	return s.pending
}

func newSubscriptions(heartbeatTimeout time.Duration) *subscriptions {
	return &subscriptions{
		subsBySubID:         make(map[string]*subscription),
		subsByChanID:        make(map[int64]*subscription),
		hbTimeout:           heartbeatTimeout,
		hbParentDisconnect:  make(chan error),
		hbChannelDisconnect: make(chan error),
	}
}

type subscriptions struct {
	lock sync.Mutex

	subsBySubID  map[string]*subscription // subscription map indexed by subscription ID
	subsByChanID map[int64]*subscription  // subscription map indexed by channel ID

	hbTimeout           time.Duration
	hbChannelDisconnect chan error // message sent when a subscription fails a heartbeat check
	hbParentDisconnect  chan error // parent listens to this channel to receive disconnect events
}

func (s *subscriptions) Empty() bool {
	s.lock.Lock()
	defer s.lock.Unlock()
	return len(s.subsBySubID) == 0
}

// SubscriptionSet is a typed version of an array of subscription pointers, intended to meet the sortable interface.
// We need to sort Reset()'s return values for tests with more than 1 subscription (range map order is undefined)
type SubscriptionSet []*subscription

func (s SubscriptionSet) Len() int {
	return len(s)
}
func (s SubscriptionSet) Less(i, j int) bool {
	return strings.Compare(s[i].SubID(), s[j].SubID()) < 0
}
func (s SubscriptionSet) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

// Reset clears all subscriptions from the currently managed list, and returns
// a slice of the existing subscriptions prior to reset.  Returns nil if no subscriptions exist.
func (s *subscriptions) Reset() []*subscription {
	s.lock.Lock()
	defer s.lock.Unlock()
	var subs []*subscription
	if len(s.subsBySubID) > 0 {
		subs = make([]*subscription, 0, len(s.subsBySubID))
		for _, sub := range s.subsBySubID {
			sub.stopHeartbeatTimeout()
			subs = append(subs, sub)
		}
		sort.Sort(SubscriptionSet(subs))
	}
	close(s.hbChannelDisconnect)
	close(s.hbParentDisconnect)
	s.subsBySubID = make(map[string]*subscription)
	s.subsByChanID = make(map[int64]*subscription)
	s.hbChannelDisconnect = make(chan error)
	s.hbParentDisconnect = make(chan error)
	go s.listenHeartbeats()
	return subs
}

func (s *subscriptions) listenHeartbeats() {
	if err := <-s.hbChannelDisconnect; err != nil {
		s.hbParentDisconnect <- err
	}
}

func (s *subscriptions) ListenDisconnect() <-chan error {
	return s.hbParentDisconnect
}

func (s *subscriptions) add(sub *SubscriptionRequest) *subscription {
	s.lock.Lock()
	defer s.lock.Unlock()
	subscription := newSubscription(sub, s.hbTimeout, s.hbChannelDisconnect)
	s.subsBySubID[sub.SubID] = subscription
	return subscription
}

func (s *subscriptions) removeByChannelID(chanID int64) error {
	s.lock.Lock()
	defer s.lock.Unlock()
	sub, ok := s.subsByChanID[chanID]
	if !ok {
		return fmt.Errorf("could not find channel ID %d", chanID)
	}
	delete(s.subsByChanID, chanID)
	if _, ok = s.subsBySubID[sub.SubID()]; ok {
		delete(s.subsBySubID, sub.SubID())
	}
	return nil
}

func (s *subscriptions) removeBySubscriptionID(subID string) error {
	s.lock.Lock()
	defer s.lock.Unlock()
	sub, ok := s.subsBySubID[subID]
	if !ok {
		return fmt.Errorf("could not find subscription ID %s to remove", subID)
	}
	// exists, remove both indices
	delete(s.subsBySubID, subID)
	if _, ok = s.subsByChanID[sub.ChanID]; ok {
		delete(s.subsByChanID, sub.ChanID)
	}
	return nil
}

func (s *subscriptions) activate(subID string, chanID int64) error {
	s.lock.Lock()
	defer s.lock.Unlock()
	if sub, ok := s.subsBySubID[subID]; ok {
		sub.pending = false
		sub.ChanID = chanID
		s.subsByChanID[chanID] = sub
		sub.activate()
		return nil
	}
	return fmt.Errorf("could not find subscription ID %s to activate", subID)
}

func (s *subscriptions) lookupByChannelID(chanID int64) (*subscription, error) {
	s.lock.Lock()
	defer s.lock.Unlock()
	if sub, ok := s.subsByChanID[chanID]; ok {
		return sub, nil
	}
	return nil, fmt.Errorf("could not find subscription for channel ID %d", chanID)
}

func (s *subscriptions) lookupBySubscriptionID(subID string) (*subscription, error) {
	s.lock.Lock()
	defer s.lock.Unlock()
	if sub, ok := s.subsBySubID[subID]; ok {
		return sub, nil
	}
	return nil, fmt.Errorf("could not find subscription ID %s", subID)
}

func (s *subscriptions) heartbeat(chanID int64) error {
	s.lock.Lock()
	defer s.lock.Unlock()
	if sub, ok := s.subsByChanID[chanID]; ok {
		sub.heartbeat()
	}
	return fmt.Errorf("could not find channel ID to update heartbeat %d", chanID)
}
<commit_msg>import error<commit_after>package websocket

import (
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"
)

type SubscriptionRequest struct {
	SubID string `json:"subId"`
	Event string `json:"event"`

	// authenticated
	APIKey      string   `json:"apiKey,omitempty"`
	AuthSig     string   `json:"authSig,omitempty"`
	AuthPayload string   `json:"authPayload,omitempty"`
	AuthNonce   string   `json:"authNonce,omitempty"`
	Filter      []string `json:"filter,omitempty"`

	// unauthenticated
	Channel   string `json:"channel,omitempty"`
	Symbol    string `json:"symbol,omitempty"`
	Precision string `json:"prec,omitempty"`
	Frequency string `json:"freq,omitempty"`
	Key       string `json:"key,omitempty"`
	Len       string `json:"len,omitempty"`
	Pair      string `json:"pair,omitempty"`
}

func (s *SubscriptionRequest) String() string {
	if s.Key == "" && s.Channel != "" && s.Symbol != "" {
		return fmt.Sprintf("%s %s", s.Channel, s.Symbol)
	}
	if s.Channel != "" && s.Symbol != "" && s.Precision != "" && s.Frequency != "" {
		return fmt.Sprintf("%s %s %s %s", s.Channel, s.Symbol, s.Precision, s.Frequency)
	}
	if s.Channel != "" && s.Symbol != "" {
		return fmt.Sprintf("%s %s", s.Channel, s.Symbol)
	}
	return ""
}

type UnsubscribeRequest struct {
	Event  string `json:"event"`
	ChanID int64  `json:"chanId"`
}

type messageFactory func(chanID int64, raw []interface{}) (interface{}, error)

type subscription struct {
	ChanID  int64
	pending bool
	Public  bool

	Request *SubscriptionRequest

	// heartbeat timer
	hbInterval time.Duration
	hb         *time.Timer
	die        chan bool

	parentDisconnect chan error
}

func (s *subscription) activate() {
	s.heartbeat()
}

func (s *subscription) timeout() {
	s.parentDisconnect <- fmt.Errorf("heartbeat timed out on channel %d", s.ChanID)
}

// returns death channel
func (s *subscription) heartbeat() {
	if s.hb != nil {
		s.hb.Stop()
	}
	close(s.die) // terminate previous hb timeout
	s.die = make(chan bool)
	s.hb = time.AfterFunc(s.hbInterval, s.timeout)
	go func() {
		select {
		case <-s.die:
			s.hb.Stop()
			return
		}
	}()
}

func (s *subscription) stopHeartbeatTimeout() {
	s.die <- true
	close(s.die)
}

func isPublic(request *SubscriptionRequest) bool {
	switch request.Channel {
	case ChanBook:
		return true
	case ChanCandles:
		return true
	case ChanTicker:
		return true
	case ChanTrades:
		return true
	}
	return false
}

func newSubscription(request *SubscriptionRequest, interval time.Duration, parentDisconnect chan error) *subscription {
	return &subscription{
		Request:          request,
		pending:          true,
		Public:           isPublic(request),
		die:              make(chan bool), // kill pending heartbeats
		hbInterval:       interval,
		parentDisconnect: parentDisconnect, // disconnect parent in hb timeout
	}
}

func (s subscription) SubID() string {
	return s.Request.SubID
}

func (s subscription) Pending() bool {
	return s.pending
}

func newSubscriptions(heartbeatTimeout time.Duration) *subscriptions {
	return &subscriptions{
		subsBySubID:         make(map[string]*subscription),
		subsByChanID:        make(map[int64]*subscription),
		hbTimeout:           heartbeatTimeout,
		hbParentDisconnect:  make(chan error),
		hbChannelDisconnect: make(chan error),
	}
}

type subscriptions struct {
	lock sync.Mutex

	subsBySubID  map[string]*subscription // subscription map indexed by subscription ID
	subsByChanID map[int64]*subscription  // subscription map indexed by channel ID

	hbTimeout           time.Duration
	hbChannelDisconnect chan error // message sent when a subscription fails a heartbeat check
	hbParentDisconnect  chan error // parent listens to this channel to receive disconnect events
}

func (s *subscriptions) Empty() bool {
	s.lock.Lock()
	defer s.lock.Unlock()
	return len(s.subsBySubID) == 0
}

// SubscriptionSet is a typed version of an array of subscription pointers, intended to meet the sortable interface.
// We need to sort Reset()'s return values for tests with more than 1 subscription (range map order is undefined)
type SubscriptionSet []*subscription

func (s SubscriptionSet) Len() int {
	return len(s)
}
func (s SubscriptionSet) Less(i, j int) bool {
	return strings.Compare(s[i].SubID(), s[j].SubID()) < 0
}
func (s SubscriptionSet) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

// Reset clears all subscriptions from the currently managed list, and returns
// a slice of the existing subscriptions prior to reset.  Returns nil if no subscriptions exist.
func (s *subscriptions) Reset() []*subscription {
	s.lock.Lock()
	defer s.lock.Unlock()
	var subs []*subscription
	if len(s.subsBySubID) > 0 {
		subs = make([]*subscription, 0, len(s.subsBySubID))
		for _, sub := range s.subsBySubID {
			sub.stopHeartbeatTimeout()
			subs = append(subs, sub)
		}
		sort.Sort(SubscriptionSet(subs))
	}
	close(s.hbChannelDisconnect)
	close(s.hbParentDisconnect)
	s.subsBySubID = make(map[string]*subscription)
	s.subsByChanID = make(map[int64]*subscription)
	s.hbChannelDisconnect = make(chan error)
	s.hbParentDisconnect = make(chan error)
	go s.listenHeartbeats()
	return subs
}

func (s *subscriptions) listenHeartbeats() {
	if err := <-s.hbChannelDisconnect; err != nil {
		s.hbParentDisconnect <- err
	}
}

func (s *subscriptions) ListenDisconnect() <-chan error {
	return s.hbParentDisconnect
}

func (s *subscriptions) add(sub *SubscriptionRequest) *subscription {
	s.lock.Lock()
	defer s.lock.Unlock()
	subscription := newSubscription(sub, s.hbTimeout, s.hbChannelDisconnect)
	s.subsBySubID[sub.SubID] = subscription
	return subscription
}

func (s *subscriptions) removeByChannelID(chanID int64) error {
	s.lock.Lock()
	defer s.lock.Unlock()
	sub, ok := s.subsByChanID[chanID]
	if !ok {
		return fmt.Errorf("could not find channel ID %d", chanID)
	}
	delete(s.subsByChanID, chanID)
	if _, ok = s.subsBySubID[sub.SubID()]; ok {
		delete(s.subsBySubID, sub.SubID())
	}
	return nil
}

func (s *subscriptions) removeBySubscriptionID(subID string) error {
	s.lock.Lock()
	defer s.lock.Unlock()
	sub, ok := s.subsBySubID[subID]
	if !ok {
		return fmt.Errorf("could not find subscription ID %s to remove", subID)
	}
	// exists, remove both indices
	delete(s.subsBySubID, subID)
	if _, ok = s.subsByChanID[sub.ChanID]; ok {
		delete(s.subsByChanID, sub.ChanID)
	}
	return nil
}

func (s *subscriptions) activate(subID string, chanID int64) error {
	s.lock.Lock()
	defer s.lock.Unlock()
	if sub, ok := s.subsBySubID[subID]; ok {
		sub.pending = false
		sub.ChanID = chanID
		s.subsByChanID[chanID] = sub
		sub.activate()
		return nil
	}
	return fmt.Errorf("could not find subscription ID %s to activate", subID)
}

func (s *subscriptions) lookupByChannelID(chanID int64) (*subscription, error) {
	s.lock.Lock()
	defer s.lock.Unlock()
	if sub, ok := s.subsByChanID[chanID]; ok {
		return sub, nil
	}
	return nil, fmt.Errorf("could not find subscription for channel ID %d", chanID)
}

func (s *subscriptions) lookupBySubscriptionID(subID string) (*subscription, error) {
	s.lock.Lock()
	defer s.lock.Unlock()
	if sub, ok := s.subsBySubID[subID]; ok {
		return sub, nil
	}
	return nil, fmt.Errorf("could not find subscription ID %s", subID)
}

func (s *subscriptions) heartbeat(chanID int64) error {
	s.lock.Lock()
	defer s.lock.Unlock()
	if sub, ok := s.subsByChanID[chanID]; ok {
		sub.heartbeat()
	}
	return fmt.Errorf("could not find channel ID to update heartbeat %d", chanID)
}
<|endoftext|>