<commit_before>package main

import "bufio"
import "bytes"
import "encoding/json"
import "fmt"
import "net/http"
import "os"
import "runtime"
import "time"

const BatchSize int = 2
const WorkerCount int = 2

var ESUrl string

// Represents a bulk header & document to be inserted into ES
type LogDocument struct {
	Header string
	Body   string
}

// Used to parse the LogDocument Body to build the header for it
type LineInfo struct {
	UUID string `json:"id"` // `json:"@uuid"` // _id: json['@uuid'],
	// Type      string `json:"@type"` // _type: json['@type'],
	// Timestamp string `json:"@timestamp"` // _timestamp: json['@timestamp']
}

// Used to create the header JSON string for LogDocument
type LogHeader struct {
	Index     string `json:"_index"`
	Id        string `json:"_id"`
	Type      string `json:"_type"`
	Timestamp string `json:"_timestamp"`
}

func (doc *LogDocument) String() string {
	var buf bytes.Buffer

	buf.WriteString(doc.Header)
	buf.WriteString("\n")
	buf.WriteString(doc.Body)
	buf.WriteString("\n")

	return buf.String()
}

// Takes in log lines from a channel, parses the required info out & creates
// LogDocument instances for the upsertWorkers to insert into ES
func lineWorker(lines <-chan string, documents chan LogDocument, done chan bool) {
	for line := range lines {
		// Unpack the line for the bits we need to build the header
		lineInfo := LineInfo{}
		err := json.Unmarshal([]byte(line), &lineInfo)
		if err != nil {
			fmt.Println(line)
			fmt.Println(lineInfo)
			panic(err)
		}

		// Build a header now we have the info
		logHeader := LogHeader{Id: lineInfo.UUID}
		header, _ := json.Marshal(logHeader)

		// Create a LogDocument and queue it up
		document := LogDocument{Header: string(header), Body: line}
		documents <- document
	}

	done <- true
}

func upsertWorker(documents <-chan LogDocument, docCounter chan int, done chan bool) {
	docs := []LogDocument{}

	for doc := range documents {
		docs = append(docs, doc)

		if len(docs) >= BatchSize {
			upsertToES(docs)
			docCounter <- len(docs)
			docs = []LogDocument{} // reset array
		}
	}

	// Handle any remaining documents as well
	if len(docs) > 0 {
		upsertToES(docs)
	}

	// docCounter <- len(docs)
	done <- true
}

func upsertToES(documents []LogDocument) {
	body := new(bytes.Buffer)
	for _, doc := range documents {
		_, err := body.WriteString(doc.String())
		if err != nil {
			fmt.Println("Error upserting to ES")
			panic(err)
		}
	}

	resp, err := http.Post(ESUrl, "application/json", body)
	if err != nil {
		panic(err)
	}

	defer resp.Body.Close()
	// fmt.Println(resp)
}

func outputStats(docCounter <-chan int, done chan bool) {
	totalCount := 0
	currentCount := 0

	startTime := time.Now()
	ticker := time.NewTicker(time.Second)

	for {
		select {
		case i := <-docCounter:
			totalCount += i
			currentCount += i

		case <-ticker.C:
			c := currentCount
			currentCount = 0
			fmt.Printf("%d/s uploaded\n", c)

		case <-done:
			ticker.Stop()

			// Little summary to finish
			duration := time.Since(startTime).Seconds()
			fmt.Printf("Upserted %d documents in %.2f seconds at %.0f/s\n", totalCount, duration, float64(totalCount)/duration)

			return
		}
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU() + 1)

	// TODO: handle cli arguments better
	argv := os.Args[1:]
	indexName := argv[0]

	ESUrl = fmt.Sprintf("http://localhost:8080/%s/_bulk", indexName)
	fmt.Println(ESUrl)

	// TODO: make this a proper ARGF
	argf, err := os.Open("/dev/stdin")
	if err != nil {
		fmt.Println("Error reading stdin")
		panic(err)
	}

	// Setup our channels for dishing out/waiting on work
	lines := make(chan string)
	documents := make(chan LogDocument, 3)
	docCounter := make(chan int)

	statsDone := make(chan bool)
	lineDone := make(chan bool)
	upsertDone := make(chan bool)

	// Stats outputter
	go outputStats(docCounter, statsDone)

	// Kick off our worker routines
	for i := 0; i < WorkerCount; i++ {
		go lineWorker(lines, documents, lineDone)
		go upsertWorker(documents, docCounter, upsertDone)
	}

	// Read the input & buffer into LogDocuments for workers
	scanner := bufio.NewScanner(argf)
	for scanner.Scan() {
		lines <- scanner.Text()
	}
	close(lines)

	// Once the line workers are finished, close off documents channel
	for i := 0; i < WorkerCount; i++ {
		<-lineDone
	}
	close(documents)

	// Wait for upsert workers to do their thing
	for i := 0; i < WorkerCount; i++ {
		<-upsertDone
	}

	close(docCounter)
	statsDone <- true
}
<commit_msg>Holyfuckitworks<commit_after>package main

import "bufio"
import "bytes"
import "io/ioutil"
import "encoding/json"
import "fmt"
import "net/http"
import "os"
import "runtime"
import "time"

const BatchSize int = 2
const WorkerCount int = 2

// const ESUrl string = "http://localhost:8080/_bulk"
const ESUrl string = "http://localhost:9200/_bulk"

var ESIndex string

// Represents a bulk header & document to be inserted into ES
type LogDocument struct {
	Header string
	Body   string
}

// Used to parse the LogDocument Body to build the header for it
type LineInfo struct {
	UUID      string `json:"@uuid"`
	Type      string `json:"@type"`
	Timestamp string `json:"@timestamp"`
}

// Used to create the header JSON string for LogDocument
type LogHeader struct {
	Index     string `json:"_index"`
	Id        string `json:"_id"`
	Type      string `json:"_type"`
	Timestamp string `json:"_timestamp"`
}

type IndexHeader struct {
	Index LogHeader `json:"index"`
}

func (doc *LogDocument) String() string {
	var buf bytes.Buffer

	buf.WriteString(doc.Header)
	buf.WriteString("\n")
	buf.WriteString(doc.Body)
	buf.WriteString("\n")

	return buf.String()
}

// Takes in log lines from a channel, parses the required info out & creates
// LogDocument instances for the upsertWorkers to insert into ES
func lineWorker(lines <-chan string, documents chan LogDocument, done chan bool) {
	for line := range lines {
		// Unpack the line for the bits we need to build the header
		lineInfo := LineInfo{}
		err := json.Unmarshal([]byte(line), &lineInfo)
		if err != nil {
			fmt.Println(line)
			fmt.Println(lineInfo)
			panic(err)
		}

		// Build a header now we have the info
		logHeader := LogHeader{
			Index:     ESIndex,
			Id:        lineInfo.UUID,
			Type:      lineInfo.Type,
			Timestamp: lineInfo.Timestamp,
		}
		header, _ := json.Marshal(IndexHeader{logHeader})

		// Create a LogDocument and queue it up
		document := LogDocument{Header: string(header), Body: line}
		documents <- document
	}

	done <- true
}

func upsertWorker(documents <-chan LogDocument, docCounter chan int, done chan bool) {
	docs := []LogDocument{}

	for doc := range documents {
		docs = append(docs, doc)

		if len(docs) >= BatchSize {
			upsertToES(docs)
			docCounter <- len(docs)
			docs = []LogDocument{} // reset array
		}
	}

	// Handle any remaining documents as well
	if len(docs) > 0 {
		upsertToES(docs)
	}

	// docCounter <- len(docs)
	done <- true
}

func upsertToES(documents []LogDocument) {
	body := new(bytes.Buffer)
	for _, doc := range documents {
		_, err := body.WriteString(doc.String())
		if err != nil {
			fmt.Println("Error upserting to ES")
			panic(err)
		}
	}

	resp, err := http.Post(ESUrl, "application/json", body)
	if err != nil {
		panic(err)
	}

	if resp.Status != "200" {
		b, _ := ioutil.ReadAll(resp.Body)
		fmt.Println(string(b))
	}

	defer resp.Body.Close()
}

func outputStats(docCounter <-chan int, done chan bool) {
	totalCount := 0
	currentCount := 0

	startTime := time.Now()
	ticker := time.NewTicker(time.Second)

	for {
		select {
		case i := <-docCounter:
			totalCount += i
			currentCount += i

		case <-ticker.C:
			c := currentCount
			currentCount = 0
			fmt.Printf("%d/s uploaded\n", c)

		case <-done:
			ticker.Stop()

			// Little summary to finish
			duration := time.Since(startTime).Seconds()
			fmt.Printf("Upserted %d documents in %.2f seconds at %.0f/s\n", totalCount, duration, float64(totalCount)/duration)

			return
		}
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU() + 1)

	// TODO: handle cli arguments better
	argv := os.Args[1:]
	ESIndex = argv[0]

	// TODO: make this a proper ARGF
	argf, err := os.Open("/dev/stdin")
	if err != nil {
		fmt.Println("Error reading stdin")
		panic(err)
	}

	// Setup our channels for dishing out/waiting on work
	lines := make(chan string)
	documents := make(chan LogDocument, 3)
	docCounter := make(chan int)

	statsDone := make(chan bool)
	lineDone := make(chan bool)
	upsertDone := make(chan bool)

	// Stats outputter
	go outputStats(docCounter, statsDone)

	// Kick off our worker routines
	for i := 0; i < WorkerCount; i++ {
		go lineWorker(lines, documents, lineDone)
		go upsertWorker(documents, docCounter, upsertDone)
	}

	// Read the input & buffer into LogDocuments for workers
	scanner := bufio.NewScanner(argf)
	for scanner.Scan() {
		lines <- scanner.Text()
	}
	close(lines)

	// Once the line workers are finished, close off documents channel
	for i := 0; i < WorkerCount; i++ {
		<-lineDone
	}
	close(documents)

	// Wait for upsert workers to do their thing
	for i := 0; i < WorkerCount; i++ {
		<-upsertDone
	}

	close(docCounter)
	statsDone <- true
}
<|endoftext|>


<commit_before><commit_msg>refactor(sql) simplify flags<commit_after><|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 integration provides functionality that needs to be shared between all
// integration tests.
//
// Integration tests are implemented through Go's test framework, as test
// functions that create and execute pipelines using the ptest package. Tests
// should be placed in smaller sub-packages for organizational purposes and
// parallelism (tests are only run in parallel across different packages).
// Integration tests should always begin with a call to CheckFilters to ensure
// test filters can be applied, and each package containing integration tests
// should call ptest.Main in a TestMain function if it uses ptest.
//
// Running integration tests can be done with a go test call with any flags that
// are required by the test pipelines, such as --runner or --endpoint.
// Example:
//    go test -v ./sdks/go/test/integration/... --runner=portable --endpoint=localhost:8099
//
// Alternatively, tests can be executed by running the
// run_validatesrunner_tests.sh script, which also performs much of the
// environment setup, or by calling gradle commands in :sdks:go:test.
package integration

import (
	"fmt"
	"regexp"
	"strings"
	"testing"

	// common runner flag.
	"github.com/apache/beam/sdks/v2/go/pkg/beam/options/jobopts"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
)

// Filters for temporarily skipping integration tests. All filters are regex
// matchers that must match the full name of a test at the point where
// CheckFilters is called. Multiple tests can be skipped by using regex
// wildcards. (ex. "TestXLang_.*" filters all tests starting with TestXLang_)
//
// It is strongly recommended to include, TODOs, GitHub issues, or just comments
// describing why tests are being skipped.

// sickbay filters tests that fail due to Go SDK errors. These tests will not
// execute on any runners.
var sickbay = []string{}

// Runner-specific test filters, for features that are not yet supported on
// specific runners.

var directFilters = []string{
	// The direct runner does not yet support cross-language.
	"TestXLang.*",
	"TestKafkaIO.*",
	"TestBigQueryIO.*",
	"TestDebeziumIO_BasicRead",
	"TestJDBCIO_BasicReadWrite",
	"TestJDBCIO_PostgresReadWrite",
	// Triggers, Panes are not yet supported
	"TestTrigger.*",
	"TestPanes",
	// The direct runner does not support the TestStream primitive
	"TestTestStream.*",
	// (https://github.com/apache/beam/issues/21130): The direct runner does not support windowed side inputs
	"TestValidateWindowedSideInputs",
	// (https://github.com/apache/beam/issues/21130): The direct runner does not currently support multimap side inputs
	"TestParDoMultiMapSideInput",
	"TestLargeWordcount_Loopback",
	// The direct runner does not support self-checkpointing
	"TestCheckpointing",
	// The direct runner does not support pipeline drain for SDF.
	"TestDrain",
	// FhirIO currently only supports Dataflow runner
	"TestFhirIO_.*",
}

var portableFilters = []string{
	// The portable runner does not support the TestStream primitive
	"TestTestStream.*",
	// The trigger and pane tests uses TestStream
	"TestTrigger.*",
	"TestPanes",
	// TODO(https://github.com/apache/beam/issues/21058): Python portable runner times out on Kafka reads.
	"TestKafkaIO.*",
	// TODO(BEAM-13215): GCP IOs currently do not work in non-Dataflow portable runners.
	"TestBigQueryIO.*",
	// The portable runner does not support self-checkpointing
	"TestCheckpointing",
	// The portable runner does not support pipeline drain for SDF.
	"TestDrain",
	// FhirIO currently only supports Dataflow runner
	"TestFhirIO_.*",
}

var flinkFilters = []string{
	// TODO(https://github.com/apache/beam/issues/20723): Flink tests timing out on reads.
	"TestXLang_Combine.*",
	// TODO(https://github.com/apache/beam/issues/21094): Test fails on post commits: "Insufficient number of network buffers".
	"TestXLang_Multi",
	"TestDebeziumIO_BasicRead",
	// TODO(BEAM-13215): GCP IOs currently do not work in non-Dataflow portable runners.
	"TestBigQueryIO.*",
	// The number of produced outputs in AfterSynchronizedProcessingTime varies in different runs.
	"TestTriggerAfterSynchronizedProcessingTime",
	// The flink runner does not support pipeline drain for SDF.
	"TestDrain",
	// FhirIO currently only supports Dataflow runner
	"TestFhirIO_.*",
}

var samzaFilters = []string{
	// TODO(https://github.com/apache/beam/issues/20987): Samza tests invalid encoding.
	"TestReshuffle",
	"TestReshuffleKV",
	// The Samza runner does not support the TestStream primitive
	"TestTestStream.*",
	// The trigger and pane tests uses TestStream
	"TestTrigger.*",
	"TestPanes",
	// TODO(https://github.com/apache/beam/issues/21244): Samza doesn't yet support post job metrics, used by WordCount
	"TestWordCount.*",
	// TODO(BEAM-13215): GCP IOs currently do not work in non-Dataflow portable runners.
	"TestBigQueryIO.*",
	// The Samza runner does not support self-checkpointing
	"TestCheckpointing",
	// The samza runner does not support pipeline drain for SDF.
	"TestDrain",
	// FhirIO currently only supports Dataflow runner
	"TestFhirIO_.*",
}

var sparkFilters = []string{
	// TODO(BEAM-11498): XLang tests broken with Spark runner.
	"TestXLang.*",
	"TestParDoSideInput",
	"TestParDoKVSideInput",
	// The Spark runner does not support the TestStream primitive
	"TestTestStream.*",
	// The trigger and pane tests uses TestStream
	"TestTrigger.*",
	"TestPanes",
	// [BEAM-13921]: Spark doesn't support side inputs to executable stages
	"TestDebeziumIO_BasicRead",
	// TODO(BEAM-13215): GCP IOs currently do not work in non-Dataflow portable runners.
	"TestBigQueryIO.*",
	// The spark runner does not support self-checkpointing
	"TestCheckpointing",
	// The spark runner does not support pipeline drain for SDF.
	"TestDrain",
	// FhirIO currently only supports Dataflow runner
	"TestFhirIO_.*",
}

var dataflowFilters = []string{
	// The Dataflow runner doesn't work with tests using testcontainers locally.
	"TestJDBCIO_BasicReadWrite",
	"TestJDBCIO_PostgresReadWrite",
	"TestDebeziumIO_BasicRead",
	// TODO(BEAM-11576): TestFlattenDup failing on this runner.
	"TestFlattenDup",
	// The Dataflow runner does not support the TestStream primitive
	"TestTestStream.*",
	// The trigger and pane tests uses TestStream
	"TestTrigger.*",
	"TestPanes",
	// There is no infrastructure for running KafkaIO tests with Dataflow.
	"TestKafkaIO.*",
	// Dataflow doesn't support any test that requires loopback.
	// Eg. For FileIO examples.
	".*Loopback.*",
	// Dataflow does not automatically terminate the TestCheckpointing pipeline when
	// complete.
	"TestCheckpointing",
	// TODO(21761): This test needs to provide GCP project to expansion service.
	"TestBigQueryIO_BasicWriteQueryRead",
	// Dataflow does not drain jobs by itself.
	"TestDrain",
}

// CheckFilters checks if an integration test is filtered to be skipped, either
// because the intended runner does not support it, or the test is sickbayed.
// This method should be called at the beginning of any integration test. If
// t.Run is used, CheckFilters should be called within the t.Run callback, so
// that sub-tests can be skipped individually.
func CheckFilters(t *testing.T) {
	if !ptest.MainCalled() {
		panic("ptest.Main() has not been called: please override TestMain to ensure that the integration test runs properly.")
	}

	// Check for sickbaying first.
	n := t.Name()
	for _, f := range sickbay {
		// Add start and end of string regexp matchers so only a full match is
		// counted.
		f = fmt.Sprintf("^%v$", f)
		match, err := regexp.MatchString(f, n)
		if err != nil {
			t.Errorf("Matching of regex '%v' with test '%v' failed: %v", f, n, err)
		}
		if match {
			t.Skipf("Test %v is currently sickbayed on all runners", n)
		}
	}
	// TODO(lostluck): Improve default job names.
	*jobopts.JobName = fmt.Sprintf("go-%v", strings.ToLower(n))

	// Test for runner-specific skipping second.
	var filters []string
	runner := *ptest.Runner
	if runner == "" {
		runner = ptest.DefaultRunner()
	}
	switch runner {
	case "direct", "DirectRunner":
		filters = directFilters
	case "portable", "PortableRunner":
		filters = portableFilters
	case "flink", "FlinkRunner":
		filters = flinkFilters
	case "samza", "SamzaRunner":
		filters = samzaFilters
	case "spark", "SparkRunner":
		filters = sparkFilters
	case "dataflow", "DataflowRunner":
		filters = dataflowFilters
	default:
		return
	}

	for _, f := range filters {
		// Add start and end of string regexp matchers so only a full match is
		// counted.
		f = fmt.Sprintf("^%v$", f)
		match, err := regexp.MatchString(f, n)
		if err != nil {
			t.Errorf("Matching of regex '%v' with test '%v' failed: %v", f, n, err)
		}
		if match {
			t.Skipf("Test %v is currently filtered for runner %v", n, runner)
		}
	}
}
<commit_msg>Add randomness to integration test job names to avoid collisions (#22065)<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 integration provides functionality that needs to be shared between all
// integration tests.
//
// Integration tests are implemented through Go's test framework, as test
// functions that create and execute pipelines using the ptest package. Tests
// should be placed in smaller sub-packages for organizational purposes and
// parallelism (tests are only run in parallel across different packages).
// Integration tests should always begin with a call to CheckFilters to ensure
// test filters can be applied, and each package containing integration tests
// should call ptest.Main in a TestMain function if it uses ptest.
//
// Running integration tests can be done with a go test call with any flags that
// are required by the test pipelines, such as --runner or --endpoint.
// Example:
//    go test -v ./sdks/go/test/integration/... --runner=portable --endpoint=localhost:8099
//
// Alternatively, tests can be executed by running the
// run_validatesrunner_tests.sh script, which also performs much of the
// environment setup, or by calling gradle commands in :sdks:go:test.
package integration

import (
	"fmt"
	"math/rand"
	"regexp"
	"strings"
	"testing"
	"time"

	// common runner flag.
	"github.com/apache/beam/sdks/v2/go/pkg/beam/options/jobopts"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
)

// Filters for temporarily skipping integration tests. All filters are regex
// matchers that must match the full name of a test at the point where
// CheckFilters is called. Multiple tests can be skipped by using regex
// wildcards. (ex. "TestXLang_.*" filters all tests starting with TestXLang_)
//
// It is strongly recommended to include, TODOs, GitHub issues, or just comments
// describing why tests are being skipped.

// sickbay filters tests that fail due to Go SDK errors. These tests will not
// execute on any runners.
var sickbay = []string{}

// Runner-specific test filters, for features that are not yet supported on
// specific runners.

var directFilters = []string{
	// The direct runner does not yet support cross-language.
	"TestXLang.*",
	"TestKafkaIO.*",
	"TestBigQueryIO.*",
	"TestDebeziumIO_BasicRead",
	"TestJDBCIO_BasicReadWrite",
	"TestJDBCIO_PostgresReadWrite",
	// Triggers, Panes are not yet supported
	"TestTrigger.*",
	"TestPanes",
	// The direct runner does not support the TestStream primitive
	"TestTestStream.*",
	// (https://github.com/apache/beam/issues/21130): The direct runner does not support windowed side inputs
	"TestValidateWindowedSideInputs",
	// (https://github.com/apache/beam/issues/21130): The direct runner does not currently support multimap side inputs
	"TestParDoMultiMapSideInput",
	"TestLargeWordcount_Loopback",
	// The direct runner does not support self-checkpointing
	"TestCheckpointing",
	// The direct runner does not support pipeline drain for SDF.
	"TestDrain",
	// FhirIO currently only supports Dataflow runner
	"TestFhirIO_.*",
}

var portableFilters = []string{
	// The portable runner does not support the TestStream primitive
	"TestTestStream.*",
	// The trigger and pane tests uses TestStream
	"TestTrigger.*",
	"TestPanes",
	// TODO(https://github.com/apache/beam/issues/21058): Python portable runner times out on Kafka reads.
	"TestKafkaIO.*",
	// TODO(BEAM-13215): GCP IOs currently do not work in non-Dataflow portable runners.
	"TestBigQueryIO.*",
	// The portable runner does not support self-checkpointing
	"TestCheckpointing",
	// The portable runner does not support pipeline drain for SDF.
	"TestDrain",
	// FhirIO currently only supports Dataflow runner
	"TestFhirIO_.*",
}

var flinkFilters = []string{
	// TODO(https://github.com/apache/beam/issues/20723): Flink tests timing out on reads.
	"TestXLang_Combine.*",
	// TODO(https://github.com/apache/beam/issues/21094): Test fails on post commits: "Insufficient number of network buffers".
	"TestXLang_Multi",
	"TestDebeziumIO_BasicRead",
	// TODO(BEAM-13215): GCP IOs currently do not work in non-Dataflow portable runners.
	"TestBigQueryIO.*",
	// The number of produced outputs in AfterSynchronizedProcessingTime varies in different runs.
	"TestTriggerAfterSynchronizedProcessingTime",
	// The flink runner does not support pipeline drain for SDF.
	"TestDrain",
	// FhirIO currently only supports Dataflow runner
	"TestFhirIO_.*",
}

var samzaFilters = []string{
	// TODO(https://github.com/apache/beam/issues/20987): Samza tests invalid encoding.
	"TestReshuffle",
	"TestReshuffleKV",
	// The Samza runner does not support the TestStream primitive
	"TestTestStream.*",
	// The trigger and pane tests uses TestStream
	"TestTrigger.*",
	"TestPanes",
	// TODO(https://github.com/apache/beam/issues/21244): Samza doesn't yet support post job metrics, used by WordCount
	"TestWordCount.*",
	// TODO(BEAM-13215): GCP IOs currently do not work in non-Dataflow portable runners.
	"TestBigQueryIO.*",
	// The Samza runner does not support self-checkpointing
	"TestCheckpointing",
	// The samza runner does not support pipeline drain for SDF.
	"TestDrain",
	// FhirIO currently only supports Dataflow runner
	"TestFhirIO_.*",
}

var sparkFilters = []string{
	// TODO(BEAM-11498): XLang tests broken with Spark runner.
	"TestXLang.*",
	"TestParDoSideInput",
	"TestParDoKVSideInput",
	// The Spark runner does not support the TestStream primitive
	"TestTestStream.*",
	// The trigger and pane tests uses TestStream
	"TestTrigger.*",
	"TestPanes",
	// [BEAM-13921]: Spark doesn't support side inputs to executable stages
	"TestDebeziumIO_BasicRead",
	// TODO(BEAM-13215): GCP IOs currently do not work in non-Dataflow portable runners.
	"TestBigQueryIO.*",
	// The spark runner does not support self-checkpointing
	"TestCheckpointing",
	// The spark runner does not support pipeline drain for SDF.
	"TestDrain",
	// FhirIO currently only supports Dataflow runner
	"TestFhirIO_.*",
}

var dataflowFilters = []string{
	// The Dataflow runner doesn't work with tests using testcontainers locally.
	"TestJDBCIO_BasicReadWrite",
	"TestJDBCIO_PostgresReadWrite",
	"TestDebeziumIO_BasicRead",
	// TODO(BEAM-11576): TestFlattenDup failing on this runner.
	"TestFlattenDup",
	// The Dataflow runner does not support the TestStream primitive
	"TestTestStream.*",
	// The trigger and pane tests uses TestStream
	"TestTrigger.*",
	"TestPanes",
	// There is no infrastructure for running KafkaIO tests with Dataflow.
	"TestKafkaIO.*",
	// Dataflow doesn't support any test that requires loopback.
	// Eg. For FileIO examples.
	".*Loopback.*",
	// Dataflow does not automatically terminate the TestCheckpointing pipeline when
	// complete.
	"TestCheckpointing",
	// TODO(21761): This test needs to provide GCP project to expansion service.
	"TestBigQueryIO_BasicWriteQueryRead",
	// Dataflow does not drain jobs by itself.
	"TestDrain",
}

// CheckFilters checks if an integration test is filtered to be skipped, either
// because the intended runner does not support it, or the test is sickbayed.
// This method should be called at the beginning of any integration test. If
// t.Run is used, CheckFilters should be called within the t.Run callback, so
// that sub-tests can be skipped individually.
func CheckFilters(t *testing.T) {
	if !ptest.MainCalled() {
		panic("ptest.Main() has not been called: please override TestMain to ensure that the integration test runs properly.")
	}

	// Check for sickbaying first.
	n := t.Name()
	for _, f := range sickbay {
		// Add start and end of string regexp matchers so only a full match is
		// counted.
		f = fmt.Sprintf("^%v$", f)
		match, err := regexp.MatchString(f, n)
		if err != nil {
			t.Errorf("Matching of regex '%v' with test '%v' failed: %v", f, n, err)
		}
		if match {
			t.Skipf("Test %v is currently sickbayed on all runners", n)
		}
	}
	s1 := rand.NewSource(time.Now().UnixNano())
	r1 := rand.New(s1)
	*jobopts.JobName = fmt.Sprintf("go-%v-%v", strings.ToLower(n), r1.Intn(1000))

	// Test for runner-specific skipping second.
	var filters []string
	runner := *ptest.Runner
	if runner == "" {
		runner = ptest.DefaultRunner()
	}
	switch runner {
	case "direct", "DirectRunner":
		filters = directFilters
	case "portable", "PortableRunner":
		filters = portableFilters
	case "flink", "FlinkRunner":
		filters = flinkFilters
	case "samza", "SamzaRunner":
		filters = samzaFilters
	case "spark", "SparkRunner":
		filters = sparkFilters
	case "dataflow", "DataflowRunner":
		filters = dataflowFilters
	default:
		return
	}

	for _, f := range filters {
		// Add start and end of string regexp matchers so only a full match is
		// counted.
		f = fmt.Sprintf("^%v$", f)
		match, err := regexp.MatchString(f, n)
		if err != nil {
			t.Errorf("Matching of regex '%v' with test '%v' failed: %v", f, n, err)
		}
		if match {
			t.Skipf("Test %v is currently filtered for runner %v", n, runner)
		}
	}
}
<|endoftext|>


<commit_before>// Discordgo - Discord bindings for Go
// Available at https://github.com/bwmarrin/discordgo

// Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains code related to Discord voice suppport

package discordgo

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"net"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// ------------------------------------------------------------------------------------------------
// Code related to both Voice Websocket and UDP connections.
// ------------------------------------------------------------------------------------------------

// A Voice struct holds all data and functions related to Discord Voice support.
type Voice struct {
	sync.Mutex             // future use
	Ready      bool        // If true, voice is ready to send/receive audio
	Debug      bool        // If true, print extra logging
	OP2        *voiceOP2   // exported for dgvoice, may change.
	Opus       chan []byte // Chan for sending opus audio
	//	FrameRate  int         // This can be used to set the FrameRate of Opus data
	//	FrameSize  int         // This can be used to set the FrameSize of Opus data

	wsConn  *websocket.Conn
	UDPConn *net.UDPConn // this will become unexported soon.

	sessionID string
	token     string
	endpoint  string
	guildID   string
	channelID string
	userID    string
}

// ------------------------------------------------------------------------------------------------
// Code related to the Voice websocket connection
// ------------------------------------------------------------------------------------------------

// A voiceOP2 stores the data for the voice operation 2 websocket event
// which is sort of like the voice READY packet
type voiceOP2 struct {
	SSRC              uint32        `json:"ssrc"`
	Port              int           `json:"port"`
	Modes             []string      `json:"modes"`
	HeartbeatInterval time.Duration `json:"heartbeat_interval"`
}

type voiceHandshakeData struct {
	ServerID  string `json:"server_id"`
	UserID    string `json:"user_id"`
	SessionID string `json:"session_id"`
	Token     string `json:"token"`
}

type voiceHandshakeOp struct {
	Op   int                `json:"op"` // Always 0
	Data voiceHandshakeData `json:"d"`
}

// Open opens a voice connection.  This should be called
// after VoiceChannelJoin is used and the data VOICE websocket events
// are captured.
func (v *Voice) Open() (err error) {

	// TODO: How do we handle changing channels?

	// Don't open a websocket if one is already open
	if v.wsConn != nil {
		return
	}

	// Connect to Voice Websocket
	vg := fmt.Sprintf("wss://%s", strings.TrimSuffix(v.endpoint, ":80"))
	v.wsConn, _, err = websocket.DefaultDialer.Dial(vg, nil)
	if err != nil {
		fmt.Println("VOICE error opening websocket:", err)
		return
	}

	data := voiceHandshakeOp{0, voiceHandshakeData{v.guildID, v.userID, v.sessionID, v.token}}

	err = v.wsConn.WriteJSON(data)
	if err != nil {
		fmt.Println("VOICE error sending init packet:", err)
		return
	}

	// Start a listening for voice websocket events
	// TODO add a check here to make sure Listen worked by monitoring
	// a chan or bool?
	go v.wsListen()

	return
}

// Close closes the voice connection
func (v *Voice) Close() {

	if v.UDPConn != nil {
		v.UDPConn.Close()
	}

	if v.wsConn != nil {
		v.wsConn.Close()
	}
}

// wsListen listens on the voice websocket for messages and passes them
// to the voice event handler.  This is automaticly called by the Open func
func (v *Voice) wsListen() {

	for {
		messageType, message, err := v.wsConn.ReadMessage()
		if err != nil {
			// TODO: Handle this problem better.
			// TODO: needs proper logging
			fmt.Println("Voice Listen Error:", err)
			break
		}

		// Pass received message to voice event handler
		go v.wsEvent(messageType, message)
	}

	return
}

// wsEvent handles any voice websocket events. This is only called by the
// wsListen() function.
func (v *Voice) wsEvent(messageType int, message []byte) {

	if v.Debug {
		fmt.Println("wsEvent received: ", messageType)
		printJSON(message)
	}

	var e Event
	if err := json.Unmarshal(message, &e); err != nil {
		fmt.Println("wsEvent Unmarshall error: ", err)
		return
	}

	switch e.Operation {

	case 2: // READY

		v.OP2 = &voiceOP2{}
		if err := json.Unmarshal(e.RawData, v.OP2); err != nil {
			fmt.Println("voiceWS.onEvent OP2 Unmarshall error: ", err)
			printJSON(e.RawData) // TODO: Better error logging
			return
		}

		// Start the voice websocket heartbeat to keep the connection alive
		go v.wsHeartbeat(v.OP2.HeartbeatInterval)
		// TODO monitor a chan/bool to verify this was successful

		// Start the UDP connection
		err := v.udpOpen()
		if err != nil {
			fmt.Println("Error opening udp connection: ", err)
			return
		}

		// Start the opusSender.
		// TODO: Should we allow 48000/960 values to be user defined?
		v.Opus = make(chan []byte, 2)
		go v.opusSender(v.Opus, 48000, 960)

		return

	case 3: // HEARTBEAT response
		// add code to use this to track latency?
		return

	case 4:
		// TODO

	case 5:
		// SPEAKING TRUE/FALSE NOTIFICATION
		/*
			{
				"user_id": "1238921738912",
				"ssrc": 2,
				"speaking": false
			}
		*/

	default:
		fmt.Println("UNKNOWN VOICE OP: ", e.Operation)
		printJSON(e.RawData)
	}

	return
}

type voiceHeartbeatOp struct {
	Op   int `json:"op"` // Always 3
	Data int `json:"d"`
}

// wsHeartbeat sends regular heartbeats to voice Discord so it knows the client
// is still connected.  If you do not send these heartbeats Discord will
// disconnect the websocket connection after a few seconds.
func (v *Voice) wsHeartbeat(i time.Duration) {

	ticker := time.NewTicker(i * time.Millisecond)
	for {
		err := v.wsConn.WriteJSON(voiceHeartbeatOp{3, int(time.Now().Unix())})
		if err != nil {
			v.Ready = false
			fmt.Println("wsHeartbeat send error: ", err)
			return // TODO better logging
		}
		<-ticker.C
	}
}

type voiceSpeakingData struct {
	Speaking bool `json:"speaking"`
	Delay    int  `json:"delay"`
}

type voiceSpeakingOp struct {
	Op   int               `json:"op"` // Always 5
	Data voiceSpeakingData `json:"d"`
}

// Speaking sends a speaking notification to Discord over the voice websocket.
// This must be sent as true prior to sending audio and should be set to false
// once finished sending audio.
//  b  : Send true if speaking, false if not.
func (v *Voice) Speaking(b bool) (err error) {

	if v.wsConn == nil {
		return fmt.Errorf("No Voice websocket.")
	}

	data := voiceSpeakingOp{5, voiceSpeakingData{b, 0}}
	err = v.wsConn.WriteJSON(data)
	if err != nil {
		fmt.Println("Speaking() write json error:", err)
		return
	}

	return
}

// ------------------------------------------------------------------------------------------------
// Code related to the Voice UDP connection
// ------------------------------------------------------------------------------------------------

type voiceUDPData struct {
	Address string `json:"address"` // Public IP of machine running this code
	Port    uint16 `json:"port"`    // UDP Port of machine running this code
	Mode    string `json:"mode"`    // plain or ?  (plain or encrypted)
}

type voiceUDPD struct {
	Protocol string       `json:"protocol"` // Always "udp" ?
	Data     voiceUDPData `json:"data"`
}

type voiceUDPOp struct {
	Op   int       `json:"op"` // Always 1
	Data voiceUDPD `json:"d"`
}

// udpOpen opens a UDP connection to the voice server and completes the
// initial required handshake.  This connection is left open in the session
// and can be used to send or receive audio.  This should only be called
// from voice.wsEvent OP2
func (v *Voice) udpOpen() (err error) {

	host := fmt.Sprintf("%s:%d", strings.TrimSuffix(v.endpoint, ":80"), v.OP2.Port)
	addr, err := net.ResolveUDPAddr("udp", host)
	if err != nil {
		fmt.Println("udpOpen resolve addr error: ", err)
		// TODO better logging
		return
	}

	v.UDPConn, err = net.DialUDP("udp", nil, addr)
	if err != nil {
		fmt.Println("udpOpen dial udp error: ", err)
		// TODO better logging
		return
	}

	// Create a 70 byte array and put the SSRC code from the Op 2 Voice event
	// into it.  Then send that over the UDP connection to Discord
	sb := make([]byte, 70)
	binary.BigEndian.PutUint32(sb, v.OP2.SSRC)
	_, err = v.UDPConn.Write(sb)
	if err != nil {
		fmt.Println("udpOpen udp write error : ", err)
		// TODO better logging
		return
	}

	// Create a 70 byte array and listen for the initial handshake response
	// from Discord.  Once we get it parse the IP and PORT information out
	// of the response.  This should be our public IP and PORT as Discord
	// saw us.
	rb := make([]byte, 70)
	rlen, _, err := v.UDPConn.ReadFromUDP(rb)
	if err != nil {
		fmt.Println("udpOpen udp read error : ", err)
		// TODO better logging
		return
	}
	if rlen < 70 {
		fmt.Println("Voice RLEN should be 70 but isn't")
	}

	// Loop over position 4 though 20 to grab the IP address
	// Should never be beyond position 20.
	var ip string
	for i := 4; i < 20; i++ {
		if rb[i] == 0 {
			break
		}
		ip += string(rb[i])
	}

	// Grab port from postion 68 and 69
	port := binary.LittleEndian.Uint16(rb[68:70])

	// Take the data from above and send it back to Discord to finalize
	// the UDP connection handshake.
	data := voiceUDPOp{1, voiceUDPD{"udp", voiceUDPData{ip, port, "plain"}}}

	err = v.wsConn.WriteJSON(data)
	if err != nil {
		fmt.Println("udpOpen write json error:", err)
		return
	}

	// start udpKeepAlive
	go v.udpKeepAlive(5 * time.Second)
	// TODO: find a way to check that it fired off okay

	v.Ready = true
	return
}

// udpKeepAlive sends a udp packet to keep the udp connection open
// This is still a bit of a "proof of concept"
func (v *Voice) udpKeepAlive(i time.Duration) {

	var err error
	var sequence uint64 = 0

	packet := make([]byte, 8)

	ticker := time.NewTicker(i)
	for {
		// TODO: Add a way to break from loop

		binary.LittleEndian.PutUint64(packet, sequence)
		sequence++

		_, err = v.UDPConn.Write(packet)
		if err != nil {
			fmt.Println("udpKeepAlive udp write error : ", err)
			return
		}
		<-ticker.C
	}
}

// opusSender will listen on the given channel and send any
// pre-encoded opus audio to Discord.  Supposedly.
func (v *Voice) opusSender(opus <-chan []byte, rate, size int) {

	// TODO: Better checking to prevent this from running more than
	// one instance at a time.
	v.Lock()
	if opus == nil {
		v.Unlock()
		return
	}
	v.Unlock()

	runtime.LockOSThread()

	var sequence uint16 = 0
	var timestamp uint32 = 0
	udpHeader := make([]byte, 12)

	// build the parts that don't change in the udpHeader
	udpHeader[0] = 0x80
	udpHeader[1] = 0x78
	binary.BigEndian.PutUint32(udpHeader[8:], v.OP2.SSRC)

	// start a send loop that loops until buf chan is closed
	ticker := time.NewTicker(time.Millisecond * time.Duration(size/(rate/1000)))
	for {

		// Add sequence and timestamp to udpPacket
		binary.BigEndian.PutUint16(udpHeader[2:], sequence)
		binary.BigEndian.PutUint32(udpHeader[4:], timestamp)

		// Get data from chan.  If chan is closed, return.
		recvbuf, ok := <-opus
		if !ok {
			return
		}

		// Combine the UDP Header and the opus data
		sendbuf := append(udpHeader, recvbuf...)

		// block here until we're exactly at the right time :)
		// Then send rtp audio packet to Discord over UDP
		<-ticker.C
		v.UDPConn.Write(sendbuf)

		if (sequence) == 0xFFFF {
			sequence = 0
		} else {
			sequence += 1
		}

		if (timestamp + uint32(size)) >= 0xFFFFFFFF {
			timestamp = 0
		} else {
			timestamp += uint32(size)
		}
	}
}
<commit_msg>Moved Voice.Ready = true into opusSender<commit_after>// Discordgo - Discord bindings for Go
// Available at https://github.com/bwmarrin/discordgo

// Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains code related to Discord voice suppport

package discordgo

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"net"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// ------------------------------------------------------------------------------------------------
// Code related to both Voice Websocket and UDP connections.
// ------------------------------------------------------------------------------------------------

// A Voice struct holds all data and functions related to Discord Voice support.
type Voice struct {
	sync.Mutex             // future use
	Ready      bool        // If true, voice is ready to send/receive audio
	Debug      bool        // If true, print extra logging
	OP2        *voiceOP2   // exported for dgvoice, may change.
	Opus       chan []byte // Chan for sending opus audio
	//	FrameRate  int         // This can be used to set the FrameRate of Opus data
	//	FrameSize  int         // This can be used to set the FrameSize of Opus data

	wsConn  *websocket.Conn
	UDPConn *net.UDPConn // this will become unexported soon.

	sessionID string
	token     string
	endpoint  string
	guildID   string
	channelID string
	userID    string
}

// ------------------------------------------------------------------------------------------------
// Code related to the Voice websocket connection
// ------------------------------------------------------------------------------------------------

// A voiceOP2 stores the data for the voice operation 2 websocket event
// which is sort of like the voice READY packet
type voiceOP2 struct {
	SSRC              uint32        `json:"ssrc"`
	Port              int           `json:"port"`
	Modes             []string      `json:"modes"`
	HeartbeatInterval time.Duration `json:"heartbeat_interval"`
}

type voiceHandshakeData struct {
	ServerID  string `json:"server_id"`
	UserID    string `json:"user_id"`
	SessionID string `json:"session_id"`
	Token     string `json:"token"`
}

type voiceHandshakeOp struct {
	Op   int                `json:"op"` // Always 0
	Data voiceHandshakeData `json:"d"`
}

// Open opens a voice connection.  This should be called
// after VoiceChannelJoin is used and the data VOICE websocket events
// are captured.
func (v *Voice) Open() (err error) {

	// TODO: How do we handle changing channels?

	// Don't open a websocket if one is already open
	if v.wsConn != nil {
		return
	}

	// Connect to Voice Websocket
	vg := fmt.Sprintf("wss://%s", strings.TrimSuffix(v.endpoint, ":80"))
	v.wsConn, _, err = websocket.DefaultDialer.Dial(vg, nil)
	if err != nil {
		fmt.Println("VOICE error opening websocket:", err)
		return
	}

	data := voiceHandshakeOp{0, voiceHandshakeData{v.guildID, v.userID, v.sessionID, v.token}}

	err = v.wsConn.WriteJSON(data)
	if err != nil {
		fmt.Println("VOICE error sending init packet:", err)
		return
	}

	// Start a listening for voice websocket events
	// TODO add a check here to make sure Listen worked by monitoring
	// a chan or bool?
	go v.wsListen()

	return
}

// Close closes the voice connection
func (v *Voice) Close() {

	if v.UDPConn != nil {
		v.UDPConn.Close()
	}

	if v.wsConn != nil {
		v.wsConn.Close()
	}
}

// wsListen listens on the voice websocket for messages and passes them
// to the voice event handler.  This is automaticly called by the Open func
func (v *Voice) wsListen() {

	for {
		messageType, message, err := v.wsConn.ReadMessage()
		if err != nil {
			// TODO: Handle this problem better.
			// TODO: needs proper logging
			fmt.Println("Voice Listen Error:", err)
			break
		}

		// Pass received message to voice event handler
		go v.wsEvent(messageType, message)
	}

	return
}

// wsEvent handles any voice websocket events. This is only called by the
// wsListen() function.
func (v *Voice) wsEvent(messageType int, message []byte) {

	if v.Debug {
		fmt.Println("wsEvent received: ", messageType)
		printJSON(message)
	}

	var e Event
	if err := json.Unmarshal(message, &e); err != nil {
		fmt.Println("wsEvent Unmarshall error: ", err)
		return
	}

	switch e.Operation {

	case 2: // READY

		v.OP2 = &voiceOP2{}
		if err := json.Unmarshal(e.RawData, v.OP2); err != nil {
			fmt.Println("voiceWS.onEvent OP2 Unmarshall error: ", err)
			printJSON(e.RawData) // TODO: Better error logging
			return
		}

		// Start the voice websocket heartbeat to keep the connection alive
		go v.wsHeartbeat(v.OP2.HeartbeatInterval)
		// TODO monitor a chan/bool to verify this was successful

		// Start the UDP connection
		err := v.udpOpen()
		if err != nil {
			fmt.Println("Error opening udp connection: ", err)
			return
		}

		// Start the opusSender.
		// TODO: Should we allow 48000/960 values to be user defined?
		v.Opus = make(chan []byte, 2)
		go v.opusSender(v.Opus, 48000, 960)

		return

	case 3: // HEARTBEAT response
		// add code to use this to track latency?
		return

	case 4:
		// TODO

	case 5:
		// SPEAKING TRUE/FALSE NOTIFICATION
		/*
			{
				"user_id": "1238921738912",
				"ssrc": 2,
				"speaking": false
			}
		*/

	default:
		fmt.Println("UNKNOWN VOICE OP: ", e.Operation)
		printJSON(e.RawData)
	}

	return
}

type voiceHeartbeatOp struct {
	Op   int `json:"op"` // Always 3
	Data int `json:"d"`
}

// wsHeartbeat sends regular heartbeats to voice Discord so it knows the client
// is still connected.  If you do not send these heartbeats Discord will
// disconnect the websocket connection after a few seconds.
func (v *Voice) wsHeartbeat(i time.Duration) {

	ticker := time.NewTicker(i * time.Millisecond)
	for {
		err := v.wsConn.WriteJSON(voiceHeartbeatOp{3, int(time.Now().Unix())})
		if err != nil {
			v.Ready = false
			fmt.Println("wsHeartbeat send error: ", err)
			return // TODO better logging
		}
		<-ticker.C
	}
}

type voiceSpeakingData struct {
	Speaking bool `json:"speaking"`
	Delay    int  `json:"delay"`
}

type voiceSpeakingOp struct {
	Op   int               `json:"op"` // Always 5
	Data voiceSpeakingData `json:"d"`
}

// Speaking sends a speaking notification to Discord over the voice websocket.
// This must be sent as true prior to sending audio and should be set to false
// once finished sending audio.
//  b  : Send true if speaking, false if not.
func (v *Voice) Speaking(b bool) (err error) {

	if v.wsConn == nil {
		return fmt.Errorf("No Voice websocket.")
	}

	data := voiceSpeakingOp{5, voiceSpeakingData{b, 0}}
	err = v.wsConn.WriteJSON(data)
	if err != nil {
		fmt.Println("Speaking() write json error:", err)
		return
	}

	return
}

// ------------------------------------------------------------------------------------------------
// Code related to the Voice UDP connection
// ------------------------------------------------------------------------------------------------

type voiceUDPData struct {
	Address string `json:"address"` // Public IP of machine running this code
	Port    uint16 `json:"port"`    // UDP Port of machine running this code
	Mode    string `json:"mode"`    // plain or ?  (plain or encrypted)
}

type voiceUDPD struct {
	Protocol string       `json:"protocol"` // Always "udp" ?
	Data     voiceUDPData `json:"data"`
}

type voiceUDPOp struct {
	Op   int       `json:"op"` // Always 1
	Data voiceUDPD `json:"d"`
}

// udpOpen opens a UDP connection to the voice server and completes the
// initial required handshake.  This connection is left open in the session
// and can be used to send or receive audio.  This should only be called
// from voice.wsEvent OP2
func (v *Voice) udpOpen() (err error) {

	host := fmt.Sprintf("%s:%d", strings.TrimSuffix(v.endpoint, ":80"), v.OP2.Port)
	addr, err := net.ResolveUDPAddr("udp", host)
	if err != nil {
		fmt.Println("udpOpen resolve addr error: ", err)
		// TODO better logging
		return
	}

	v.UDPConn, err = net.DialUDP("udp", nil, addr)
	if err != nil {
		fmt.Println("udpOpen dial udp error: ", err)
		// TODO better logging
		return
	}

	// Create a 70 byte array and put the SSRC code from the Op 2 Voice event
	// into it.  Then send that over the UDP connection to Discord
	sb := make([]byte, 70)
	binary.BigEndian.PutUint32(sb, v.OP2.SSRC)
	_, err = v.UDPConn.Write(sb)
	if err != nil {
		fmt.Println("udpOpen udp write error : ", err)
		// TODO better logging
		return
	}

	// Create a 70 byte array and listen for the initial handshake response
	// from Discord.  Once we get it parse the IP and PORT information out
	// of the response.  This should be our public IP and PORT as Discord
	// saw us.
	rb := make([]byte, 70)
	rlen, _, err := v.UDPConn.ReadFromUDP(rb)
	if err != nil {
		fmt.Println("udpOpen udp read error : ", err)
		// TODO better logging
		return
	}
	if rlen < 70 {
		fmt.Println("Voice RLEN should be 70 but isn't")
	}

	// Loop over position 4 though 20 to grab the IP address
	// Should never be beyond position 20.
	var ip string
	for i := 4; i < 20; i++ {
		if rb[i] == 0 {
			break
		}
		ip += string(rb[i])
	}

	// Grab port from postion 68 and 69
	port := binary.LittleEndian.Uint16(rb[68:70])

	// Take the data from above and send it back to Discord to finalize
	// the UDP connection handshake.
	data := voiceUDPOp{1, voiceUDPD{"udp", voiceUDPData{ip, port, "plain"}}}

	err = v.wsConn.WriteJSON(data)
	if err != nil {
		fmt.Println("udpOpen write json error:", err)
		return
	}

	// start udpKeepAlive
	go v.udpKeepAlive(5 * time.Second)
	// TODO: find a way to check that it fired off okay

	return
}

// udpKeepAlive sends a udp packet to keep the udp connection open
// This is still a bit of a "proof of concept"
func (v *Voice) udpKeepAlive(i time.Duration) {

	var err error
	var sequence uint64 = 0

	packet := make([]byte, 8)

	ticker := time.NewTicker(i)
	for {
		// TODO: Add a way to break from loop

		binary.LittleEndian.PutUint64(packet, sequence)
		sequence++

		_, err = v.UDPConn.Write(packet)
		if err != nil {
			fmt.Println("udpKeepAlive udp write error : ", err)
			return
		}
		<-ticker.C
	}
}

// opusSender will listen on the given channel and send any
// pre-encoded opus audio to Discord.  Supposedly.
func (v *Voice) opusSender(opus <-chan []byte, rate, size int) {

	// TODO: Better checking to prevent this from running more than
	// one instance at a time.
	v.Lock()
	if opus == nil {
		v.Unlock()
		return
	}
	v.Unlock()

	runtime.LockOSThread()

	// Voice is now ready to receive audio packets
	// TODO: this needs reviewed as I think there must be a better way.
	v.Ready = true

	var sequence uint16 = 0
	var timestamp uint32 = 0
	udpHeader := make([]byte, 12)

	// build the parts that don't change in the udpHeader
	udpHeader[0] = 0x80
	udpHeader[1] = 0x78
	binary.BigEndian.PutUint32(udpHeader[8:], v.OP2.SSRC)

	// start a send loop that loops until buf chan is closed
	ticker := time.NewTicker(time.Millisecond * time.Duration(size/(rate/1000)))
	for {

		// Add sequence and timestamp to udpPacket
		binary.BigEndian.PutUint16(udpHeader[2:], sequence)
		binary.BigEndian.PutUint32(udpHeader[4:], timestamp)

		// Get data from chan.  If chan is closed, return.
		recvbuf, ok := <-opus
		if !ok {
			return
		}

		// Combine the UDP Header and the opus data
		sendbuf := append(udpHeader, recvbuf...)

		// block here until we're exactly at the right time :)
		// Then send rtp audio packet to Discord over UDP
		<-ticker.C
		v.UDPConn.Write(sendbuf)

		if (sequence) == 0xFFFF {
			sequence = 0
		} else {
			sequence += 1
		}

		if (timestamp + uint32(size)) >= 0xFFFFFFFF {
			timestamp = 0
		} else {
			timestamp += uint32(size)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/crosbymichael/skydock/docker"
	"log"
)

func (app *Application) eventHandler(c chan *docker.Event) {
	defer app.Status.Done()

	for e := range c {
		log.Printf("received (%s) %s %s", e.Status, e.ContainerId, e.Image)

		cont, err := app.Docker.FetchContainer(e.ContainerId, e.Image)
		if err != nil {
			log.Println(e.ContainerId, err)
			continue
		}

		switch e.Status {
		case "die", "stop", "kill":
			app.Remove(cont)
		case "start", "restart":
			app.Add(cont)
		}
	}
}

func (app *Application) Add(c *docker.Container) {
	cId := ContainerID(c.Id)
	ip := IPAddress(c.NetworkSettings.IpAddress)
	u := Upstream("http://" + ip + ":80")
	app.IPs[cId] = ip

	for _, h := range getHostnames(c) {
		exists, err := h.Exists(app.Redis)
		if err != nil {
			log.Println(err)
		}

		if !exists {
			if err := h.Create(app.Redis); err != nil {
				log.Println(err)
				continue
			}
		}

		if err := u.Register(app.Redis, h); err != nil {
			log.Println(err)
		}
	}
}

func (app *Application) Remove(c *docker.Container) {
	cId := ContainerID(c.Id)
	ip, ok := app.IPs[cId]
	if !ok {
		return
	}
	delete(app.IPs, cId)
	u := Upstream("http://" + ip + ":80")

	for _, h := range getHostnames(c) {
		exists, err := h.Exists(app.Redis)
		if err != nil {
			log.Println(err)
		}

		if !exists {
			continue
		}

		if err := u.Unregister(app.Redis, h); err != nil {
			log.Println(err)
		}
	}
}

// TODO
// Watch for stop/start events on containers, removing/adding them as needed
func (app *Application) watch() {
	e := app.Docker.GetEvents()

	app.Status.Add(1)
	go app.eventHandler(e)

	log.Printf("Starting main process")
	app.Status.Wait()
	log.Printf("Stopping cleanly via EOF")
}
<commit_msg>Remove completed todo<commit_after>package main

import (
	"github.com/crosbymichael/skydock/docker"
	"log"
)

func (app *Application) eventHandler(c chan *docker.Event) {
	defer app.Status.Done()

	for e := range c {
		log.Printf("received (%s) %s %s", e.Status, e.ContainerId, e.Image)

		cont, err := app.Docker.FetchContainer(e.ContainerId, e.Image)
		if err != nil {
			log.Println(e.ContainerId, err)
			continue
		}

		switch e.Status {
		case "die", "stop", "kill":
			app.Remove(cont)
		case "start", "restart":
			app.Add(cont)
		}
	}
}

func (app *Application) Add(c *docker.Container) {
	cId := ContainerID(c.Id)
	ip := IPAddress(c.NetworkSettings.IpAddress)
	u := Upstream("http://" + ip + ":80")
	app.IPs[cId] = ip

	for _, h := range getHostnames(c) {
		exists, err := h.Exists(app.Redis)
		if err != nil {
			log.Println(err)
		}

		if !exists {
			if err := h.Create(app.Redis); err != nil {
				log.Println(err)
				continue
			}
		}

		if err := u.Register(app.Redis, h); err != nil {
			log.Println(err)
		}
	}
}

func (app *Application) Remove(c *docker.Container) {
	cId := ContainerID(c.Id)
	ip, ok := app.IPs[cId]
	if !ok {
		return
	}
	delete(app.IPs, cId)
	u := Upstream("http://" + ip + ":80")

	for _, h := range getHostnames(c) {
		exists, err := h.Exists(app.Redis)
		if err != nil {
			log.Println(err)
		}

		if !exists {
			continue
		}

		if err := u.Unregister(app.Redis, h); err != nil {
			log.Println(err)
		}
	}
}

func (app *Application) watch() {
	e := app.Docker.GetEvents()

	app.Status.Add(1)
	go app.eventHandler(e)

	log.Printf("Starting main process")
	app.Status.Wait()
	log.Printf("Stopping cleanly via EOF")
}
<|endoftext|>


<commit_before>// Copyright (c) Ilia Kravets, 2015. All rights reserved. PROVIDED "AS IS"
// WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED. See LICENSE file for details.

package pcap2log

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/jessevdk/go-flags"
	"github.com/kr/pretty"
	"io"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
)

var textFrameSeparator []byte = []byte("\nFrame ")
var textFrameSeparator1 []byte = []byte("Frame ")

func splitTextFrames(data []byte, atEOF bool) (advance int, token []byte, err error) {
	separatorIndex := bytes.Index(data, textFrameSeparator)
	if separatorIndex == -1 {
		if atEOF {
			log.Println("WARNING skipping before EOF:", string(data))
			return len(data), nil, nil
		} else {
			return 0, nil, nil
		}
	}
	if separatorIndex != 0 {
		if bytes.HasPrefix(data, textFrameSeparator1) {
			separatorIndex = 0
		} else {
			log.Println("WARNING skipping prefix:", string(data[:separatorIndex]))
			return separatorIndex, nil, nil
		}
	}
	// find start of the next frame
	const skip1 = 5
	separatorIndex = bytes.Index(data[skip1:], textFrameSeparator)
	if separatorIndex == -1 {
		if atEOF {
			return len(data), data, nil
		} else {
			return 0, nil, nil
		}
	}
	separatorIndex += skip1

	return separatorIndex, data[:separatorIndex], nil
}

type translator struct {
	r   io.Reader
	w   io.Writer
	sim simulator
	// current message data
	kvStr       map[string]string
	kvInt       map[string]uint
	msgType     byte
	refNumDelta []uint // for Block Single Side Delete Message
	qom         QOMessage
}

func NewTranslator(r io.Reader, w io.Writer) translator {
	return translator{
		r:   r,
		w:   w,
		sim: NewSimulator(w),
	}
}

type MarketSide byte

const (
	MasketSideUnknown MarketSide = 0
	MarketSideBuy                = 'B'
	MarketSideSell               = 'S'
)

type MessageType byte

const (
	MessageTypeUnknown MessageType = iota
	MessageTypeQuoteAdd
	MessageTypeQuoteReplace
	MessageTypeQuoteDelete
	MessageTypeOrderAdd
	MessageTypeOrderExecute
	MessageTypeOrderExecuteWPrice
	MessageTypeOrderCancel
	MessageTypeOrderUpdate
	MessageTypeOrderReplace
	MessageTypeOrderDelete
	MessageTypeBlockOrderDelete
)

type OptionId uint

const OptionIdUnknown OptionId = 0

type OrderSide struct {
	refNumDelta     uint
	origRefNumDelta uint
	price           uint
	size            uint
	side            MarketSide
}
type QOMessage struct {
	typ          MessageType
	timestamp    uint
	optionId     OptionId
	side1        OrderSide
	side2        OrderSide
	sseCrossNum  uint
	sseMatchNum  uint
	ssePrintable bool
	ssuReason    byte
	bssdNum      uint
	bssdRefs     []uint
}

var charToMessageType = []MessageType{
	'j': MessageTypeQuoteAdd,
	'J': MessageTypeQuoteAdd,
	'k': MessageTypeQuoteReplace,
	'K': MessageTypeQuoteReplace,
	'Y': MessageTypeQuoteDelete,
	'a': MessageTypeOrderAdd,
	'A': MessageTypeOrderAdd,
	'E': MessageTypeOrderExecute,
	'C': MessageTypeOrderExecuteWPrice,
	'X': MessageTypeOrderCancel,
	'G': MessageTypeOrderUpdate,
	'u': MessageTypeOrderReplace,
	'U': MessageTypeOrderReplace,
	'D': MessageTypeOrderDelete,
	'Z': MessageTypeBlockOrderDelete,
}

func (t *translator) translateQOMessage() {
	t.qom = QOMessage{
		typ:       charToMessageType[t.msgType],
		timestamp: t.kvInt["Timestamp"],
	}
	if oid := t.kvInt["Option ID"]; oid != 0 {
		t.qom.optionId = OptionId(oid)
	} else {
		t.qom.optionId = OptionIdUnknown
	}
	switch t.msgType {
	case 'T', 'L', 'S', 'H', 'O', 'Q', 'I': // ignore Seconds, Base Reference, System,  Options Trading Action, Option Open, Cross Trade, NOII
	case 'j': // Add Quote
		t.qom.side1 = OrderSide{
			side:        MarketSideBuy,
			refNumDelta: t.kvInt["Bid Reference Number Delta"],
			size:        t.kvInt["Bid Size"],
			price:       t.kvInt["Bid Price"],
		}
		t.qom.side2 = OrderSide{
			side:        MarketSideSell,
			refNumDelta: t.kvInt["Ask Reference Number Delta"],
			size:        t.kvInt["Ask Size"],
			price:       t.kvInt["Ask Price"],
		}
	case 'J': // Add Quote
		t.qom.side1 = OrderSide{
			side:        MarketSideBuy,
			refNumDelta: t.kvInt["Bid Reference Number Delta"],
			size:        t.kvInt["Bid Size"],
			price:       t.kvInt["Bid"],
		}
		t.qom.side2 = OrderSide{
			side:        MarketSideSell,
			refNumDelta: t.kvInt["Ask Reference Number Delta"],
			size:        t.kvInt["Ask Size"],
			price:       t.kvInt["Ask"],
		}
	case 'k', 'K': // Quote Replace
		t.qom.side1 = OrderSide{
			side:            MarketSideBuy,
			refNumDelta:     t.kvInt["Bid Reference Number Delta"],
			origRefNumDelta: t.kvInt["Original Bid Reference Number Delta"],
			size:            t.kvInt["Bid Size"],
			price:           t.kvInt["Bid Price"],
		}
		t.qom.side2 = OrderSide{
			side:            MarketSideSell,
			refNumDelta:     t.kvInt["Ask Reference Delta Number"],
			origRefNumDelta: t.kvInt["Original Ask Reference Number Delta"],
			size:            t.kvInt["Ask Size"],
			price:           t.kvInt["Ask Price"],
		}
	case 'Y': // Quote Delete
		t.qom.side1 = OrderSide{
			side:            MarketSideBuy,
			origRefNumDelta: t.kvInt["Bid Reference Number Delta"],
		}
		t.qom.side2 = OrderSide{
			side:            MarketSideSell,
			origRefNumDelta: t.kvInt["Ask Reference Number Delta"],
		}
	case 'a', 'A': // Add Order
		t.qom.side1 = OrderSide{
			side:        MarketSide(t.kvInt["Market Side"]),
			refNumDelta: t.kvInt["Order Reference Number Delta"],
			size:        t.kvInt["Volume"],
			price:       t.kvInt["Price"],
		}
	case 'E': // Single Side Executed
		t.qom.side1 = OrderSide{
			origRefNumDelta: t.kvInt["Reference Number Delta"],
			size:            t.kvInt["Executed Contracts"],
		}
		t.qom.sseCrossNum = t.kvInt["Cross Number"]
		t.qom.sseMatchNum = t.kvInt["Match Number"]
	case 'C': // Single Side Executed with Price
		t.qom.side1 = OrderSide{
			origRefNumDelta: t.kvInt["Reference Number Delta"],
			size:            t.kvInt["Volume"],
			price:           t.kvInt["Price"],
		}
		t.qom.sseCrossNum = t.kvInt["Cross Number"]
		t.qom.sseMatchNum = t.kvInt["Match Number"]
		t.qom.ssePrintable = t.kvStr["Printable"] == "Y"
	case 'X': //  Order Cancel
		t.qom.side1 = OrderSide{
			origRefNumDelta: t.kvInt["Order Reference Number Delta"],
			size:            t.kvInt["Cancelled Contracts"],
		}
	case 'G': // Single Side Update
		t.qom.side1 = OrderSide{
			origRefNumDelta: t.kvInt["Reference Number Delta"],
			price:           t.kvInt["Price"],
			size:            t.kvInt["Volume"],
		}
	case 'u', 'U': // Single Side Replace
		t.qom.side1 = OrderSide{
			refNumDelta:     t.kvInt["New Reference Number Delta"],
			origRefNumDelta: t.kvInt["Original Reference Number Delta"],
			price:           t.kvInt["Price"],
			size:            t.kvInt["Volume"],
		}
	case 'D': // Single Side Delete
		t.qom.side1 = OrderSide{
			origRefNumDelta: t.kvInt["Reference Number Delta"],
		}
	case 'Z': // Block Single Side Delete
		t.qom.bssdNum = t.kvInt["Total Number of Reference Number Deltas."]
		if uint(len(t.refNumDelta)) != t.qom.bssdNum {
			pretty.Println(t.kvInt)
			log.Fatalf("Unexpected number of refs in Z message (%d != %d)\n", t.qom.bssdNum, len(t.refNumDelta))
		}
		t.qom.bssdRefs = append([]uint(nil), t.refNumDelta...)
	default:
		s := pretty.Sprintf("%v", t)
		//log.Fatalf("Unknown message type %d (%c)\n%s\n", t.msgType, t.msgType, s)
		log.Printf("Unknown message type %d (%c)\n%s\n", t.msgType, t.msgType, s)
	}
}

func (t *translator) translate() {
	kvRegexp := regexp.MustCompile("(?m)^            ([^:]*): (.*)$")
	parValueRegexp := regexp.MustCompile(".*\\((\\d+)\\)")
	scanner := bufio.NewScanner(t.r)
	scanner.Split(splitTextFrames)
	for scanner.Scan() {
		//fmt.Println("=====================")
		//fmt.Println(scanner.Text())
		ittoMessages := strings.Split(scanner.Text(), "        ITTO ")
		if len(ittoMessages) == 1 {
			continue
		}
		for _, ittoMessage := range ittoMessages[1:] {
			matches := kvRegexp.FindAllStringSubmatch(ittoMessage, -1)
			t.kvStr = make(map[string]string)
			t.kvInt = make(map[string]uint)
			t.refNumDelta = nil
			t.msgType = 0
			for _, m := range matches {
				k := m[1]
				v := m[2]
				if t.msgType == 'Z' && k == "Reference Number Delta" {
					vInt, err := strconv.ParseUint(v, 0, 32)
					if err != nil {
						log.Fatal("Can't parse", v)
					}
					t.refNumDelta = append(t.refNumDelta, uint(vInt))
				} else {
					if _, ok := t.kvStr[k]; ok {
						pretty.Println(ittoMessage)
						pretty.Println(matches)
						pretty.Println(m)
						log.Fatal("Duplicate key ", k)
					}
					t.kvStr[k] = v
					vInt, err := strconv.ParseUint(v, 0, 32)
					if err == nil {
						t.kvInt[k] = uint(vInt)
					} else if matches := parValueRegexp.FindStringSubmatch(v); matches != nil {
						vInt, err := strconv.ParseUint(matches[1], 0, 32)
						t.kvInt[k] = uint(vInt)
						if err != nil {
							log.Fatal("Can't parse", v)
						}
						if k == "Message Type" {
							t.msgType = byte(vInt)
						}
					}
				}
			}
			t.translateQOMessage()
			t.sim.addMessage(&t.qom, t.msgType)
		}
	}
}

func getTsharkDump(fileName string, args []string) (reader io.Reader, finisher func()) {
	//pretty.Println(fileName, args)
	cmdArgs := []string{
		"-d", "udp.port==18000:10,moldudp64",
		"-V",
		"-r",
		fileName,
	}
	cmdArgs = append(cmdArgs, args...)
	cmd := exec.Command("tshark", cmdArgs...)
	reader, err := cmd.StdoutPipe()
	if err != nil {
		log.Fatal(err)
	}
	if err := cmd.Start(); err != nil {
		log.Fatal(err)
	}

	finisher = func() {
		if err := cmd.Wait(); err != nil {
			log.Fatal(err)
		}
	}
	return
}

type pcap2log struct {
	InputFileName  string                        `long:"input" short:"i" required:"y" value-name:"PCAP_FILE" description:"input pcap file to read"`
	OutputFileName string                        `long:"output" short:"o" value-name:"FILE" default:"/dev/stdout" default-mask:"stdout" description:"output file"`
	Args           struct{ TsharkArgs []string } `positional-args:"y"`
	shouldExecute  bool
}

func (p *pcap2log) Execute(args []string) error {
	p.shouldExecute = true
	return nil
}

func (p *pcap2log) maybeRun() error {
	if !p.shouldExecute {
		return nil
	}
	dumpReader, finisher := getTsharkDump(p.InputFileName, p.Args.TsharkArgs)
	defer finisher()
	outFile, err := os.OpenFile(p.OutputFileName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatal(err)
	}
	defer outFile.Close()
	t := NewTranslator(dumpReader, outFile)
	t.translate()
	return nil
}

func InitArgv(parser *flags.Parser) func() error {
	var command pcap2log
	parser.AddCommand("pcap2log",
		"convert pcap file to simulator output",
		"",
		&command)
	return command.maybeRun
}

/*****************************************************************************/
// experiments and debugging

func main() {
	t := NewTranslator(os.Stdin, os.Stdout)
	t.translate()
	_ = pretty.Print
	_ = fmt.Print

}
<commit_msg>pcap2log: fix typo<commit_after>// Copyright (c) Ilia Kravets, 2015. All rights reserved. PROVIDED "AS IS"
// WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED. See LICENSE file for details.

package pcap2log

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/jessevdk/go-flags"
	"github.com/kr/pretty"
	"io"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
)

var textFrameSeparator []byte = []byte("\nFrame ")
var textFrameSeparator1 []byte = []byte("Frame ")

func splitTextFrames(data []byte, atEOF bool) (advance int, token []byte, err error) {
	separatorIndex := bytes.Index(data, textFrameSeparator)
	if separatorIndex == -1 {
		if atEOF {
			log.Println("WARNING skipping before EOF:", string(data))
			return len(data), nil, nil
		} else {
			return 0, nil, nil
		}
	}
	if separatorIndex != 0 {
		if bytes.HasPrefix(data, textFrameSeparator1) {
			separatorIndex = 0
		} else {
			log.Println("WARNING skipping prefix:", string(data[:separatorIndex]))
			return separatorIndex, nil, nil
		}
	}
	// find start of the next frame
	const skip1 = 5
	separatorIndex = bytes.Index(data[skip1:], textFrameSeparator)
	if separatorIndex == -1 {
		if atEOF {
			return len(data), data, nil
		} else {
			return 0, nil, nil
		}
	}
	separatorIndex += skip1

	return separatorIndex, data[:separatorIndex], nil
}

type translator struct {
	r   io.Reader
	w   io.Writer
	sim simulator
	// current message data
	kvStr       map[string]string
	kvInt       map[string]uint
	msgType     byte
	refNumDelta []uint // for Block Single Side Delete Message
	qom         QOMessage
}

func NewTranslator(r io.Reader, w io.Writer) translator {
	return translator{
		r:   r,
		w:   w,
		sim: NewSimulator(w),
	}
}

type MarketSide byte

const (
	MarketSideUnknown MarketSide = 0
	MarketSideBuy                = 'B'
	MarketSideSell               = 'S'
)

type MessageType byte

const (
	MessageTypeUnknown MessageType = iota
	MessageTypeQuoteAdd
	MessageTypeQuoteReplace
	MessageTypeQuoteDelete
	MessageTypeOrderAdd
	MessageTypeOrderExecute
	MessageTypeOrderExecuteWPrice
	MessageTypeOrderCancel
	MessageTypeOrderUpdate
	MessageTypeOrderReplace
	MessageTypeOrderDelete
	MessageTypeBlockOrderDelete
)

type OptionId uint

const OptionIdUnknown OptionId = 0

type OrderSide struct {
	refNumDelta     uint
	origRefNumDelta uint
	price           uint
	size            uint
	side            MarketSide
}
type QOMessage struct {
	typ          MessageType
	timestamp    uint
	optionId     OptionId
	side1        OrderSide
	side2        OrderSide
	sseCrossNum  uint
	sseMatchNum  uint
	ssePrintable bool
	ssuReason    byte
	bssdNum      uint
	bssdRefs     []uint
}

var charToMessageType = []MessageType{
	'j': MessageTypeQuoteAdd,
	'J': MessageTypeQuoteAdd,
	'k': MessageTypeQuoteReplace,
	'K': MessageTypeQuoteReplace,
	'Y': MessageTypeQuoteDelete,
	'a': MessageTypeOrderAdd,
	'A': MessageTypeOrderAdd,
	'E': MessageTypeOrderExecute,
	'C': MessageTypeOrderExecuteWPrice,
	'X': MessageTypeOrderCancel,
	'G': MessageTypeOrderUpdate,
	'u': MessageTypeOrderReplace,
	'U': MessageTypeOrderReplace,
	'D': MessageTypeOrderDelete,
	'Z': MessageTypeBlockOrderDelete,
}

func (t *translator) translateQOMessage() {
	t.qom = QOMessage{
		typ:       charToMessageType[t.msgType],
		timestamp: t.kvInt["Timestamp"],
	}
	if oid := t.kvInt["Option ID"]; oid != 0 {
		t.qom.optionId = OptionId(oid)
	} else {
		t.qom.optionId = OptionIdUnknown
	}
	switch t.msgType {
	case 'T', 'L', 'S', 'H', 'O', 'Q', 'I': // ignore Seconds, Base Reference, System,  Options Trading Action, Option Open, Cross Trade, NOII
	case 'j': // Add Quote
		t.qom.side1 = OrderSide{
			side:        MarketSideBuy,
			refNumDelta: t.kvInt["Bid Reference Number Delta"],
			size:        t.kvInt["Bid Size"],
			price:       t.kvInt["Bid Price"],
		}
		t.qom.side2 = OrderSide{
			side:        MarketSideSell,
			refNumDelta: t.kvInt["Ask Reference Number Delta"],
			size:        t.kvInt["Ask Size"],
			price:       t.kvInt["Ask Price"],
		}
	case 'J': // Add Quote
		t.qom.side1 = OrderSide{
			side:        MarketSideBuy,
			refNumDelta: t.kvInt["Bid Reference Number Delta"],
			size:        t.kvInt["Bid Size"],
			price:       t.kvInt["Bid"],
		}
		t.qom.side2 = OrderSide{
			side:        MarketSideSell,
			refNumDelta: t.kvInt["Ask Reference Number Delta"],
			size:        t.kvInt["Ask Size"],
			price:       t.kvInt["Ask"],
		}
	case 'k', 'K': // Quote Replace
		t.qom.side1 = OrderSide{
			side:            MarketSideBuy,
			refNumDelta:     t.kvInt["Bid Reference Number Delta"],
			origRefNumDelta: t.kvInt["Original Bid Reference Number Delta"],
			size:            t.kvInt["Bid Size"],
			price:           t.kvInt["Bid Price"],
		}
		t.qom.side2 = OrderSide{
			side:            MarketSideSell,
			refNumDelta:     t.kvInt["Ask Reference Delta Number"],
			origRefNumDelta: t.kvInt["Original Ask Reference Number Delta"],
			size:            t.kvInt["Ask Size"],
			price:           t.kvInt["Ask Price"],
		}
	case 'Y': // Quote Delete
		t.qom.side1 = OrderSide{
			side:            MarketSideBuy,
			origRefNumDelta: t.kvInt["Bid Reference Number Delta"],
		}
		t.qom.side2 = OrderSide{
			side:            MarketSideSell,
			origRefNumDelta: t.kvInt["Ask Reference Number Delta"],
		}
	case 'a', 'A': // Add Order
		t.qom.side1 = OrderSide{
			side:        MarketSide(t.kvInt["Market Side"]),
			refNumDelta: t.kvInt["Order Reference Number Delta"],
			size:        t.kvInt["Volume"],
			price:       t.kvInt["Price"],
		}
	case 'E': // Single Side Executed
		t.qom.side1 = OrderSide{
			origRefNumDelta: t.kvInt["Reference Number Delta"],
			size:            t.kvInt["Executed Contracts"],
		}
		t.qom.sseCrossNum = t.kvInt["Cross Number"]
		t.qom.sseMatchNum = t.kvInt["Match Number"]
	case 'C': // Single Side Executed with Price
		t.qom.side1 = OrderSide{
			origRefNumDelta: t.kvInt["Reference Number Delta"],
			size:            t.kvInt["Volume"],
			price:           t.kvInt["Price"],
		}
		t.qom.sseCrossNum = t.kvInt["Cross Number"]
		t.qom.sseMatchNum = t.kvInt["Match Number"]
		t.qom.ssePrintable = t.kvStr["Printable"] == "Y"
	case 'X': //  Order Cancel
		t.qom.side1 = OrderSide{
			origRefNumDelta: t.kvInt["Order Reference Number Delta"],
			size:            t.kvInt["Cancelled Contracts"],
		}
	case 'G': // Single Side Update
		t.qom.side1 = OrderSide{
			origRefNumDelta: t.kvInt["Reference Number Delta"],
			price:           t.kvInt["Price"],
			size:            t.kvInt["Volume"],
		}
	case 'u', 'U': // Single Side Replace
		t.qom.side1 = OrderSide{
			refNumDelta:     t.kvInt["New Reference Number Delta"],
			origRefNumDelta: t.kvInt["Original Reference Number Delta"],
			price:           t.kvInt["Price"],
			size:            t.kvInt["Volume"],
		}
	case 'D': // Single Side Delete
		t.qom.side1 = OrderSide{
			origRefNumDelta: t.kvInt["Reference Number Delta"],
		}
	case 'Z': // Block Single Side Delete
		t.qom.bssdNum = t.kvInt["Total Number of Reference Number Deltas."]
		if uint(len(t.refNumDelta)) != t.qom.bssdNum {
			pretty.Println(t.kvInt)
			log.Fatalf("Unexpected number of refs in Z message (%d != %d)\n", t.qom.bssdNum, len(t.refNumDelta))
		}
		t.qom.bssdRefs = append([]uint(nil), t.refNumDelta...)
	default:
		s := pretty.Sprintf("%v", t)
		//log.Fatalf("Unknown message type %d (%c)\n%s\n", t.msgType, t.msgType, s)
		log.Printf("Unknown message type %d (%c)\n%s\n", t.msgType, t.msgType, s)
	}
}

func (t *translator) translate() {
	kvRegexp := regexp.MustCompile("(?m)^            ([^:]*): (.*)$")
	parValueRegexp := regexp.MustCompile(".*\\((\\d+)\\)")
	scanner := bufio.NewScanner(t.r)
	scanner.Split(splitTextFrames)
	for scanner.Scan() {
		//fmt.Println("=====================")
		//fmt.Println(scanner.Text())
		ittoMessages := strings.Split(scanner.Text(), "        ITTO ")
		if len(ittoMessages) == 1 {
			continue
		}
		for _, ittoMessage := range ittoMessages[1:] {
			matches := kvRegexp.FindAllStringSubmatch(ittoMessage, -1)
			t.kvStr = make(map[string]string)
			t.kvInt = make(map[string]uint)
			t.refNumDelta = nil
			t.msgType = 0
			for _, m := range matches {
				k := m[1]
				v := m[2]
				if t.msgType == 'Z' && k == "Reference Number Delta" {
					vInt, err := strconv.ParseUint(v, 0, 32)
					if err != nil {
						log.Fatal("Can't parse", v)
					}
					t.refNumDelta = append(t.refNumDelta, uint(vInt))
				} else {
					if _, ok := t.kvStr[k]; ok {
						pretty.Println(ittoMessage)
						pretty.Println(matches)
						pretty.Println(m)
						log.Fatal("Duplicate key ", k)
					}
					t.kvStr[k] = v
					vInt, err := strconv.ParseUint(v, 0, 32)
					if err == nil {
						t.kvInt[k] = uint(vInt)
					} else if matches := parValueRegexp.FindStringSubmatch(v); matches != nil {
						vInt, err := strconv.ParseUint(matches[1], 0, 32)
						t.kvInt[k] = uint(vInt)
						if err != nil {
							log.Fatal("Can't parse", v)
						}
						if k == "Message Type" {
							t.msgType = byte(vInt)
						}
					}
				}
			}
			t.translateQOMessage()
			t.sim.addMessage(&t.qom, t.msgType)
		}
	}
}

func getTsharkDump(fileName string, args []string) (reader io.Reader, finisher func()) {
	//pretty.Println(fileName, args)
	cmdArgs := []string{
		"-d", "udp.port==18000:10,moldudp64",
		"-V",
		"-r",
		fileName,
	}
	cmdArgs = append(cmdArgs, args...)
	cmd := exec.Command("tshark", cmdArgs...)
	reader, err := cmd.StdoutPipe()
	if err != nil {
		log.Fatal(err)
	}
	if err := cmd.Start(); err != nil {
		log.Fatal(err)
	}

	finisher = func() {
		if err := cmd.Wait(); err != nil {
			log.Fatal(err)
		}
	}
	return
}

type pcap2log struct {
	InputFileName  string                        `long:"input" short:"i" required:"y" value-name:"PCAP_FILE" description:"input pcap file to read"`
	OutputFileName string                        `long:"output" short:"o" value-name:"FILE" default:"/dev/stdout" default-mask:"stdout" description:"output file"`
	Args           struct{ TsharkArgs []string } `positional-args:"y"`
	shouldExecute  bool
}

func (p *pcap2log) Execute(args []string) error {
	p.shouldExecute = true
	return nil
}

func (p *pcap2log) maybeRun() error {
	if !p.shouldExecute {
		return nil
	}
	dumpReader, finisher := getTsharkDump(p.InputFileName, p.Args.TsharkArgs)
	defer finisher()
	outFile, err := os.OpenFile(p.OutputFileName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatal(err)
	}
	defer outFile.Close()
	t := NewTranslator(dumpReader, outFile)
	t.translate()
	return nil
}

func InitArgv(parser *flags.Parser) func() error {
	var command pcap2log
	parser.AddCommand("pcap2log",
		"convert pcap file to simulator output",
		"",
		&command)
	return command.maybeRun
}

/*****************************************************************************/
// experiments and debugging

func main() {
	t := NewTranslator(os.Stdin, os.Stdout)
	t.translate()
	_ = pretty.Print
	_ = fmt.Print

}
<|endoftext|>


<commit_before>// Copyright 2011, Bryan Matsuo. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package main
/*
 *  Filename:    file.go
 *  Package:     main
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Sun Jul  3 16:57:42 PDT 2011
 *  Description:
 */
import (
    "os"
)

type File struct {
    Name    string
    User    string
    Pkg     string
	License LicenseType
    Repo    RepoType
    Host    RepoHost
}
type TestFile struct {
    Name    string
    Pkg     string
    License LicenseType
    Repo    RepoType
    Host    RepoHost
}

func (t TestFile) GenerateDictionary() map[string]string {
    var (
        test = t.Name + "_test.go"
        dict = map[string]string{
            "file":test,
            "name":AppConfig.Name,
            "email":AppConfig.Email,
            "date":DateString(),
            "year":YearString(),
            "gotarget":t.Pkg}
    )
    return dict
}
func (t TestFile) TemplatePath() []string {
    return []string{"testfiles", "pkg.t"}
}
func (t TestFile) Create() os.Error {
    var (
        dict = t.GenerateDictionary()
        errWrite = WriteTemplate(dict["file"], "library", dict, t.TemplatePath()...)
    )
    if errWrite != nil {
        return errWrite
    }
    // TODO: check the new file into git under certain conditions...
    return nil
}

func (f File) GenerateDictionary() map[string]string {
    var (
        lib = f.Name + ".go"
        dict = map[string]string{
            "file":lib,
            "name":AppConfig.Name,
            "email":AppConfig.Email,
            "date":DateString(),
            "year":YearString(),
            "gotarget":f.Pkg}
    )
    return dict
}
func (f File) TemplatePath() []string {
    return []string{"gofiles", "lib.t"}
}
func (f File) Create() os.Error {
    var (
        dict = f.GenerateDictionary()
        errWrite = WriteTemplate(dict["file"], "library", dict, f.TemplatePath()...)
    )
    if errWrite != nil {
        return errWrite
    }
    // TODO: check the new file into git under certain conditions...
    // Create a test for the new file.
    var (
        test = f.TestFile()
        errTestCreate = test.Create()
    )
    if errTestCreate != nil {
        return errTestCreate
    }
    return nil
}
func (f File) TestFile() TestFile {
    return TestFile{Name:f.Name, Pkg:f.Pkg, Repo:f.Repo, Host:f.Host}
}
<commit_msg>Clean up file.go<commit_after>// Copyright 2011, Bryan Matsuo. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package main
/*
 *  Filename:    file.go
 *  Package:     main
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Sun Jul  3 16:57:42 PDT 2011
 *  Description:
 */
import (
    "os"
)

type TestFile struct {
    Name    string
    Pkg     string
    License LicenseType
    Repo    RepoType
    Host    RepoHost
}

func (t TestFile) GenerateDictionary() map[string]string {
    var (
        test = t.Name + "_test.go"
        dict = map[string]string{
            "file":test,
            "name":AppConfig.Name,
            "email":AppConfig.Email,
            "date":DateString(),
            "year":YearString(),
            "gotarget":t.Pkg}
    )
    return dict
}

func (t TestFile) TemplatePath() []string { return []string{"testfiles", "pkg.t"} }

func (t TestFile) Create() os.Error {
    dict := t.GenerateDictionary()
    errWrite := WriteTemplate(dict["file"], "library", dict, t.TemplatePath()...)
    if errWrite != nil {
        return errWrite
    }
    // TODO: check the new file into git under certain conditions...
    return nil
}

type File struct {
    Name    string
    User    string
    Pkg     string
	License LicenseType
    Repo    RepoType
    Host    RepoHost
}

func (f File) GenerateDictionary() map[string]string {
    lib := f.Name + ".go"
    dict := map[string]string{
        "file":lib,
        "name":AppConfig.Name,
        "email":AppConfig.Email,
        "date":DateString(),
        "year":YearString(),
        "gotarget":f.Pkg}
    return dict
}

func (f File) TemplatePath() []string { return []string{"gofiles", "lib.t"} }

func (f File) Create() os.Error {
    dict := f.GenerateDictionary()
    errWrite := WriteTemplate(dict["file"], "library", dict, f.TemplatePath()...)
    if errWrite != nil {
        return errWrite
    }

    // TODO: check the new file into git under certain conditions...

    // Create a test for the new file.
    test := f.TestFile()
    if err := test.Create(); err != nil {
        return err
    }
    return nil
}

func (f File) TestFile() TestFile {
    return TestFile{Name:f.Name, Pkg:f.Pkg, Repo:f.Repo, Host:f.Host}
}
<|endoftext|>


<commit_before>package cf

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

	"code.cloudfoundry.org/lager"

	"autoscaler/models"
)

const (
	TokenTypeBearer = "Bearer"
	PathApp         = "/v2/apps"
	CFAppNotFound   = "CF-AppNotFound"
)

func (c *cfClient) GetApp(appId string) (*models.AppEntity, error) {
	url := c.conf.API + path.Join(PathApp, appId, "summary")
	c.logger.Debug("get-app-instances", lager.Data{"url": url})

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		c.logger.Error("get-app-instances-new-request", err)
		return nil, err
	}
	req.Header.Set("Authorization", TokenTypeBearer+" "+c.GetTokensWithRefresh().AccessToken)

	var resp *http.Response
	resp, err = c.httpClient.Do(req)

	if err != nil {
		c.logger.Error("get-app-instances-do-request", err)
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		if resp.StatusCode == 404 {
			respBody, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				c.logger.Error("failed-to-read-response-body-while-getting-app-summary", err, lager.Data{"appid": appId})
				return nil, err
			}
			var bodydata map[string]interface{}
			err = json.Unmarshal([]byte(respBody), &bodydata)
			if err != nil {
				c.logger.Error("failed-to-unmarshal-response-body-while-getting-app-summary", err, lager.Data{"appid": appId, "error_response": string(respBody)})
				return nil, err
			}
			errorDescription := bodydata["description"].(string)
			errorCode := bodydata["error_code"].(string)
			code := bodydata["code"].(float64)

			if errorCode == CFAppNotFound && code == 100004 {
				// Application does not exists
				err = models.NewAppNotFoundErr(errorDescription)
			} else {
				err = fmt.Errorf("failed getting application summary: [%d] %s: %s", resp.StatusCode, errorCode, errorDescription)
			}
			c.logger.Error("get-app-summary-response", err, lager.Data{"appid": appId, "statusCode": resp.StatusCode, "description": errorDescription, "errorCode": errorCode})
			return nil, err
		}
		// For Non 404 Error type
		err = fmt.Errorf("failed getting application summary: %s [%d] %s", url, resp.StatusCode, resp.Status)
		c.logger.Error("get-app-instances-response", err)
		return nil, err
	}

	appEntity := &models.AppEntity{}
	err = json.NewDecoder(resp.Body).Decode(appEntity)
	if err != nil {
		c.logger.Error("get-app-instances-decode", err)
		return nil, err
	}
	return appEntity, nil
}

func (c *cfClient) SetAppInstances(appID string, num int) error {
	url := c.conf.API + path.Join(PathApp, appID)
	c.logger.Debug("set-app-instances", lager.Data{"url": url})

	appEntity := models.AppEntity{
		Instances: num,
	}
	body, err := json.Marshal(appEntity)
	if err != nil {
		c.logger.Error("set-app-instances-marshal", err, lager.Data{"appid": appID, "appEntity": appEntity})
		return err
	}

	var req *http.Request
	req, err = http.NewRequest("PUT", url, bytes.NewReader(body))
	if err != nil {
		c.logger.Error("set-app-instances-new-request", err)
		return err
	}
	req.Header.Set("Authorization", TokenTypeBearer+" "+c.GetTokensWithRefresh().AccessToken)
	req.Header.Set("Content-Type", "application/json")

	var resp *http.Response
	resp, err = c.httpClient.Do(req)
	if err != nil {
		c.logger.Error("set-app-instances-do-request", err)
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusCreated {
		respBody, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			c.logger.Error("failed-to-read-response-body-while-setting-app-instance", err, lager.Data{"appid": appID})
			return err
		}
		var bodydata map[string]interface{}
		err = json.Unmarshal([]byte(respBody), &bodydata)
		if err != nil {
			c.logger.Error("failed-to-unmarshal-response-body-while-setting-app-instance", err, lager.Data{"appid": appID, "error_response": string(respBody)})
			return err
		}
		errorDescription := bodydata["description"].(string)
		errorCode := bodydata["error_code"].(string)
		err = fmt.Errorf("failed setting application instances: [%d] %s: %s", resp.StatusCode, errorCode, errorDescription)
		c.logger.Error("set-app-instances-response", err, lager.Data{"appid": appID, "statusCode": resp.StatusCode, "description": errorDescription, "errorCode": errorCode})
		return err
	}
	return nil
}
<commit_msg>update per comment<commit_after>package cf

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

	"code.cloudfoundry.org/lager"

	"autoscaler/models"
)

const (
	TokenTypeBearer = "Bearer"
	PathApp         = "/v2/apps"
	CFAppNotFound   = "CF-AppNotFound"
)

func (c *cfClient) GetApp(appID string) (*models.AppEntity, error) {
	url := c.conf.API + path.Join(PathApp, appID, "summary")
	c.logger.Debug("get-app-instances", lager.Data{"url": url})

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		c.logger.Error("get-app-instances-new-request", err)
		return nil, err
	}
	req.Header.Set("Authorization", TokenTypeBearer+" "+c.GetTokensWithRefresh().AccessToken)

	var resp *http.Response
	resp, err = c.httpClient.Do(req)

	if err != nil {
		c.logger.Error("get-app-instances-do-request", err)
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		if resp.StatusCode == 404 {
			respBody, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				c.logger.Error("failed-to-read-response-body-while-getting-app-summary", err, lager.Data{"appID": appID})
				return nil, err
			}
			var bodydata map[string]interface{}
			err = json.Unmarshal([]byte(respBody), &bodydata)
			if err != nil {
				err = fmt.Errorf("%s", string(respBody))
				c.logger.Error("failed-to-get-application-summary", err, lager.Data{"appID": appID})
				return nil, err
			}
			errorDescription := bodydata["description"].(string)
			errorCode := bodydata["error_code"].(string)
			code := bodydata["code"].(float64)

			if errorCode == CFAppNotFound && code == 100004 {
				// Application does not exists
				err = models.NewAppNotFoundErr(errorDescription)
			} else {
				err = fmt.Errorf("failed getting application summary: [%d] %s: %s", resp.StatusCode, errorCode, errorDescription)
			}
			c.logger.Error("get-app-summary-response", err, lager.Data{"appID": appID, "statusCode": resp.StatusCode, "description": errorDescription, "errorCode": errorCode})
			return nil, err
		}
		// For Non 404 Error type
		err = fmt.Errorf("failed getting application summary: %s [%d] %s", url, resp.StatusCode, resp.Status)
		c.logger.Error("get-app-instances-response", err)
		return nil, err
	}

	appEntity := &models.AppEntity{}
	err = json.NewDecoder(resp.Body).Decode(appEntity)
	if err != nil {
		c.logger.Error("get-app-instances-decode", err)
		return nil, err
	}
	return appEntity, nil
}

func (c *cfClient) SetAppInstances(appID string, num int) error {
	url := c.conf.API + path.Join(PathApp, appID)
	c.logger.Debug("set-app-instances", lager.Data{"url": url})

	appEntity := models.AppEntity{
		Instances: num,
	}
	body, err := json.Marshal(appEntity)
	if err != nil {
		c.logger.Error("set-app-instances-marshal", err, lager.Data{"appID": appID, "appEntity": appEntity})
		return err
	}

	var req *http.Request
	req, err = http.NewRequest("PUT", url, bytes.NewReader(body))
	if err != nil {
		c.logger.Error("set-app-instances-new-request", err)
		return err
	}
	req.Header.Set("Authorization", TokenTypeBearer+" "+c.GetTokensWithRefresh().AccessToken)
	req.Header.Set("Content-Type", "application/json")

	var resp *http.Response
	resp, err = c.httpClient.Do(req)
	if err != nil {
		c.logger.Error("set-app-instances-do-request", err)
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusCreated {
		respBody, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			c.logger.Error("failed-to-read-response-body-while-setting-app-instance", err, lager.Data{"appID": appID})
			return err
		}
		var bodydata map[string]interface{}
		err = json.Unmarshal([]byte(respBody), &bodydata)
		if err != nil {
			err = fmt.Errorf("%s", string(respBody))
			c.logger.Error("faileded-to-set-application-instances", err, lager.Data{"appID": appID})
			return err
		}
		errorDescription := bodydata["description"].(string)
		errorCode := bodydata["error_code"].(string)
		err = fmt.Errorf("failed setting application instances: [%d] %s: %s", resp.StatusCode, errorCode, errorDescription)
		c.logger.Error("set-app-instances-response", err, lager.Data{"appID": appID, "statusCode": resp.StatusCode, "description": errorDescription, "errorCode": errorCode})
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package fire

import (
	"errors"

	"github.com/kr/pretty"
	"github.com/ory-am/fosite"
	"golang.org/x/crypto/bcrypt"
	"golang.org/x/net/context"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type authenticatorStorage struct {
	db                 *mgo.Database
	ownerModel         Model
	ownerIDAttr        attribute
	ownerSecretAttr    attribute
	clientModel        Model
	clientIDAttr       attribute
	clientSecretAttr   attribute
	clientCallableAttr attribute
}

func (s *authenticatorStorage) GetClient(id string) (fosite.Client, error) {
	// prepare object
	obj := newStructPointer(s.clientModel)

	// query db
	err := s.db.C(s.clientModel.Collection()).Find(bson.M{
		s.clientIDAttr.dbField: id,
	}).One(obj)
	if err != nil {
		return nil, err
	}

	// initialize model
	_client := Init(obj.(Model))

	// TODO: We shouldn't use Attribute() as the field might be hidden.

	return &fosite.DefaultClient{
		ID:            id,
		Secret:        _client.Attribute(s.clientSecretAttr.name).([]byte),
		GrantTypes:    []string{"password", "client_credentials", "implicit"},
		ResponseTypes: []string{"token"},
		RedirectURIs:  []string{_client.Attribute(s.clientCallableAttr.name).(string)},
	}, nil
}

func (s *authenticatorStorage) CreateAccessTokenSession(ctx context.Context, signature string, request fosite.Requester) error {
	// create access token
	accessToken := Init(&AccessToken{
		Signature:     signature,
		RequestedAt:   request.GetRequestedAt(),
		GrantedScopes: request.GetGrantedScopes(),
	})

	// TODO: Save Client Id.

	// save access token
	return s.db.C(accessTokenModel.Collection()).Insert(accessToken)
}

func (s *authenticatorStorage) GetAccessTokenSession(ctx context.Context, signature string, session interface{}) (fosite.Requester, error) {
	// fetch access token
	var accessToken AccessToken
	err := s.db.C(accessTokenModel.Collection()).Find(bson.M{
		"signature": signature,
	}).One(&accessToken)
	if err != nil {
		return nil, err
	}

	// create request
	req := fosite.NewRequest()
	req.RequestedAt = accessToken.RequestedAt
	req.GrantedScopes = accessToken.GrantedScopes
	req.Session = session

	return req, nil
}

func (s *authenticatorStorage) DeleteAccessTokenSession(ctx context.Context, signature string) error {
	pretty.Println("DeleteAccessTokenSession", ctx, signature)
	return nil
}

func (s *authenticatorStorage) CreateRefreshTokenSession(ctx context.Context, signature string, request fosite.Requester) error {
	pretty.Println("CreateRefreshTokenSession", ctx, signature, request)
	return nil
}

func (s *authenticatorStorage) GetRefreshTokenSession(ctx context.Context, signature string, session interface{}) (fosite.Requester, error) {
	pretty.Println("GetRefreshTokenSession", ctx, signature, session)
	return nil, errors.New("error get refresh token session")
}

func (s *authenticatorStorage) DeleteRefreshTokenSession(ctx context.Context, signature string) error {
	pretty.Println("DeleteRefreshTokenSession", ctx, signature)
	return nil
}

func (s *authenticatorStorage) PersistRefreshTokenGrantSession(ctx context.Context, requestRefreshSignature, accessSignature, refreshSignature string, request fosite.Requester) error {
	pretty.Println("PersistRefreshTokenGrantSession", ctx, requestRefreshSignature, accessSignature, refreshSignature, request)
	return nil
}

func (s *authenticatorStorage) Authenticate(ctx context.Context, id string, secret string) error {
	// prepare object
	obj := newStructPointer(s.ownerModel)

	// query db
	err := s.db.C(s.ownerModel.Collection()).Find(bson.M{
		s.ownerIDAttr.dbField: id,
	}).One(obj)
	if err != nil {
		return err
	}

	// initialize model
	owner := Init(obj.(Model))

	// check secret
	return bcrypt.CompareHashAndPassword(owner.Attribute(s.ownerSecretAttr.name).([]byte), []byte(secret))
}
<commit_msg>refresh token is currently not implemented<commit_after>package fire

import (
	"github.com/kr/pretty"
	"github.com/ory-am/fosite"
	"golang.org/x/crypto/bcrypt"
	"golang.org/x/net/context"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type authenticatorStorage struct {
	db                 *mgo.Database
	ownerModel         Model
	ownerIDAttr        attribute
	ownerSecretAttr    attribute
	clientModel        Model
	clientIDAttr       attribute
	clientSecretAttr   attribute
	clientCallableAttr attribute
}

func (s *authenticatorStorage) GetClient(id string) (fosite.Client, error) {
	// prepare object
	obj := newStructPointer(s.clientModel)

	// query db
	err := s.db.C(s.clientModel.Collection()).Find(bson.M{
		s.clientIDAttr.dbField: id,
	}).One(obj)
	if err != nil {
		return nil, err
	}

	// initialize model
	_client := Init(obj.(Model))

	// TODO: We shouldn't use Attribute() as the field might be hidden.

	return &fosite.DefaultClient{
		ID:            id,
		Secret:        _client.Attribute(s.clientSecretAttr.name).([]byte),
		GrantTypes:    []string{"password", "client_credentials", "implicit"},
		ResponseTypes: []string{"token"},
		RedirectURIs:  []string{_client.Attribute(s.clientCallableAttr.name).(string)},
	}, nil
}

func (s *authenticatorStorage) CreateAccessTokenSession(ctx context.Context, signature string, request fosite.Requester) error {
	// create access token
	accessToken := Init(&AccessToken{
		Signature:     signature,
		RequestedAt:   request.GetRequestedAt(),
		GrantedScopes: request.GetGrantedScopes(),
	})

	// TODO: Save Client Id.

	// save access token
	return s.db.C(accessTokenModel.Collection()).Insert(accessToken)
}

func (s *authenticatorStorage) GetAccessTokenSession(ctx context.Context, signature string, session interface{}) (fosite.Requester, error) {
	// fetch access token
	var accessToken AccessToken
	err := s.db.C(accessTokenModel.Collection()).Find(bson.M{
		"signature": signature,
	}).One(&accessToken)
	if err != nil {
		return nil, err
	}

	// create request
	req := fosite.NewRequest()
	req.RequestedAt = accessToken.RequestedAt
	req.GrantedScopes = accessToken.GrantedScopes
	req.Session = session

	return req, nil
}

func (s *authenticatorStorage) DeleteAccessTokenSession(ctx context.Context, signature string) error {
	pretty.Println("DeleteAccessTokenSession", ctx, signature)
	return nil
}

func (s *authenticatorStorage) Authenticate(ctx context.Context, id string, secret string) error {
	// prepare object
	obj := newStructPointer(s.ownerModel)

	// query db
	err := s.db.C(s.ownerModel.Collection()).Find(bson.M{
		s.ownerIDAttr.dbField: id,
	}).One(obj)
	if err != nil {
		return err
	}

	// initialize model
	owner := Init(obj.(Model))

	// check secret
	return bcrypt.CompareHashAndPassword(owner.Attribute(s.ownerSecretAttr.name).([]byte), []byte(secret))
}
<|endoftext|>