<commit_before>/*
   Copyright 2015 Daniel Gruber, dgruber@univa.com

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package ugego

import (
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

// UserList is a Univa Grid Engine access control list or Department.
type UserList struct {
	Name    string
	Type    string
	FShare  int
	OTicket int
	Entries []string
}

// ParseUserList parses an Univa Grid Engine ACL by the given name into an struct.
// The expected structure of ul is a multiline string in the Univa Grid Engine format:
// name    exclude
// type    ACL
// fshare  0
// oticket 0
// entries daniel,root,%wheel
func ParseUserList(ul string) (userList *UserList, err error) {
	var ol UserList
	el := strings.Split(ul, "\n")
	if len(el) != 5 {
		return nil, errors.New(fmt.Sprintf("User list does not have 5 lines, it has %d.\n", len(el)))
	}
	name := strings.Split(el[0], "name")
	if len(name) != 2 {
		return nil, errors.New("Error during name parsing.")
	}
	ol.Name = strings.TrimSpace(name[1])
	t := strings.Split(el[1], "type")
	if len(t) != 2 {
		return nil, errors.New("Error during type parsing.")
	}
	ol.Type = strings.TrimSpace(t[1])
	fs := strings.Split(el[2], "fshare")
	if len(fs) != 2 {
		return nil, errors.New("Error during fshare parsing.")
	}
	ot := strings.Split(el[3], "oticket")
	if len(ot) != 2 {
		return nil, errors.New("Error during oticket parsing.")
	}
	ent := strings.Split(el[4], "entries")
	if len(ent) != 2 {
		return nil, errors.New("Error during entries parsing.")
	}
	if ol.FShare, err = strconv.Atoi(strings.TrimSpace(fs[1])); err != nil {
		return nil, err
	}
	if ol.OTicket, err = strconv.Atoi(strings.TrimSpace(ot[1])); err != nil {
		return nil, err
	}
	ol.Entries = strings.Split(strings.TrimSpace(ent[1]), ",")
	return &ol, nil
}

// GetUserLists calls qconf -su <listOfUl> and parses the output
// into UserList structs.
func GetUserLists(userlist ...string) ([]UserList, error) {
	rootPath := os.Getenv("SGE_ROOT")
	if rootPath == "" {
		return nil, errors.New("$SGE_ROOT environment variable not set")
	}
	qconf := fmt.Sprintf("%s/bin/lx-amd64/qconf", rootPath)

	// create comma separated list of user list names
	var csvUserList string
	for k, v := range userlist {
		if k == 0 {
			csvUserList = fmt.Sprintf("%s", v)
		} else {
			csvUserList = fmt.Sprintf("%s,%s", csvUserList, v)
		}
	}
	log.Printf("Calling: qconf -su %s\n", csvUserList)

	cmd := exec.Command(qconf, "-su", csvUserList)
	if out, errOut := cmd.CombinedOutput(); errOut == nil {
		// empty line is the delimiter
		uls := strings.Split(string(out), "\n\n")
		if len(uls) == 0 {
			return nil, errors.New(fmt.Sprintf("Could not split output: %s", string(out)))
		}
		outputList := make([]UserList, len(uls), len(uls))
		for i, ul := range uls {
			parsedUserList, errParse := ParseUserList(strings.TrimSpace(uls[i]))
			if errParse != nil {
				log.Printf("Error during parsing user list: %s\n%s\n", errParse, ul)
				log.Printf("%s", ul)
				return nil, errParse
			}
			outputList[i] = *parsedUserList
		}
		return outputList, nil
	} else {
		log.Printf("Error: %s\n", out)
		return nil, errOut
	}
}
<commit_msg>DG: Removed logging.<commit_after>/*
   Copyright 2015 Daniel Gruber, dgruber@univa.com

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package ugego

import (
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

// UserList is a Univa Grid Engine access control list or Department.
type UserList struct {
	Name    string
	Type    string
	FShare  int
	OTicket int
	Entries []string
}

// ParseUserList parses an Univa Grid Engine ACL by the given name into an struct.
// The expected structure of ul is a multiline string in the Univa Grid Engine format:
// name    exclude
// type    ACL
// fshare  0
// oticket 0
// entries daniel,root,%wheel
func ParseUserList(ul string) (userList *UserList, err error) {
	var ol UserList
	el := strings.Split(ul, "\n")
	if len(el) != 5 {
		return nil, errors.New(fmt.Sprintf("User list does not have 5 lines, it has %d.\n", len(el)))
	}
	name := strings.Split(el[0], "name")
	if len(name) != 2 {
		return nil, errors.New("Error during name parsing.")
	}
	ol.Name = strings.TrimSpace(name[1])
	t := strings.Split(el[1], "type")
	if len(t) != 2 {
		return nil, errors.New("Error during type parsing.")
	}
	ol.Type = strings.TrimSpace(t[1])
	fs := strings.Split(el[2], "fshare")
	if len(fs) != 2 {
		return nil, errors.New("Error during fshare parsing.")
	}
	ot := strings.Split(el[3], "oticket")
	if len(ot) != 2 {
		return nil, errors.New("Error during oticket parsing.")
	}
	ent := strings.Split(el[4], "entries")
	if len(ent) != 2 {
		return nil, errors.New("Error during entries parsing.")
	}
	if ol.FShare, err = strconv.Atoi(strings.TrimSpace(fs[1])); err != nil {
		return nil, err
	}
	if ol.OTicket, err = strconv.Atoi(strings.TrimSpace(ot[1])); err != nil {
		return nil, err
	}
	ol.Entries = strings.Split(strings.TrimSpace(ent[1]), ",")
	return &ol, nil
}

// GetUserLists calls qconf -su <listOfUl> and parses the output
// into UserList structs.
func GetUserLists(userlist ...string) ([]UserList, error) {
	rootPath := os.Getenv("SGE_ROOT")
	if rootPath == "" {
		return nil, errors.New("$SGE_ROOT environment variable not set")
	}
	qconf := fmt.Sprintf("%s/bin/lx-amd64/qconf", rootPath)

	// create comma separated list of user list names
	var csvUserList string
	for k, v := range userlist {
		if k == 0 {
			csvUserList = fmt.Sprintf("%s", v)
		} else {
			csvUserList = fmt.Sprintf("%s,%s", csvUserList, v)
		}
	}

	cmd := exec.Command(qconf, "-su", csvUserList)
	if out, errOut := cmd.CombinedOutput(); errOut == nil {
		// empty line is the delimiter
		uls := strings.Split(string(out), "\n\n")
		if len(uls) == 0 {
			return nil, errors.New(fmt.Sprintf("Could not split output: %s", string(out)))
		}
		outputList := make([]UserList, len(uls), len(uls))
		for i, ul := range uls {
			parsedUserList, errParse := ParseUserList(strings.TrimSpace(uls[i]))
			if errParse != nil {
				log.Printf("Error during parsing user list: %s\n%s\n", errParse, ul)
				log.Printf("%s", ul)
				return nil, errParse
			}
			outputList[i] = *parsedUserList
		}
		return outputList, nil
	} else {
		log.Printf("Error: %s\n", out)
		return nil, errOut
	}
}
<|endoftext|>


<commit_before>package exec

import (
	"gosubli.me/mg"
)

type Switch struct {
	Case        string
	Attr        map[string]interface{}
	Negative    bool
	Discard     bool
	Fallthrough bool
}

type Exec struct {
	Cid    string
	Stream string

	Wd            string
	Input         string
	DiscardStdout bool
	DiscardStderr bool
	Env           map[string]string
	Cmd           string
	Args          []string
	Switch        []*Switch

	b *mg.Broker
}

func (e *Exec) Call() (interface{}, string) {
	return nil, "N/I"
}

func init() {
	mg.Register("exec", func(b *mg.Broker) mg.Caller {
		return &Exec{
			b: b,
		}
	})
}
<commit_msg>* initial re-impl if meth/exec (minus the 9o features like switch)<commit_after>package exec

import (
	"gosubli.me/mg"
	"gosubli.me/sink"
	"os/exec"
	"strings"
	"sync"
	"time"
)

var (
	zChunk  = []byte{}
	zChunks = [][]byte{}
)

type StreamResp struct {
	Chunks [][]byte
	Stream string
	End    bool
}

type Resp struct {
	Chunks [][]byte
	Ok     bool
	Dur    string
}

type Switch struct {
	Case        string
	Attr        map[string]interface{}
	Negative    bool
	Discard     bool
	Fallthrough bool
}

type Exec struct {
	Cid    string
	Stream string

	Wd            string
	Input         string
	DiscardStdout bool
	DiscardStderr bool
	Env           map[string]string
	Cmd           string
	Args          []string
	Switch        []*Switch

	brk  *mg.Broker
	sink *sink.Chan
	buf  struct {
		s     [][]byte
		n     int
		merge bool
	}
}

func (e *Exec) Call() (interface{}, string) {
	c := exec.Command(e.Cmd, e.Args...)
	c.Dir = e.Wd
	c.Env = mg.Env(e.Env)
	if !e.DiscardStdout {
		c.Stdout = e.sink
	}
	if !e.DiscardStderr {
		c.Stderr = e.sink
	}
	if e.Input != "" {
		c.Stdin = strings.NewReader(e.Input)
	}

	wg := &sync.WaitGroup{}
	wg.Add(1)
	go e.stream(wg)
	err, dur := procs.Run(e.Cid, c)
	e.sink.Close()
	wg.Wait()
	res := Resp{
		Dur:    dur.String(),
		Chunks: e.chunks(),
		Ok:     c.ProcessState != nil && c.ProcessState.Success(),
	}

	if err != nil {
		return res, err.Error()
	}
	return res, ""
}

func (e *Exec) stream(wg *sync.WaitGroup) {
	defer wg.Done()

	tck := time.NewTicker(500 * time.Millisecond)
	for {
		select {
		case s, ok := <-e.sink.C:
			eof := !ok
			e.put(s, eof)
			if eof {
				return
			}
		case <-tck.C:
			if len(e.buf.s) != 0 {
				e.flush(false)
			}
		}
	}
}

func (e *Exec) put(s []byte, eof bool) {
	switch {
	case crSuffix(s):
		e.buf.merge = false
		e.buf.s = append(e.buf.s, s)
	case !e.buf.merge || len(e.buf.s) == 0:
		e.buf.merge = true
		e.buf.s = append(e.buf.s, zChunk)
		fallthrough
	default:
		i := len(e.buf.s) - 1
		e.buf.s[i] = append(e.buf.s[i], s...)
	}

	e.buf.n += len(s)
	if eof || e.buf.n >= 32*1024 {
		e.flush(eof)
	}
}

func (e *Exec) chunks() [][]byte {
	if e.buf.s != nil {
		return e.buf.s
	}
	// be nice to the client by not sending them `null` when they expect an array
	return zChunks
}

func (e *Exec) flush(eof bool) {
	if e.Stream == "" {
		return
	}

	e.brk.Send(mg.Response{
		Token: e.Stream,
		Data: StreamResp{
			Chunks: e.chunks(),
			End:    eof,
		},
	})

	for i, s := range e.buf.s {
		e.buf.s[i] = s[:0]
	}
	e.buf.s = e.buf.s[:0]
	e.buf.n = 0
}

func crSuffix(s []byte) bool {
	i := len(s) - 1
	return i >= 0 && s[i] == '\r'
}

func init() {
	mg.Register("exec", func(b *mg.Broker) mg.Caller {
		return &Exec{
			brk:  b,
			sink: sink.NewBufferedChan(100),
		}
	})
}
<|endoftext|>


<commit_before>package core

import (
	"fmt"

	"testing"
)
import (
	gstring "github.com/hraberg/cljs.go/goog/string"
	"github.com/hraberg/cljs.go/js"
	"github.com/hraberg/cljs.go/js/Math"
	"github.com/stretchr/testify/assert"
)

/*
;; IFn

;; Used in core.cljs apply, which uses apply-to built by core.clj gen-apply-to.
;; It falls back to JS .apply it applyTo doesn't exist on the passed in fn.
.-cljs$lang$maxFixedArity ;; a field on the fn
.-cljs$lang$applyTo ;; accessed as field to see if it's there
.cljs$lang$applyTo  ;; then called as fn if it exists.

;; Used by the dispatch fn to actually invoke the various overloaded fns, uses JS arguments in a switch.
;; See emit* :fn and emit* :invoke
;; emit* :fn will emit a single fn or a dispatch fn around the real overlaoded fns.
;; At times these are called directly, like cljs.core.str.cljs$core$IFn$_invoke$arity$1(~{}) in cljs.core/str
.cljs$core$IFn$_invoke$arity$variadic
.cljs$core$IFn$_invoke$arity$N

;; defprotocol

.-cljs$lang$protocol_mask$partitionN$

;; deftype

.-cljs$lang$type
.-cljs$lang$ctorStr
.-cljs$lang$ctorPrWriter

*/

func Test_Main(t *testing.T) {
	mainWasCalled := false
	STAR_main_cli_fn_STAR_ = Fn(func(args ...interface{}) interface{} {
		mainWasCalled = true
		return nil
	})
	Main()
	assert.True(t, mainWasCalled)
}

var Baz = Fn(func(args ...interface{}) interface{} {
	x := args[0]
	xs := args[1:] // this should be an array-seq (an IndexedSeq backed by slices or arrays)
	_ = x
	return xs
}, func(x interface{}) interface{} {
	return x
})

func Test_Invoke(t *testing.T) {
	assert.True(t, NativeSatisifes_QMARK_.Invoke_Arity2(Symbol.Invoke_Arity2("cljs.core", "IFn"), Baz).(bool))
	PanicsWith(t, "Invalid arity: 0", func() { Baz.(*AFn).Call() })
	PanicsWith(t, "Invalid arity: 0", func() { Baz.Invoke_Arity0() })
	assert.Equal(t, 1, Baz.(*AFn).MaxFixedArity)
	assert.True(t, Baz.(*AFn).IsVariadic())
	assert.Equal(t, "Hello", Baz.Invoke_Arity1("Hello"))
	assert.Equal(t, "Hello", Invoke_.Invoke_Arity2(Baz, "Hello"))
	assert.Equal(t, []interface{}{"World"}, Baz.(*AFn).Call("Hello", "World"))
	assert.Equal(t, []interface{}{"World"}, Baz.Invoke_ArityVariadic("Hello", "World"))
	assert.Equal(t, []interface{}{"World"}, Invoke_.Invoke_ArityVariadic(Baz, "Hello", "World"))
	assert.Equal(t, []interface{}{"World"}, Invoke_.(*AFn).Call(Baz, "Hello", "World"))
	assert.Equal(t, []interface{}{"World"}, Apply.Invoke_ArityVariadic(Baz, "Hello", []interface{}{"World"}))

	assert.Equal(t, []interface{}{"World"}, Baz.Invoke_Arity2("Hello", "World"))
	assert.Equal(t, []interface{}{"World", "Space"}, Baz.Invoke_Arity3("Hello", "World", "Space"))

	assert.Equal(t, "", Str.Invoke_Arity0())
	assert.Equal(t, "Hello", Str.Invoke_Arity1("Hello"))
	assert.Equal(t, "1", Str.Invoke_Arity1(1))
	assert.Equal(t, "HelloClojureWorld", Str.Invoke_ArityVariadic("Hello", "Clojure", "World"))
}

func Test_PrimitiveFn(t *testing.T) {
	fib := func(this *AFnPrimtive) IFn {
		return Fn(this, func(n float64) float64 {
			if n == 0.0 {
				return 0.0
			} else if n == 1.0 {
				return 1.0
			} else {
				return this.Arity1FF(n-1.0) + this.Arity1FF(n-2.0)
			}
		})
	}(&AFnPrimtive{})
	assert.NotNil(t, fib.(*AFnPrimtive).Arity1FF)
	assert.NotNil(t, fib.(*AFnPrimtive).Arity1)
	assert.Nil(t, fib.(*AFnPrimtive).Arity0F)
	assert.Equal(t, 832040, fib.Invoke_Arity1(30.0))

	odd := Fn(&AFnPrimtive{}, func(n interface{}) bool {
		return int(n.(float64))%2 != 0
	})
	assert.NotNil(t, odd.(*AFnPrimtive).Arity1IB)
	assert.NotNil(t, odd.(*AFnPrimtive).Arity1)
	assert.True(t, odd.Invoke_Arity1(1.0).(bool))
	assert.False(t, odd.(*AFnPrimtive).Arity1IB(2.0))
}

func Test_Protocols(t *testing.T) {
	symbol := Symbol.Invoke_Arity2("foo", "bar")

	assert.True(t, NativeSatisifes_QMARK_.Invoke_Arity2(Symbol.Invoke_Arity2("cljs.core", "INamed"), symbol).(bool))
	assert.True(t, NativeSatisifes_QMARK_.Invoke_Arity2(Symbol.Invoke_Arity2("cljs.core", "IFn"), symbol).(bool))
	assert.Equal(t, "foo", symbol.(INamed).Namespace_Arity1())
	assert.Equal(t, "foo", Namespace_.Invoke_Arity1(symbol))
	PanicsWith(t, "Invalid arity: 0", func() { symbol.(*CljsCoreSymbol).Invoke_Arity0() })
	PanicsWith(t, "Invalid arity: 0", func() { Invoke_.Invoke_Arity1(symbol) })

	assert.Equal(t, "foo", Namespace.Invoke_Arity1(symbol))
	PanicsWith(t, "Doesn't support namespace: 2", func() { Namespace.Invoke_Arity1(2) })

	assert.Equal(t, "bar", Name.Invoke_Arity1(symbol))
	PanicsWith(t, "Doesn't support name: 2", func() { Name.Invoke_Arity1(2) })
	assert.Equal(t, "baz", Name.Invoke_Arity1("baz"))

	assert.Equal(t, "foo", Invoke_.Invoke_Arity2(Namespace_, symbol))
	assert.Equal(t, "bar", symbol.(INamed).Name_Arity1())
	assert.Equal(t, "foo/bar", symbol.(fmt.Stringer).String())

	foo, bar := Symbol.Invoke_Arity1("foo"), Symbol.Invoke_Arity1("bar")
	m := ObjMap(map[interface{}]interface{}{foo: "bar"})

	assert.True(t, NativeSatisifes_QMARK_.Invoke_Arity2(Symbol.Invoke_Arity2("cljs.core", "ILookup"), m).(bool))
	assert.Equal(t, "bar", m.Lookup_Arity2(foo))
	assert.Nil(t, Lookup_.(*AFn).Call(m, bar))
	assert.Equal(t, "baz", m.Lookup_Arity3(bar, "baz"))

	assert.Equal(t, "bar", foo.(IFn).Invoke_Arity1(m))
	assert.Equal(t, "bar", Invoke_.Invoke_Arity2(foo, m))
	assert.Nil(t, bar.(IFn).Invoke_Arity1(m))
	assert.Equal(t, "baz", bar.(IFn).Invoke_Arity2(m, "baz"))
}

func Test_InteropViaReflection(t *testing.T) {
	sb := &gstring.StringBuffer{js.JSString("foo")}
	assert.Equal(t, "foo", NativeGetInstanceField.Invoke_Arity2(sb, "Buffer"),
		"(.-buffer sb)")
	NativeSetInstanceField.Invoke_Arity3(sb, "Buffer", "bar")
	assert.Equal(t, "bar", sb.Buffer,
		"(set! (.-buffer sb) \"bar\")")
	assert.Equal(t, sb, NativeInvokeInstanceMethod.Invoke_Arity3(sb, "Append", []interface{}{"baz"}))
	assert.Equal(t, "barbaz", sb.Buffer,
		"(.append sb \"baz\")")

	assert.Equal(t, js.Number.MAX_VALUE, NativeGetInstanceField.Invoke_Arity2(js.Number, "MAX_VALUE"),
		"(.-MAX-VALUE js/Number)")

	assert.Equal(t, 3, NativeInvokeFunc.Invoke_Arity2(Math.Floor, []interface{}{3.14}),
		"(Math/floor 3.14)")
	assert.Equal(t, "3.14", NativeInvokeFunc.Invoke_Arity2(fmt.Sprint, []interface{}{3.14}),
		"(fmt/Sprint 3.14)")
	assert.Equal(t, 3.14, NativeInvokeFunc.Invoke_Arity2(js.ParseFloat, []interface{}{js.JSString("3.14")}),
		"(js/parseFloat \"3.14\")")

	assert.Equal(t, "ABC", NativeInvokeFunc.Invoke_Arity2(js.String.FromCharCode, []interface{}{65, 66, 67}),
		"(.fromCharCode js/String 65 66 67)")
}

func PanicsWith(t *testing.T, message string, f assert.PanicTestFunc) {
	assert.Equal(t, message, func() (message string) {
		defer func() { message = fmt.Sprint(recover()) }()
		f()
		assert.Fail(t, "should panic")
		return
	}())
}
<commit_msg>Ok, last JS type stuff for a while<commit_after>package core

import (
	"fmt"

	"testing"
)
import (
	gstring "github.com/hraberg/cljs.go/goog/string"
	"github.com/hraberg/cljs.go/js"
	"github.com/hraberg/cljs.go/js/Math"
	"github.com/stretchr/testify/assert"
)

/*
;; IFn

;; Used in core.cljs apply, which uses apply-to built by core.clj gen-apply-to.
;; It falls back to JS .apply it applyTo doesn't exist on the passed in fn.
.-cljs$lang$maxFixedArity ;; a field on the fn
.-cljs$lang$applyTo ;; accessed as field to see if it's there
.cljs$lang$applyTo  ;; then called as fn if it exists.

;; Used by the dispatch fn to actually invoke the various overloaded fns, uses JS arguments in a switch.
;; See emit* :fn and emit* :invoke
;; emit* :fn will emit a single fn or a dispatch fn around the real overlaoded fns.
;; At times these are called directly, like cljs.core.str.cljs$core$IFn$_invoke$arity$1(~{}) in cljs.core/str
.cljs$core$IFn$_invoke$arity$variadic
.cljs$core$IFn$_invoke$arity$N

;; defprotocol

.-cljs$lang$protocol_mask$partitionN$

;; deftype

.-cljs$lang$type
.-cljs$lang$ctorStr
.-cljs$lang$ctorPrWriter

*/

func Test_Main(t *testing.T) {
	mainWasCalled := false
	STAR_main_cli_fn_STAR_ = Fn(func(args ...interface{}) interface{} {
		mainWasCalled = true
		return nil
	})
	Main()
	assert.True(t, mainWasCalled)
}

var Baz = Fn(func(args ...interface{}) interface{} {
	x := args[0]
	xs := args[1:] // this should be an array-seq (an IndexedSeq backed by slices or arrays)
	_ = x
	return xs
}, func(x interface{}) interface{} {
	return x
})

func Test_Invoke(t *testing.T) {
	assert.True(t, NativeSatisifes_QMARK_.Invoke_Arity2(Symbol.Invoke_Arity2("cljs.core", "IFn"), Baz).(bool))
	PanicsWith(t, "Invalid arity: 0", func() { Baz.(*AFn).Call() })
	PanicsWith(t, "Invalid arity: 0", func() { Baz.Invoke_Arity0() })
	assert.Equal(t, 1, Baz.(*AFn).MaxFixedArity)
	assert.True(t, Baz.(*AFn).IsVariadic())
	assert.Equal(t, "Hello", Baz.Invoke_Arity1("Hello"))
	assert.Equal(t, "Hello", Invoke_.Invoke_Arity2(Baz, "Hello"))
	assert.Equal(t, []interface{}{"World"}, Baz.(*AFn).Call("Hello", "World"))
	assert.Equal(t, []interface{}{"World"}, Baz.Invoke_ArityVariadic("Hello", "World"))
	assert.Equal(t, []interface{}{"World"}, Invoke_.Invoke_ArityVariadic(Baz, "Hello", "World"))
	assert.Equal(t, []interface{}{"World"}, Invoke_.(*AFn).Call(Baz, "Hello", "World"))
	assert.Equal(t, []interface{}{"World"}, Apply.Invoke_ArityVariadic(Baz, "Hello", []interface{}{"World"}))

	assert.Equal(t, []interface{}{"World"}, Baz.Invoke_Arity2("Hello", "World"))
	assert.Equal(t, []interface{}{"World", "Space"}, Baz.Invoke_Arity3("Hello", "World", "Space"))

	assert.Equal(t, "", Str.Invoke_Arity0())
	assert.Equal(t, "Hello", Str.Invoke_Arity1("Hello"))
	assert.Equal(t, "1", Str.Invoke_Arity1(1))
	assert.Equal(t, "HelloClojureWorld", Str.Invoke_ArityVariadic("Hello", "Clojure", "World"))
}

func Test_PrimitiveFn(t *testing.T) {
	fib := func(this *AFnPrimtive) IFn {
		return Fn(this, func(n float64) float64 {
			if n == 0.0 {
				return 0.0
			} else if n == 1.0 {
				return 1.0
			} else {
				return this.Arity1FF(n-1.0) + this.Arity1FF(n-2.0)
			}
		})
	}(&AFnPrimtive{})
	assert.NotNil(t, fib.(*AFnPrimtive).Arity1FF)
	assert.NotNil(t, fib.(*AFnPrimtive).Arity1)
	assert.Nil(t, fib.(*AFnPrimtive).Arity0F)
	assert.Equal(t, 832040, fib.Invoke_Arity1(30.0))

	odd := Fn(&AFnPrimtive{}, func(n interface{}) bool {
		return int(n.(float64))%2 != 0
	})
	assert.NotNil(t, odd.(*AFnPrimtive).Arity1IB)
	assert.NotNil(t, odd.(*AFnPrimtive).Arity1)
	assert.True(t, odd.Invoke_Arity1(1.0).(bool))
	assert.False(t, odd.(*AFnPrimtive).Arity1IB(2.0))
}

func Test_Protocols(t *testing.T) {
	symbol := Symbol.Invoke_Arity2("foo", "bar")

	assert.True(t, NativeSatisifes_QMARK_.Invoke_Arity2(Symbol.Invoke_Arity2("cljs.core", "INamed"), symbol).(bool))
	assert.True(t, NativeSatisifes_QMARK_.Invoke_Arity2(Symbol.Invoke_Arity2("cljs.core", "IFn"), symbol).(bool))
	assert.Equal(t, "foo", symbol.(INamed).Namespace_Arity1())
	assert.Equal(t, "foo", Namespace_.Invoke_Arity1(symbol))
	PanicsWith(t, "Invalid arity: 0", func() { symbol.(*CljsCoreSymbol).Invoke_Arity0() })
	PanicsWith(t, "Invalid arity: 0", func() { Invoke_.Invoke_Arity1(symbol) })

	assert.Equal(t, "foo", Namespace.Invoke_Arity1(symbol))
	PanicsWith(t, "Doesn't support namespace: 2", func() { Namespace.Invoke_Arity1(2) })

	assert.Equal(t, "bar", Name.Invoke_Arity1(symbol))
	PanicsWith(t, "Doesn't support name: 2", func() { Name.Invoke_Arity1(2) })
	assert.Equal(t, "baz", Name.Invoke_Arity1("baz"))

	assert.Equal(t, "foo", Invoke_.Invoke_Arity2(Namespace_, symbol))
	assert.Equal(t, "bar", symbol.(INamed).Name_Arity1())
	assert.Equal(t, "foo/bar", symbol.(fmt.Stringer).String())

	foo, bar := Symbol.Invoke_Arity1("foo"), Symbol.Invoke_Arity1("bar")
	m := ObjMap(map[interface{}]interface{}{foo: "bar"})

	assert.True(t, NativeSatisifes_QMARK_.Invoke_Arity2(Symbol.Invoke_Arity2("cljs.core", "ILookup"), m).(bool))
	assert.Equal(t, "bar", m.Lookup_Arity2(foo))
	assert.Nil(t, Lookup_.(*AFn).Call(m, bar))
	assert.Equal(t, "baz", m.Lookup_Arity3(bar, "baz"))

	assert.Equal(t, "bar", foo.(IFn).Invoke_Arity1(m))
	assert.Equal(t, "bar", Invoke_.Invoke_Arity2(foo, m))
	assert.Nil(t, bar.(IFn).Invoke_Arity1(m))
	assert.Equal(t, "baz", bar.(IFn).Invoke_Arity2(m, "baz"))
}

func Test_InteropViaReflection(t *testing.T) {
	sb := &gstring.StringBuffer{js.JSString("foo")}
	assert.Equal(t, "foo", NativeGetInstanceField.Invoke_Arity2(sb, "Buffer"),
		"(.-buffer sb)")
	NativeSetInstanceField.Invoke_Arity3(sb, "Buffer", "bar")
	assert.Equal(t, "bar", sb.Buffer,
		"(set! (.-buffer sb) \"bar\")")
	assert.Equal(t, sb, NativeInvokeInstanceMethod.Invoke_Arity3(sb, "Append", []interface{}{"baz"}))
	assert.Equal(t, "barbaz", sb.Buffer,
		"(.append sb \"baz\")")

	assert.Equal(t, js.Number.MAX_VALUE, NativeGetInstanceField.Invoke_Arity2(js.Number, "MAX_VALUE"),
		"(.-MAX-VALUE js/Number)")

	assert.Equal(t, 3, NativeInvokeFunc.Invoke_Arity2(Math.Floor, []interface{}{3.14}),
		"(Math/floor 3.14)")
	assert.Equal(t, "3.14", NativeInvokeFunc.Invoke_Arity2(fmt.Sprint, []interface{}{3.14}),
		"(fmt/Sprint 3.14)")
	assert.Equal(t, 3.14, NativeInvokeFunc.Invoke_Arity2(js.ParseFloat, []interface{}{js.JSString("3.14")}),
		"(js/parseFloat \"3.14\")")

	assert.Equal(t, "ABC", NativeInvokeFunc.Invoke_Arity2(js.String.FromCharCode, []interface{}{65.0, 66.0, 67.0}),
		"(.fromCharCode js/String 65 66 67)")
}

func PanicsWith(t *testing.T, message string, f assert.PanicTestFunc) {
	assert.Equal(t, message, func() (message string) {
		defer func() { message = fmt.Sprint(recover()) }()
		f()
		assert.Fail(t, "should panic")
		return
	}())
}
<|endoftext|>


<commit_before>package main

import (
	"./arch"
	"./loader"
)

type Usercorn struct {
	*Unicorn
	loader    loader.Loader
	Entry     uint64
	OS        string
	StackBase uint64
}

func NewUsercorn(exe string) (*Usercorn, error) {
	l, err := loader.LoadFile(exe)
	if err != nil {
		return nil, err
	}
	a, err := arch.GetArch(l.Arch(), l.OS())
	if err != nil {
		return nil, err
	}
	uc, err := NewUnicorn(a)
	if err != nil {
		return nil, err
	}
	u := &Usercorn{
		Unicorn: uc,
		loader:  l,
		OS:      l.OS(),
		Entry:   l.Entry(),
	}
	return u, nil
}

func (u *Usercorn) Run(args ...string) error {
	if err := u.mapMemory(); err != nil {
		return err
	}
	if err := u.setupStack(); err != nil {
		return err
	}
	// envp
	u.Push(0)
	// argv
	if err := u.pushStrings(args...); err != nil {
		return err
	}
	// argc
	u.Push(uint64(len(args)))
	return u.Uc.Start(u.Entry, 0)
}

func (u *Usercorn) mapMemory() error {
	segments, err := u.loader.Segments()
	if err != nil {
		return err
	}
	for _, seg := range segments {
		if err := u.MemMap(seg.Addr, uint64(len(seg.Data))); err != nil {
			return err
		}
		if err := u.MemWrite(seg.Addr, seg.Data); err != nil {
			return err
		}
	}
	return nil
}

func (u *Usercorn) setupStack() error {
	stack, err := u.Mmap(STACK_BASE, STACK_SIZE)
	if err != nil {
		return err
	}
	u.StackBase = stack
	if err := u.RegWrite(u.Arch.SP, stack+STACK_SIZE-uint64(u.Bsz)); err != nil {
		return err
	}
	return nil
}

func (u *Usercorn) pushStrings(args ...string) error {
	argvSize := 0
	for _, v := range args {
		argvSize += len(v) + 1
	}
	argvAddr, err := u.Mmap(0, uint64(argvSize))
	if err != nil {
		return err
	}
	buf := make([]byte, argvSize)
	addrs := make([]uint64, 0, len(args)+1)
	var pos uint64
	for i := len(args) - 1; i >= 0; i-- {
		copy(buf[pos:], []byte(args[i]))
		addrs = append(addrs, argvAddr+pos)
		pos += uint64(len(args[i]) + 1)
	}
	u.Push(0)
	for _, v := range addrs {
		u.Push(v)
	}
	return nil
}
<commit_msg>place SP more correctly<commit_after>package main

import (
	"./arch"
	"./loader"
)

type Usercorn struct {
	*Unicorn
	loader    loader.Loader
	Entry     uint64
	OS        string
	StackBase uint64
}

func NewUsercorn(exe string) (*Usercorn, error) {
	l, err := loader.LoadFile(exe)
	if err != nil {
		return nil, err
	}
	a, err := arch.GetArch(l.Arch(), l.OS())
	if err != nil {
		return nil, err
	}
	uc, err := NewUnicorn(a)
	if err != nil {
		return nil, err
	}
	u := &Usercorn{
		Unicorn: uc,
		loader:  l,
		OS:      l.OS(),
		Entry:   l.Entry(),
	}
	return u, nil
}

func (u *Usercorn) Run(args ...string) error {
	if err := u.mapMemory(); err != nil {
		return err
	}
	if err := u.setupStack(); err != nil {
		return err
	}
	// envp
	u.Push(0)
	// argv
	if err := u.pushStrings(args...); err != nil {
		return err
	}
	// argc
	u.Push(uint64(len(args)))
	return u.Uc.Start(u.Entry, 0)
}

func (u *Usercorn) mapMemory() error {
	segments, err := u.loader.Segments()
	if err != nil {
		return err
	}
	for _, seg := range segments {
		if err := u.MemMap(seg.Addr, uint64(len(seg.Data))); err != nil {
			return err
		}
		if err := u.MemWrite(seg.Addr, seg.Data); err != nil {
			return err
		}
	}
	return nil
}

func (u *Usercorn) setupStack() error {
	stack, err := u.Mmap(STACK_BASE, STACK_SIZE)
	if err != nil {
		return err
	}
	u.StackBase = stack
	if err := u.RegWrite(u.Arch.SP, stack+STACK_SIZE); err != nil {
		return err
	}
	return nil
}

func (u *Usercorn) pushStrings(args ...string) error {
	argvSize := 0
	for _, v := range args {
		argvSize += len(v) + 1
	}
	argvAddr, err := u.Mmap(0, uint64(argvSize))
	if err != nil {
		return err
	}
	buf := make([]byte, argvSize)
	addrs := make([]uint64, 0, len(args)+1)
	var pos uint64
	for i := len(args) - 1; i >= 0; i-- {
		copy(buf[pos:], []byte(args[i]))
		addrs = append(addrs, argvAddr+pos)
		pos += uint64(len(args[i]) + 1)
	}
	u.Push(0)
	for _, v := range addrs {
		u.Push(v)
	}
	return nil
}
<|endoftext|>


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

// +build android

package util

import (
	"fmt"
	"log"
	"reflect"
	"runtime/debug"
	"sync"
	"unsafe"
)

// #include "jni_wrapper.h"
import "C"

// NewGlobalRef creates a new global reference to the object referred to by the
// obj argument.  The obj argument may be a global or local reference. Global
// references must be explicitly disposed of by calling DeleteGlobalRef().
// NOTE: Because CGO creates package-local types and because this method may be
// invoked from a different package, Java types are passed in an empty interface
// and then cast into their package local types.
func NewGlobalRef(env, obj interface{}) unsafe.Pointer {
	jEnv := getEnv(env)
	jObj := getObject(obj)
	return unsafe.Pointer(C.NewGlobalRef(jEnv, jObj))
}

// DeleteGlobalRef deletes the global reference pointed to by obj.
// NOTE: Because CGO creates package-local types and because this method may be
// invoked from a different package, Java types are passed in an empty interface
// and then cast into their package local types.
func DeleteGlobalRef(env, obj interface{}) {
	jEnv := getEnv(env)
	jObj := getObject(obj)
	C.DeleteGlobalRef(jEnv, jObj)
}

// Creates a new local reference that refers to the same object as obj. The
// given obj may be a global or local reference. Returns null if ref refers
// to null.
// NOTE: Because CGO creates package-local types and because this method may be
// invoked from a different package, Java types are passed in an empty interface
// and then cast into their package local types.
func NewLocalRef(env, obj interface{}) unsafe.Pointer {
	jEnv := getEnv(env)
	jObj := getObject(obj)
	return unsafe.Pointer(C.NewLocalRef(jEnv, jObj))
}

// IsNull returns true iff the provided object is not null.
// NOTE: Because CGO creates package-local types and because this method may be
// invoked from a different package, Java types are passed in an empty interface
// and then cast into their package local types.
func IsNull(obj interface{}) bool {
	return getObject(obj) == nil
}

// GoRef creates a new reference to the value addressed by the provided pointer.
// The value will remain referenced until it is explicitly unreferenced using
// goUnref().
func GoRef(valptr interface{}) {
	if !IsPointer(valptr) {
		panic("must pass pointer value to goRef")
	}
	goRefs.ref(valptr)
}

// GoUnref removes a previously added reference to the value addressed by the
// provided pointer.  If the value hasn't been ref-ed (a bug?), this unref will
// be a no-op.
func GoUnref(valptr interface{}) {
	if !IsPointer(valptr) {
		panic("must pass pointer value to goUnref")
	}
	goRefs.unref(valptr)
}

// IsPointer returns true iff the provided value is a pointer.
func IsPointer(val interface{}) bool {
	if _, ok := val.(unsafe.Pointer); ok {
		return true
	}
	return reflect.ValueOf(val).Kind() == reflect.Ptr
}

// PtrValue returns the value of the pointer as a uintptr.
func PtrValue(ptr interface{}) uintptr {
	v := reflect.ValueOf(ptr)
	if v.Kind() != reflect.Ptr && v.Kind() != reflect.UnsafePointer {
		panic(fmt.Sprintf("must pass pointer value to PtrValue, was %v", v.Kind()))
	}
	return v.Pointer()
}

// DerefOrDie dereferences the provided (pointer) value, or panic-s if the value
// isn't of pointer type.
func DerefOrDie(i interface{}) interface{} {
	v := reflect.ValueOf(i)
	if v.Kind() != reflect.Ptr {
		panic(fmt.Sprintf("want reflect.Ptr value for %v, have %v", i, v.Kind()))
	}
	return v.Elem().Interface()
}

// Ptr returns the value of the provided Java pointer (of type C.jlong) as an
// unsafe.Pointer.
// NOTE: Because CGO creates package-local types and because this method may be
// invoked from a different package, Java types are passed in an empty interface
// and then cast into their package local types.
func Ptr(jPtr interface{}) unsafe.Pointer {
	v := reflect.ValueOf(jPtr)
	return unsafe.Pointer(uintptr(v.Int()))
}

// goRefs stores references to instances of various Go types, namely instances
// that are referenced only by the Java code.  The only purpose of this store
// is to prevent Go runtime from garbage collecting those instances.
var goRefs = newSafeRefCounter()

type refData struct {
	instance interface{}
	count    int
}

// newSafeRefCounter returns a new instance of a thread-safe reference counter.
func newSafeRefCounter() *safeRefCounter {
	return &safeRefCounter{
		refs: make(map[uintptr]refData),
	}
}

// safeRefCounter is a thread-safe reference counter.
type safeRefCounter struct {
	lock sync.Mutex
	refs map[uintptr]refData
}

// ref increases the reference count to the given valptr by 1.
func (c *safeRefCounter) ref(valptr interface{}) {
	p := PtrValue(valptr)
	c.lock.Lock()
	defer c.lock.Unlock()
	ref, ok := c.refs[p]
	if !ok {
		c.refs[p] = refData{
			instance: valptr,
			count:    1,
		}
	} else {
		ref.count++
	}
}

func (c *safeRefCounter) unref(valptr interface{}) {
	c.lock.Lock()
	defer c.lock.Unlock()
	p := PtrValue(valptr)
	ref, ok := c.refs[p]
	if !ok {
		log.Printf("Unrefing pointer %d of type %T that hasn't been refed before, stack: %s", int64(p), valptr, string(debug.Stack()))
		return
	}
	count := ref.count
	if count == 0 {
		log.Printf("Ref count for pointer %d of type %T is zero: that shouldn't happen, stack: %s", int64(p), valptr, string(debug.Stack()))
		return
	}
	if count > 1 {
		ref.count--
		return
	}
	delete(c.refs, p)
}
<commit_msg>util: fix incrementing count on a *copy* of refData Change-Id: Ife16b388365a48254ef02e34ef456c939c3b1ab1<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build android

package util

import (
	"fmt"
	"log"
	"reflect"
	"runtime/debug"
	"sync"
	"unsafe"
)

// #include "jni_wrapper.h"
import "C"

// NewGlobalRef creates a new global reference to the object referred to by the
// obj argument.  The obj argument may be a global or local reference. Global
// references must be explicitly disposed of by calling DeleteGlobalRef().
// NOTE: Because CGO creates package-local types and because this method may be
// invoked from a different package, Java types are passed in an empty interface
// and then cast into their package local types.
func NewGlobalRef(env, obj interface{}) unsafe.Pointer {
	jEnv := getEnv(env)
	jObj := getObject(obj)
	return unsafe.Pointer(C.NewGlobalRef(jEnv, jObj))
}

// DeleteGlobalRef deletes the global reference pointed to by obj.
// NOTE: Because CGO creates package-local types and because this method may be
// invoked from a different package, Java types are passed in an empty interface
// and then cast into their package local types.
func DeleteGlobalRef(env, obj interface{}) {
	jEnv := getEnv(env)
	jObj := getObject(obj)
	C.DeleteGlobalRef(jEnv, jObj)
}

// Creates a new local reference that refers to the same object as obj. The
// given obj may be a global or local reference. Returns null if ref refers
// to null.
// NOTE: Because CGO creates package-local types and because this method may be
// invoked from a different package, Java types are passed in an empty interface
// and then cast into their package local types.
func NewLocalRef(env, obj interface{}) unsafe.Pointer {
	jEnv := getEnv(env)
	jObj := getObject(obj)
	return unsafe.Pointer(C.NewLocalRef(jEnv, jObj))
}

// IsNull returns true iff the provided object is not null.
// NOTE: Because CGO creates package-local types and because this method may be
// invoked from a different package, Java types are passed in an empty interface
// and then cast into their package local types.
func IsNull(obj interface{}) bool {
	return getObject(obj) == nil
}

// GoRef creates a new reference to the value addressed by the provided pointer.
// The value will remain referenced until it is explicitly unreferenced using
// goUnref().
func GoRef(valptr interface{}) {
	if !IsPointer(valptr) {
		panic("must pass pointer value to goRef")
	}
	goRefs.ref(valptr)
}

// GoUnref removes a previously added reference to the value addressed by the
// provided pointer.  If the value hasn't been ref-ed (a bug?), this unref will
// be a no-op.
func GoUnref(valptr interface{}) {
	if !IsPointer(valptr) {
		panic("must pass pointer value to goUnref")
	}
	goRefs.unref(valptr)
}

// IsPointer returns true iff the provided value is a pointer.
func IsPointer(val interface{}) bool {
	if _, ok := val.(unsafe.Pointer); ok {
		return true
	}
	return reflect.ValueOf(val).Kind() == reflect.Ptr
}

// PtrValue returns the value of the pointer as a uintptr.
func PtrValue(ptr interface{}) uintptr {
	v := reflect.ValueOf(ptr)
	if v.Kind() != reflect.Ptr && v.Kind() != reflect.UnsafePointer {
		panic(fmt.Sprintf("must pass pointer value to PtrValue, was %v", v.Kind()))
	}
	return v.Pointer()
}

// DerefOrDie dereferences the provided (pointer) value, or panic-s if the value
// isn't of pointer type.
func DerefOrDie(i interface{}) interface{} {
	v := reflect.ValueOf(i)
	if v.Kind() != reflect.Ptr {
		panic(fmt.Sprintf("want reflect.Ptr value for %v, have %v", i, v.Kind()))
	}
	return v.Elem().Interface()
}

// Ptr returns the value of the provided Java pointer (of type C.jlong) as an
// unsafe.Pointer.
// NOTE: Because CGO creates package-local types and because this method may be
// invoked from a different package, Java types are passed in an empty interface
// and then cast into their package local types.
func Ptr(jPtr interface{}) unsafe.Pointer {
	v := reflect.ValueOf(jPtr)
	return unsafe.Pointer(uintptr(v.Int()))
}

// goRefs stores references to instances of various Go types, namely instances
// that are referenced only by the Java code.  The only purpose of this store
// is to prevent Go runtime from garbage collecting those instances.
var goRefs = newSafeRefCounter()

type refData struct {
	instance interface{}
	count    int
}

// newSafeRefCounter returns a new instance of a thread-safe reference counter.
func newSafeRefCounter() *safeRefCounter {
	return &safeRefCounter{
		refs: make(map[uintptr]*refData),
	}
}

// safeRefCounter is a thread-safe reference counter.
type safeRefCounter struct {
	lock sync.Mutex
	refs map[uintptr]*refData
}

// ref increases the reference count to the given valptr by 1.
func (c *safeRefCounter) ref(valptr interface{}) {
	p := PtrValue(valptr)
	c.lock.Lock()
	defer c.lock.Unlock()
	ref, ok := c.refs[p]
	if !ok {
		c.refs[p] = &refData{
			instance: valptr,
			count:    1,
		}
	} else {
		ref.count++
	}
}

func (c *safeRefCounter) unref(valptr interface{}) {
	c.lock.Lock()
	defer c.lock.Unlock()
	p := PtrValue(valptr)
	ref, ok := c.refs[p]
	if !ok {
		log.Printf("Unrefing pointer %d of type %T that hasn't been refed before, stack: %s", int64(p), valptr, string(debug.Stack()))
		return
	}
	count := ref.count
	if count == 0 {
		log.Printf("Ref count for pointer %d of type %T is zero: that shouldn't happen, stack: %s", int64(p), valptr, string(debug.Stack()))
		return
	}
	if count > 1 {
		ref.count--
		return
	}
	delete(c.refs, p)
}
<|endoftext|>


<commit_before>package notifierKafka

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

	"github.com/Shopify/sarama"
	part "github.com/raintank/metrictank/cluster/partitioner"
	"github.com/raintank/metrictank/kafka"
	"github.com/raintank/metrictank/stats"
	"github.com/raintank/worldping-api/pkg/log"
	"github.com/rakyll/globalconf"
)

var Enabled bool
var brokerStr string
var brokers []string
var topic string
var offsetStr string
var dataDir string
var config *sarama.Config
var offsetDuration time.Duration
var offsetCommitInterval time.Duration
var partitionStr string
var partitions []int32
var partitioner *part.Kafka
var partitionScheme string
var bootTimeOffsets map[int32]int64
var backlogProcessTimeout time.Duration
var backlogProcessTimeoutStr string
var partitionOffset map[int32]*stats.Gauge64
var partitionLogSize map[int32]*stats.Gauge64
var partitionLag map[int32]*stats.Gauge64

// metric cluster.notifier.kafka.messages-published is a counter of messages published to the kafka cluster notifier
var messagesPublished = stats.NewCounter32("cluster.notifier.kafka.messages-published")

// metric cluster.notifier.kafka.message_size is the sizes seen of messages through the kafka cluster notifier
var messagesSize = stats.NewMeter32("cluster.notifier.kafka.message_size", false)

func init() {
	fs := flag.NewFlagSet("kafka-cluster", flag.ExitOnError)
	fs.BoolVar(&Enabled, "enabled", false, "")
	fs.StringVar(&brokerStr, "brokers", "kafka:9092", "tcp address for kafka (may be given multiple times as comma separated list)")
	fs.StringVar(&topic, "topic", "metricpersist", "kafka topic")
	fs.StringVar(&partitionStr, "partitions", "*", "kafka partitions to consume. use '*' or a comma separated list of id's. This should match the partitions used for kafka-mdm-in")
	fs.StringVar(&partitionScheme, "partition-scheme", "bySeries", "method used for partitioning metrics. This should match the settings of tsdb-gw. (byOrg|bySeries)")
	fs.StringVar(&offsetStr, "offset", "last", "Set the offset to start consuming from. Can be one of newest, oldest,last or a time duration")
	fs.StringVar(&dataDir, "data-dir", "", "Directory to store partition offsets index")
	fs.DurationVar(&offsetCommitInterval, "offset-commit-interval", time.Second*5, "Interval at which offsets should be saved.")
	fs.StringVar(&backlogProcessTimeoutStr, "backlog-process-timeout", "60s", "Maximum time backlog processing can block during metrictank startup.")
	globalconf.Register("kafka-cluster", fs)
}

func ConfigProcess(instance string) {
	if !Enabled {
		return
	}
	var err error
	switch offsetStr {
	case "last":
	case "oldest":
	case "newest":
	default:
		offsetDuration, err = time.ParseDuration(offsetStr)
		if err != nil {
			log.Fatal(4, "kafka-cluster: invalid offest format. %s", err)
		}
	}
	brokers = strings.Split(brokerStr, ",")

	config = sarama.NewConfig()
	config.ClientID = instance + "-cluster"
	config.Version = sarama.V0_10_0_0
	config.Producer.RequiredAcks = sarama.WaitForAll // Wait for all in-sync replicas to ack the message
	config.Producer.Retry.Max = 10                   // Retry up to 10 times to produce the message
	config.Producer.Compression = sarama.CompressionSnappy
	config.Producer.Return.Successes = true
	err = config.Validate()
	if err != nil {
		log.Fatal(2, "kafka-cluster invalid consumer config: %s", err)
	}

	backlogProcessTimeout, err = time.ParseDuration(backlogProcessTimeoutStr)
	if err != nil {
		log.Fatal(4, "kafka-cluster: unable to parse backlog-process-timeout. %s", err)
	}

	partitioner, err = part.NewKafka(partitionScheme)
	if err != nil {
		log.Fatal(4, "kafka-cluster: failed to initialize partitioner. %s", err)
	}

	if partitionStr != "*" {
		parts := strings.Split(partitionStr, ",")
		for _, part := range parts {
			i, err := strconv.Atoi(part)
			if err != nil {
				log.Fatal(4, "kafka-cluster: could not parse partition %q. partitions must be '*' or a comma separated list of id's", part)
			}
			partitions = append(partitions, int32(i))
		}
	}
	// validate our partitions
	client, err := sarama.NewClient(brokers, config)
	if err != nil {
		log.Fatal(4, "kafka-cluster failed to create client. %s", err)
	}
	defer client.Close()

	availParts, err := kafka.GetPartitions(client, []string{topic})
	if err != nil {
		log.Fatal(4, "kafka-cluster: %s", err.Error())
	}
	if partitionStr == "*" {
		partitions = availParts
	} else {
		missing := kafka.DiffPartitions(partitions, availParts)
		if len(missing) > 0 {
			log.Fatal(4, "kafka-cluster: configured partitions not in list of available partitions. missing %v", missing)
		}
	}

	// initialize our offset metrics
	partitionOffset = make(map[int32]*stats.Gauge64)
	partitionLogSize = make(map[int32]*stats.Gauge64)
	partitionLag = make(map[int32]*stats.Gauge64)

	// get the "newest" offset for all partitions.
	// when booting up, we will delay consuming metrics until we have
	// caught up to these offsets.
	bootTimeOffsets = make(map[int32]int64)
	for _, part := range partitions {
		offset, err := client.GetOffset(topic, part, sarama.OffsetNewest)
		if err != nil {
			log.Fatal(4, "kakfa-cluster: failed to get newest offset for %s:%d. %s", topic, part)
		}
		bootTimeOffsets[part] = offset
		partitionOffset[part] = stats.NewGauge64(fmt.Sprintf("cluster.notifier.kafka.partition.%d.offset", part))
		partitionLogSize[part] = stats.NewGauge64(fmt.Sprintf("cluster.notifier.kafka.partition.%d.log_size", part))
		partitionLag[part] = stats.NewGauge64(fmt.Sprintf("cluster.notifier.kafka.partition.%d.lag", part))
	}
	log.Info("kafka-cluster: consuming from partitions %v", partitions)
}
<commit_msg>fix error fmt.<commit_after>package notifierKafka

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

	"github.com/Shopify/sarama"
	part "github.com/raintank/metrictank/cluster/partitioner"
	"github.com/raintank/metrictank/kafka"
	"github.com/raintank/metrictank/stats"
	"github.com/raintank/worldping-api/pkg/log"
	"github.com/rakyll/globalconf"
)

var Enabled bool
var brokerStr string
var brokers []string
var topic string
var offsetStr string
var dataDir string
var config *sarama.Config
var offsetDuration time.Duration
var offsetCommitInterval time.Duration
var partitionStr string
var partitions []int32
var partitioner *part.Kafka
var partitionScheme string
var bootTimeOffsets map[int32]int64
var backlogProcessTimeout time.Duration
var backlogProcessTimeoutStr string
var partitionOffset map[int32]*stats.Gauge64
var partitionLogSize map[int32]*stats.Gauge64
var partitionLag map[int32]*stats.Gauge64

// metric cluster.notifier.kafka.messages-published is a counter of messages published to the kafka cluster notifier
var messagesPublished = stats.NewCounter32("cluster.notifier.kafka.messages-published")

// metric cluster.notifier.kafka.message_size is the sizes seen of messages through the kafka cluster notifier
var messagesSize = stats.NewMeter32("cluster.notifier.kafka.message_size", false)

func init() {
	fs := flag.NewFlagSet("kafka-cluster", flag.ExitOnError)
	fs.BoolVar(&Enabled, "enabled", false, "")
	fs.StringVar(&brokerStr, "brokers", "kafka:9092", "tcp address for kafka (may be given multiple times as comma separated list)")
	fs.StringVar(&topic, "topic", "metricpersist", "kafka topic")
	fs.StringVar(&partitionStr, "partitions", "*", "kafka partitions to consume. use '*' or a comma separated list of id's. This should match the partitions used for kafka-mdm-in")
	fs.StringVar(&partitionScheme, "partition-scheme", "bySeries", "method used for partitioning metrics. This should match the settings of tsdb-gw. (byOrg|bySeries)")
	fs.StringVar(&offsetStr, "offset", "last", "Set the offset to start consuming from. Can be one of newest, oldest,last or a time duration")
	fs.StringVar(&dataDir, "data-dir", "", "Directory to store partition offsets index")
	fs.DurationVar(&offsetCommitInterval, "offset-commit-interval", time.Second*5, "Interval at which offsets should be saved.")
	fs.StringVar(&backlogProcessTimeoutStr, "backlog-process-timeout", "60s", "Maximum time backlog processing can block during metrictank startup.")
	globalconf.Register("kafka-cluster", fs)
}

func ConfigProcess(instance string) {
	if !Enabled {
		return
	}
	var err error
	switch offsetStr {
	case "last":
	case "oldest":
	case "newest":
	default:
		offsetDuration, err = time.ParseDuration(offsetStr)
		if err != nil {
			log.Fatal(4, "kafka-cluster: invalid offest format. %s", err)
		}
	}
	brokers = strings.Split(brokerStr, ",")

	config = sarama.NewConfig()
	config.ClientID = instance + "-cluster"
	config.Version = sarama.V0_10_0_0
	config.Producer.RequiredAcks = sarama.WaitForAll // Wait for all in-sync replicas to ack the message
	config.Producer.Retry.Max = 10                   // Retry up to 10 times to produce the message
	config.Producer.Compression = sarama.CompressionSnappy
	config.Producer.Return.Successes = true
	err = config.Validate()
	if err != nil {
		log.Fatal(2, "kafka-cluster invalid consumer config: %s", err)
	}

	backlogProcessTimeout, err = time.ParseDuration(backlogProcessTimeoutStr)
	if err != nil {
		log.Fatal(4, "kafka-cluster: unable to parse backlog-process-timeout. %s", err)
	}

	partitioner, err = part.NewKafka(partitionScheme)
	if err != nil {
		log.Fatal(4, "kafka-cluster: failed to initialize partitioner. %s", err)
	}

	if partitionStr != "*" {
		parts := strings.Split(partitionStr, ",")
		for _, part := range parts {
			i, err := strconv.Atoi(part)
			if err != nil {
				log.Fatal(4, "kafka-cluster: could not parse partition %q. partitions must be '*' or a comma separated list of id's", part)
			}
			partitions = append(partitions, int32(i))
		}
	}
	// validate our partitions
	client, err := sarama.NewClient(brokers, config)
	if err != nil {
		log.Fatal(4, "kafka-cluster failed to create client. %s", err)
	}
	defer client.Close()

	availParts, err := kafka.GetPartitions(client, []string{topic})
	if err != nil {
		log.Fatal(4, "kafka-cluster: %s", err.Error())
	}
	if partitionStr == "*" {
		partitions = availParts
	} else {
		missing := kafka.DiffPartitions(partitions, availParts)
		if len(missing) > 0 {
			log.Fatal(4, "kafka-cluster: configured partitions not in list of available partitions. missing %v", missing)
		}
	}

	// initialize our offset metrics
	partitionOffset = make(map[int32]*stats.Gauge64)
	partitionLogSize = make(map[int32]*stats.Gauge64)
	partitionLag = make(map[int32]*stats.Gauge64)

	// get the "newest" offset for all partitions.
	// when booting up, we will delay consuming metrics until we have
	// caught up to these offsets.
	bootTimeOffsets = make(map[int32]int64)
	for _, part := range partitions {
		offset, err := client.GetOffset(topic, part, sarama.OffsetNewest)
		if err != nil {
			log.Fatal(4, "kakfa-cluster: failed to get newest offset for topic %s part %d: %s", topic, part, err)
		}
		bootTimeOffsets[part] = offset
		partitionOffset[part] = stats.NewGauge64(fmt.Sprintf("cluster.notifier.kafka.partition.%d.offset", part))
		partitionLogSize[part] = stats.NewGauge64(fmt.Sprintf("cluster.notifier.kafka.partition.%d.log_size", part))
		partitionLag[part] = stats.NewGauge64(fmt.Sprintf("cluster.notifier.kafka.partition.%d.lag", part))
	}
	log.Info("kafka-cluster: consuming from partitions %v", partitions)
}
<|endoftext|>


<commit_before>package main

import (
	"os"

	"github.com/cbednarski/hostess"
	"github.com/codegangsta/cli"
)

func getCommand() string {
	return os.Args[1]
}

func getArgs() []string {
	return os.Args[2:]
}

const help = `an idempotent tool for managing /etc/hosts

 * Commands will exit 0 or 1 in a sensible way to facilitate scripting.
 * Hostess operates on /etc/hosts by default. Specify the HOSTESS_PATH
   environment variable to change this.
 * Run 'hostess fix -n' to preview changes hostess will make to your hostsfile.
 * Report bugs and feedback at https://github.com/cbednarski/hostess`

func main() {
	app := cli.NewApp()
	app.Name = "hostess"
	app.Authors = []cli.Author{{Name: "Chris Bednarski", Email: "banzaimonkey@gmail.com"}}
	app.Usage = help
	app.Version = "0.1.0"

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "f",
			Usage: "operate even if there are errors or conflicts",
		},
		cli.BoolFlag{
			Name:  "n",
			Usage: "no-op. Show changes but don't write them.",
		},
		cli.BoolFlag{
			Name:  "q",
			Usage: "quiet operation -- no notices",
		},
		cli.BoolFlag{
			Name:  "s",
			Usage: "silent operation -- no errors (implies -q)",
		},
	}

	app.Commands = []cli.Command{
		{
			Name:   "add",
			Usage:  "add or replace a hosts entry",
			Action: hostess.Add,
			Flags:  app.Flags,
		},
		{
			Name:   "aff",
			Usage:  "add or replace a hosts entry in an off state",
			Action: hostess.Add,
			Flags:  app.Flags,
		},
		{
			Name:    "del",
			Aliases: []string{"rm"},
			Usage:   "delete a hosts entry",
			Action:  hostess.Del,
			Flags:   app.Flags,
		},
		{
			Name:   "has",
			Usage:  "exit 0 if entry exists, 1 if not",
			Action: hostess.Has,
			Flags:  app.Flags,
		},
		{
			Name:   "on",
			Usage:  "enable a hosts entry (if if exists)",
			Action: hostess.OnOff,
			Flags:  app.Flags,
		},
		{
			Name:   "off",
			Usage:  "disable a hosts entry (don't delete it)",
			Action: hostess.OnOff,
			Flags:  app.Flags,
		},
		{
			Name:    "list",
			Aliases: []string{"ls"},
			Usage:   "list entries in the hosts file",
			Action:  hostess.Ls,
			Flags:   app.Flags,
		},
		{
			Name:   "fix",
			Usage:  "reformat the hosts file based on hostess' rules",
			Action: hostess.Fix,
			Flags:  app.Flags,
		},
		{
			Name:   "dump",
			Usage:  "dump the hosts file as JSON",
			Action: hostess.Dump,
			Flags:  app.Flags,
		},
		{
			Name:   "apply",
			Usage:  "add hostnames from a JSON file to the hosts file",
			Action: hostess.Apply,
			Flags:  app.Flags,
		},
	}

	app.Run(os.Args)
	os.Exit(0)
}
<commit_msg>Bump version in CLI output to match the current tag<commit_after>package main

import (
	"os"

	"github.com/cbednarski/hostess"
	"github.com/codegangsta/cli"
)

func getCommand() string {
	return os.Args[1]
}

func getArgs() []string {
	return os.Args[2:]
}

const help = `an idempotent tool for managing /etc/hosts

 * Commands will exit 0 or 1 in a sensible way to facilitate scripting.
 * Hostess operates on /etc/hosts by default. Specify the HOSTESS_PATH
   environment variable to change this.
 * Run 'hostess fix -n' to preview changes hostess will make to your hostsfile.
 * Report bugs and feedback at https://github.com/cbednarski/hostess`

func main() {
	app := cli.NewApp()
	app.Name = "hostess"
	app.Authors = []cli.Author{{Name: "Chris Bednarski", Email: "banzaimonkey@gmail.com"}}
	app.Usage = help
	app.Version = "0.2.1"

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "f",
			Usage: "operate even if there are errors or conflicts",
		},
		cli.BoolFlag{
			Name:  "n",
			Usage: "no-op. Show changes but don't write them.",
		},
		cli.BoolFlag{
			Name:  "q",
			Usage: "quiet operation -- no notices",
		},
		cli.BoolFlag{
			Name:  "s",
			Usage: "silent operation -- no errors (implies -q)",
		},
	}

	app.Commands = []cli.Command{
		{
			Name:   "add",
			Usage:  "add or replace a hosts entry",
			Action: hostess.Add,
			Flags:  app.Flags,
		},
		{
			Name:   "aff",
			Usage:  "add or replace a hosts entry in an off state",
			Action: hostess.Add,
			Flags:  app.Flags,
		},
		{
			Name:    "del",
			Aliases: []string{"rm"},
			Usage:   "delete a hosts entry",
			Action:  hostess.Del,
			Flags:   app.Flags,
		},
		{
			Name:   "has",
			Usage:  "exit 0 if entry exists, 1 if not",
			Action: hostess.Has,
			Flags:  app.Flags,
		},
		{
			Name:   "on",
			Usage:  "enable a hosts entry (if if exists)",
			Action: hostess.OnOff,
			Flags:  app.Flags,
		},
		{
			Name:   "off",
			Usage:  "disable a hosts entry (don't delete it)",
			Action: hostess.OnOff,
			Flags:  app.Flags,
		},
		{
			Name:    "list",
			Aliases: []string{"ls"},
			Usage:   "list entries in the hosts file",
			Action:  hostess.Ls,
			Flags:   app.Flags,
		},
		{
			Name:   "fix",
			Usage:  "reformat the hosts file based on hostess' rules",
			Action: hostess.Fix,
			Flags:  app.Flags,
		},
		{
			Name:   "dump",
			Usage:  "dump the hosts file as JSON",
			Action: hostess.Dump,
			Flags:  app.Flags,
		},
		{
			Name:   "apply",
			Usage:  "add hostnames from a JSON file to the hosts file",
			Action: hostess.Apply,
			Flags:  app.Flags,
		},
	}

	app.Run(os.Args)
	os.Exit(0)
}
<|endoftext|>


<commit_before>package vorbis

import (
	"bytes"
	"io"
	"time"

	"github.com/dhowden/tag"
	"github.com/jfreymuth/go-vorbis/ogg"
	"github.com/jfreymuth/go-vorbis/ogg/vorbis"
	"github.com/mjibson/moggio/codec"
)

func init() {
	codec.RegisterCodec("VORBIS", []string{"OggS"}, []string{"ogg"}, NewSongs, nil)
}

func NewSongs(rf codec.Reader) (codec.Songs, error) {
	s, err := NewSong(rf)
	if err != nil {
		return nil, err
	}
	return codec.Songs{codec.None: s}, nil
}

func NewSong(rf codec.Reader) (codec.Song, error) {
	f := &Vorbis{
		Reader: rf,
	}
	return f, nil
}

type Vorbis struct {
	Reader  codec.Reader
	r       io.ReadCloser
	v       *vorbis.Vorbis
	samples []float32
	info    *codec.SongInfo
}

func (v *Vorbis) Init() (sampleRate, channels int, err error) {
	if v.v == nil {
		r, _, err := v.Reader()
		if err != nil {
			return 0, 0, err
		}
		vr, err := vorbis.Open(r)
		if err != nil {
			r.Close()
			return 0, 0, err
		}
		v.r = r
		v.v = vr
	}
	return v.v.SampleRate(), v.v.Channels(), nil
}

func (v *Vorbis) Info() (info codec.SongInfo, err error) {
	if v.info != nil {
		return *v.info, nil
	}
	si, _, b, err := v.Reader.Metadata(tag.OGG)
	if err != nil {
		return
	}
	or := ogg.NewReader(bytes.NewReader(b))
	vr, err := vorbis.OpenOgg(or)
	if err != nil {
		return
	}
	l, err := or.Length()
	if err != nil {
		return
	}
	si.Time = time.Duration(l/uint64(vr.SampleRate())) * time.Second
	v.info = si
	return *si, nil
}

func (v *Vorbis) Play(n int) ([]float32, error) {
	var err error
	for len(v.samples) < n && err == nil {
		samples, err := v.v.DecodePacket()
		if err != nil {
			break
		}
		n := len(samples[0])
		c := len(samples)
		data := make([]float32, c*n)
		for i, cs := range samples {
			for j, s := range cs {
				data[j*c+i] = s
			}
		}
		v.samples = append(v.samples, data...)
	}
	if n > len(v.samples) {
		n = len(v.samples)
	}
	ret := v.samples[:n]
	v.samples = v.samples[n:]
	return ret, err
}

func (v *Vorbis) Close() {
	if v.r != nil {
		v.r.Close()
		v.r = nil
	}
	v.v = nil
}
<commit_msg>Don't hide io.EOF error<commit_after>package vorbis

import (
	"bytes"
	"io"
	"time"

	"github.com/dhowden/tag"
	"github.com/jfreymuth/go-vorbis/ogg"
	"github.com/jfreymuth/go-vorbis/ogg/vorbis"
	"github.com/mjibson/moggio/codec"
)

func init() {
	codec.RegisterCodec("VORBIS", []string{"OggS"}, []string{"ogg"}, NewSongs, nil)
}

func NewSongs(rf codec.Reader) (codec.Songs, error) {
	s, err := NewSong(rf)
	if err != nil {
		return nil, err
	}
	return codec.Songs{codec.None: s}, nil
}

func NewSong(rf codec.Reader) (codec.Song, error) {
	f := &Vorbis{
		Reader: rf,
	}
	return f, nil
}

type Vorbis struct {
	Reader  codec.Reader
	r       io.ReadCloser
	v       *vorbis.Vorbis
	samples []float32
	info    *codec.SongInfo
}

func (v *Vorbis) Init() (sampleRate, channels int, err error) {
	if v.v == nil {
		r, _, err := v.Reader()
		if err != nil {
			return 0, 0, err
		}
		vr, err := vorbis.Open(r)
		if err != nil {
			r.Close()
			return 0, 0, err
		}
		v.r = r
		v.v = vr
	}
	return v.v.SampleRate(), v.v.Channels(), nil
}

func (v *Vorbis) Info() (info codec.SongInfo, err error) {
	if v.info != nil {
		return *v.info, nil
	}
	si, _, b, err := v.Reader.Metadata(tag.OGG)
	if err != nil {
		return
	}
	or := ogg.NewReader(bytes.NewReader(b))
	vr, err := vorbis.OpenOgg(or)
	if err != nil {
		return
	}
	l, err := or.Length()
	if err != nil {
		return
	}
	si.Time = time.Duration(l/uint64(vr.SampleRate())) * time.Second
	v.info = si
	return *si, nil
}

func (v *Vorbis) Play(n int) ([]float32, error) {
	var err error
	var samples [][]float32
	for len(v.samples) < n && err == nil {
		samples, err = v.v.DecodePacket()
		if len(samples) == 0 {
			break
		}
		n := len(samples[0])
		c := len(samples)
		data := make([]float32, c*n)
		for i, cs := range samples {
			for j, s := range cs {
				data[j*c+i] = s
			}
		}
		v.samples = append(v.samples, data...)
	}
	if n > len(v.samples) {
		n = len(v.samples)
	}
	ret := v.samples[:n]
	v.samples = v.samples[n:]
	return ret, err
}

func (v *Vorbis) Close() {
	if v.r != nil {
		v.r.Close()
		v.r = nil
	}
	v.v = nil
}
<|endoftext|>


<commit_before>package unit

import (
	"strconv"
	"sync"
)

// Uniter is an interface representing a type that can be converted
// to a unit.
type Uniter interface {
	Unit() *Unit
}

// Dimension is a type representing an SI base dimension or other
// orthogonal dimension. If a new dimension is desired for a
// domain-specific problem, NewDimension should be called.
type Dimension int

const (
	// SI Base Units
	CurrentDim Dimension = iota
	LengthDim
	LuminousIntensityDim
	MassDim
	TemperatureDim
	TimeDim
	// Start of other SI Units
	AngleDim             // e.g. radians
	lastPackageDimension // Used in create dimension
)

//TODO: Should there be some number reserved? We don't want users ever using integer literals
var lastCreatedDimension Dimension = lastPackageDimension

var dimensionToSymbol map[Dimension]string = make(map[Dimension]string)
var symbolToDimension map[string]Dimension = make(map[string]Dimension)

func init() {
	dimensionToSymbol[CurrentDim] = "A"
	symbolToDimension["A"] = CurrentDim
	dimensionToSymbol[LengthDim] = "m"
	symbolToDimension["m"] = LengthDim
	dimensionToSymbol[LuminousIntensityDim] = "cd"
	symbolToDimension["cd"] = LuminousIntensityDim
	dimensionToSymbol[MassDim] = "kg"
	symbolToDimension["kg"] = MassDim
	dimensionToSymbol[TemperatureDim] = "K"
	symbolToDimension["K"] = TemperatureDim
	dimensionToSymbol[TimeDim] = "s"
	symbolToDimension["s"] = TimeDim
	dimensionToSymbol[AngleDim] = "rad"
	symbolToDimension["rad"] = AngleDim
}

// Dimensions represent the dimensionality of the unit in powers
// of that dimension. If a key is not present, the power of that
// dimension is zero. Dimensions is used in conjuction with NewUnit
type Dimensions map[Dimension]int

var newUnitMutex *sync.Mutex = &sync.Mutex{} // so there is no race condition for dimension

// NewDimension returns a new dimension variable which will have a
// unique representation across packages to prevent accidental overlap.
// NewDimension should only be called for unit types that are orthogonal
// to the base dimensions defined in this package. For example, one unit
// that comes up in blood work is "White blood cells per microscope slide".
// NewDimension is appropriate for "White blood cells", as they are not
// representable in SI base units. However, NewDimension is not appropriate
// for "Slide", as slide is really a unit of area. Slide should instead be
// defined as a constant of type unit.Area
//  Expecting to be used only during initialization, it panics if the mapping between types and names is not a bijection.
func NewDimension(symbol string) Dimension {
	newUnitMutex.Lock()
	defer newUnitMutex.Unlock()
	lastCreatedDimension++
	_, ok := symbolToDimension[symbol]
	if ok {
		panic("unit: dimension string " + symbol + " already used")
	}
	dimensionToSymbol[lastCreatedDimension] = symbol
	symbolToDimension[symbol] = lastCreatedDimension
	return lastCreatedDimension
}

// Unit is a type a value with generic SI units. Most useful for
// translating between dimensions, for example, by multiplying
// an acceleration with a mass to get a force
type Unit struct {
	dimensions map[Dimension]int // Map for custom dimensions
	value      float64
}

// NewUnit creates a new variable of type Unit which has the value
// specified by value and the dimensions specified by the
// base units struct. The value is always in SI Units.
//
// Example: To create an acceleration of 3 m/s^2, one could do
// myvar := CreateUnit(3.0, &Dimensions{length: 1, time: -2})
func NewUnit(value float64, d Dimensions) *Unit {

	// TODO: Find most efficient way of doing this
	// I think copy is necessary in case the input
	// dimension map is changed later
	u := &Unit{
		dimensions: make(map[Dimension]int),
	}
	for key, val := range d {
		u.dimensions[key] = val
	}
	u.value = value
	return u
}

// DimensionsMatch checks if the dimensions of two Uniters are the same
func DimensionsMatch(aU, bU Uniter) bool {
	a := aU.Unit()
	b := bU.Unit()
	if len(a.dimensions) != len(b.dimensions) {
		return false
	}
	for key, val := range a.dimensions {
		if b.dimensions[key] != val {
			return false
		}
	}
	return true
}

// Add adds the function argument to the reciever. Panics if the units of
// the receiver and the argument don't match.
func (u *Unit) Add(aU Uniter) *Unit {
	a := aU.Unit()
	if !DimensionsMatch(u, a) {
		panic("Attempted to add the values of two units whose dimensions do not match.")
	}
	u.value += a.value
	return u
}

// Unit allows unit to satisfy the uniter interface
func (u *Unit) Unit() *Unit {
	return u
}

// Mul multiply the receiver by the unit changing the dimensions
// of the receiver as appropriate
func (u *Unit) Mul(aU Uniter) *Unit {
	a := aU.Unit()
	for key, val := range a.dimensions {
		u.dimensions[key] += val
	}
	u.value *= a.value
	return u
}

// Div divides the receiver by the argument changing the
// dimensions of the receiver as appropriate
func (u *Unit) Div(aU Uniter) *Unit {
	a := aU.Unit()
	u.value /= a.value
	for key, val := range a.dimensions {
		u.dimensions[key] -= val
	}
	return u
}

// Value return the raw value of the unit as a float64. Use of this
// method is not recommended, instead it is recommended to use a
// FromUnit type of a specific dimension
func (u *Unit) Value() float64 {
	return u.value
}

func (u Unit) String() string {
	str := strconv.FormatFloat(u.value, 'e', -1, 64)
	for dimension, power := range u.dimensions {
		if power != 0 {
			str += " " + dimensionToSymbol[dimension]
			if power != 1 {
				str += "^" + strconv.Itoa(power)
			}
		}
	}
	return str
}
<commit_msg>Made symbols print in alphabetical order<commit_after>package unit

import (
	"sort"
	"strconv"
	"sync"
)

// Uniter is an interface representing a type that can be converted
// to a unit.
type Uniter interface {
	Unit() *Unit
}

// Dimension is a type representing an SI base dimension or other
// orthogonal dimension. If a new dimension is desired for a
// domain-specific problem, NewDimension should be called.
type Dimension int

const (
	// SI Base Units
	CurrentDim Dimension = iota
	LengthDim
	LuminousIntensityDim
	MassDim
	TemperatureDim
	TimeDim
	// Start of other SI Units
	AngleDim             // e.g. radians
	lastPackageDimension // Used in create dimension
)

//TODO: Should there be some number reserved? We don't want users ever using integer literals
var lastCreatedDimension Dimension = lastPackageDimension

var dimensionToSymbol map[Dimension]string = make(map[Dimension]string)
var symbolToDimension map[string]Dimension = make(map[string]Dimension)

func init() {
	dimensionToSymbol[CurrentDim] = "A"
	symbolToDimension["A"] = CurrentDim
	dimensionToSymbol[LengthDim] = "m"
	symbolToDimension["m"] = LengthDim
	dimensionToSymbol[LuminousIntensityDim] = "cd"
	symbolToDimension["cd"] = LuminousIntensityDim
	dimensionToSymbol[MassDim] = "kg"
	symbolToDimension["kg"] = MassDim
	dimensionToSymbol[TemperatureDim] = "K"
	symbolToDimension["K"] = TemperatureDim
	dimensionToSymbol[TimeDim] = "s"
	symbolToDimension["s"] = TimeDim
	dimensionToSymbol[AngleDim] = "rad"
	symbolToDimension["rad"] = AngleDim
}

// Dimensions represent the dimensionality of the unit in powers
// of that dimension. If a key is not present, the power of that
// dimension is zero. Dimensions is used in conjuction with NewUnit
type Dimensions map[Dimension]int

var newUnitMutex *sync.Mutex = &sync.Mutex{} // so there is no race condition for dimension

// NewDimension returns a new dimension variable which will have a
// unique representation across packages to prevent accidental overlap.
// NewDimension should only be called for unit types that are orthogonal
// to the base dimensions defined in this package. For example, one unit
// that comes up in blood work is "White blood cells per microscope slide".
// NewDimension is appropriate for "White blood cells", as they are not
// representable in SI base units. However, NewDimension is not appropriate
// for "Slide", as slide is really a unit of area. Slide should instead be
// defined as a constant of type unit.Area
//  Expecting to be used only during initialization, it panics if the mapping between types and names is not a bijection.
func NewDimension(symbol string) Dimension {
	newUnitMutex.Lock()
	defer newUnitMutex.Unlock()
	lastCreatedDimension++
	_, ok := symbolToDimension[symbol]
	if ok {
		panic("unit: dimension string " + symbol + " already used")
	}
	dimensionToSymbol[lastCreatedDimension] = symbol
	symbolToDimension[symbol] = lastCreatedDimension
	return lastCreatedDimension
}

// Unit is a type a value with generic SI units. Most useful for
// translating between dimensions, for example, by multiplying
// an acceleration with a mass to get a force
type Unit struct {
	dimensions map[Dimension]int // Map for custom dimensions
	value      float64
}

// NewUnit creates a new variable of type Unit which has the value
// specified by value and the dimensions specified by the
// base units struct. The value is always in SI Units.
//
// Example: To create an acceleration of 3 m/s^2, one could do
// myvar := CreateUnit(3.0, &Dimensions{length: 1, time: -2})
func NewUnit(value float64, d Dimensions) *Unit {

	// TODO: Find most efficient way of doing this
	// I think copy is necessary in case the input
	// dimension map is changed later
	u := &Unit{
		dimensions: make(map[Dimension]int),
	}
	for key, val := range d {
		u.dimensions[key] = val
	}
	u.value = value
	return u
}

// DimensionsMatch checks if the dimensions of two Uniters are the same
func DimensionsMatch(aU, bU Uniter) bool {
	a := aU.Unit()
	b := bU.Unit()
	if len(a.dimensions) != len(b.dimensions) {
		return false
	}
	for key, val := range a.dimensions {
		if b.dimensions[key] != val {
			return false
		}
	}
	return true
}

// Add adds the function argument to the reciever. Panics if the units of
// the receiver and the argument don't match.
func (u *Unit) Add(aU Uniter) *Unit {
	a := aU.Unit()
	if !DimensionsMatch(u, a) {
		panic("Attempted to add the values of two units whose dimensions do not match.")
	}
	u.value += a.value
	return u
}

// Unit allows unit to satisfy the uniter interface
func (u *Unit) Unit() *Unit {
	return u
}

// Mul multiply the receiver by the unit changing the dimensions
// of the receiver as appropriate
func (u *Unit) Mul(aU Uniter) *Unit {
	a := aU.Unit()
	for key, val := range a.dimensions {
		u.dimensions[key] += val
	}
	u.value *= a.value
	return u
}

// Div divides the receiver by the argument changing the
// dimensions of the receiver as appropriate
func (u *Unit) Div(aU Uniter) *Unit {
	a := aU.Unit()
	u.value /= a.value
	for key, val := range a.dimensions {
		u.dimensions[key] -= val
	}
	return u
}

// Value return the raw value of the unit as a float64. Use of this
// method is not recommended, instead it is recommended to use a
// FromUnit type of a specific dimension
func (u *Unit) Value() float64 {
	return u.value
}

type symbolString struct {
	symbol string
	pow    int
}

type unitPrinters []symbolString

func (u unitPrinters) Len() int {
	return len(u)
}

func (u unitPrinters) Less(i, j int) bool {
	return u[i].symbol < u[j].symbol
}

func (u unitPrinters) Swap(i, j int) {
	u[i], u[j] = u[j], u[i]
}

func (u Unit) String() string {
	str := strconv.FormatFloat(u.value, 'e', -1, 64)
	// Map iterates randomly, but we should output the symbols in a logical order
	data := make(unitPrinters, 0, 10)
	for dimension, power := range u.dimensions {
		if power != 0 {
			data = append(data, symbolString{dimensionToSymbol[dimension], power})
		}
	}
	sort.Sort(data)
	for _, s := range data {
		str += " " + s.symbol
		if s.pow != 1 {
			str += "^" + strconv.Itoa(s.pow)
		}
	}
	/*
		for dimension, power := range u.dimensions {
			if power != 0 {
				str += " " + dimensionToSymbol[dimension]
				if power != 1 {
					str += "^" + strconv.Itoa(power)
				}
			}
		}
	*/
	return str
}
<|endoftext|>


<commit_before>// meek-client is the client transport plugin for the meek pluggable transport.
//
// Sample usage in torrc:
// 	Bridge meek 0.0.2.0:1
// 	ClientTransportPlugin meek exec ./meek-client --url=https://meek-reflect.appspot.com/ --front=www.google.com --log meek-client.log
// The transport ignores the bridge address 0.0.2.0:1 and instead connects to
// the URL given by --url. When --front is given, the domain in the URL is
// replaced by the front domain for the purpose of the DNS lookup, TCP
// connection, and TLS SNI, but the HTTP Host header in the request will be the
// one in --url. (For example, in the configuration above, the connection will
// appear on the outside to be going to www.google.com, but it will actually be
// dispatched to meek-reflect.appspot.com by the Google frontend server.)
//
// Most user configuration can happen either through SOCKS args (i.e., args on a
// Bridge line) or through command line options. SOCKS args take precedence
// per-connection over command line options. For example, this configuration
// using SOCKS args:
// 	Bridge meek 0.0.2.0:1 url=https://meek-reflect.appspot.com/ front=www.google.com
// 	ClientTransportPlugin meek exec ./meek-client
// is the same as this one using command line options.
// 	Bridge meek 0.0.2.0:1
// 	ClientTransportPlugin meek exec ./meek-client --url=https://meek-reflect.appspot.com/ --front=www.google.com
// The advantage of SOCKS args is that multiple Bridge lines can have different
// configurations, but it requires a newer tor.
//
// The --helper option prevents this program from doing any network operations
// itself. Rather, it will send all requests through a browser extension that
// makes HTTP requests.
package main

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"syscall"
	"time"
)

import "git.torproject.org/pluggable-transports/goptlib.git"

const (
	ptMethodName = "meek"
	// A session ID is a randomly generated string that identifies a
	// long-lived session. We split a TCP stream across multiple HTTP
	// requests, and those with the same session ID belong to the same
	// stream.
	sessionIdLength = 32
	// The size of the largest chunk of data we will read from the SOCKS
	// port before forwarding it in a request, and the maximum size of a
	// body we are willing to handle in a reply.
	maxPayloadLength = 0x10000
	// We must poll the server to see if it has anything to send; there is
	// no way for the server to push data back to us until we send an HTTP
	// request. When a timer expires, we send a request even if it has an
	// empty body. The interval starts at this value and then grows.
	initPollInterval = 100 * time.Millisecond
	// Maximum polling interval.
	maxPollInterval = 5 * time.Second
	// Geometric increase in the polling interval each time we fail to read
	// data.
	pollIntervalMultiplier = 1.5
	// Safety limits on interaction with the HTTP helper.
	maxHelperResponseLength = 10000000
	helperReadTimeout       = 60 * time.Second
	helperWriteTimeout      = 2 * time.Second
)

var ptInfo pt.ClientInfo

// Store for command line options.
var options struct {
	URL          string
	Front        string
	HTTPProxyURL *url.URL
	HelperAddr   *net.TCPAddr
}

// When a connection handler starts, +1 is written to this channel; when it
// ends, -1 is written.
var handlerChan = make(chan int)

// RequestInfo encapsulates all the configuration used for a request–response
// roundtrip, including variables that may come from SOCKS args or from the
// command line.
type RequestInfo struct {
	// What to put in the X-Session-ID header.
	SessionID string
	// The URL to request.
	URL *url.URL
	// The Host header to put in the HTTP request (optional and may be
	// different from the host name in URL).
	Host string
	// URL of an HTTP proxy to use. If nil, the default net/http library's
	// behavior is used, which is to check the HTTP_PROXY and http_proxy
	// environment for a proxy URL.
	HTTPProxyURL *url.URL
}

// Do an HTTP roundtrip using the payload data in buf and the request metadata
// in info.
func roundTripWithHTTP(buf []byte, info *RequestInfo) (*http.Response, error) {
	tr := http.DefaultTransport
	if info.HTTPProxyURL != nil {
		tr = &http.Transport{
			Proxy: http.ProxyURL(info.HTTPProxyURL),
		}
	}
	req, err := http.NewRequest("POST", info.URL.String(), bytes.NewReader(buf))
	if err != nil {
		return nil, err
	}
	if info.Host != "" {
		req.Host = info.Host
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	req.Header.Set("X-Session-Id", info.SessionID)
	return tr.RoundTrip(req)
}

// Send the data in buf to the remote URL, wait for a reply, and feed the reply
// body back into conn.
func sendRecv(buf []byte, conn net.Conn, info *RequestInfo) (int64, error) {
	roundTrip := roundTripWithHTTP
	if options.HelperAddr != nil {
		roundTrip = roundTripWithHelper
	}
	resp, err := roundTrip(buf, info)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return 0, errors.New(fmt.Sprintf("status code was %d, not %d", resp.StatusCode, http.StatusOK))
	}

	return io.Copy(conn, io.LimitReader(resp.Body, maxPayloadLength))
}

// Repeatedly read from conn, issue HTTP requests, and write the responses back
// to conn.
func copyLoop(conn net.Conn, info *RequestInfo) error {
	var interval time.Duration

	ch := make(chan []byte)

	// Read from the Conn and send byte slices on the channel.
	go func() {
		var buf [maxPayloadLength]byte
		r := bufio.NewReader(conn)
		for {
			n, err := r.Read(buf[:])
			b := make([]byte, n)
			copy(b, buf[:n])
			// log.Printf("read from local: %q", b)
			ch <- b
			if err != nil {
				log.Printf("error reading from local: %s", err)
				break
			}
		}
		close(ch)
	}()

	interval = initPollInterval
loop:
	for {
		var buf []byte
		var ok bool

		// log.Printf("waiting up to %.2f s", interval.Seconds())
		// start := time.Now()
		select {
		case buf, ok = <-ch:
			if !ok {
				break loop
			}
			// log.Printf("read %d bytes from local after %.2f s", len(buf), time.Since(start).Seconds())
		case <-time.After(interval):
			// log.Printf("read nothing from local after %.2f s", time.Since(start).Seconds())
			buf = nil
		}

		nw, err := sendRecv(buf, conn, info)
		if err != nil {
			return err
		}
		/*
			if nw > 0 {
				log.Printf("got %d bytes from remote", nw)
			} else {
				log.Printf("got nothing from remote")
			}
		*/

		if nw > 0 || len(buf) > 0 {
			// If we sent or received anything, poll again
			// immediately.
			interval = 0
		} else if interval == 0 {
			// The first time we don't send or receive anything,
			// wait a while.
			interval = initPollInterval
		} else {
			// After that, wait a little longer.
			interval = time.Duration(float64(interval) * pollIntervalMultiplier)
		}
		if interval > maxPollInterval {
			interval = maxPollInterval
		}
	}

	return nil
}

func genSessionId() string {
	buf := make([]byte, sessionIdLength)
	_, err := rand.Read(buf)
	if err != nil {
		panic(err.Error())
	}
	return base64.StdEncoding.EncodeToString(buf)
}

// Callback for new SOCKS requests.
func handler(conn *pt.SocksConn) error {
	handlerChan <- 1
	defer func() {
		handlerChan <- -1
	}()

	defer conn.Close()
	err := conn.Grant(&net.TCPAddr{IP: net.ParseIP("0.0.0.0"), Port: 0})
	if err != nil {
		return err
	}

	var info RequestInfo
	info.SessionID = genSessionId()

	// First check url= SOCKS arg, then --url option, then SOCKS target.
	urlArg, ok := conn.Req.Args.Get("url")
	if ok {
	} else if options.URL != "" {
		urlArg = options.URL
	} else {
		urlArg = (&url.URL{
			Scheme: "http",
			Host:   conn.Req.Target,
			Path:   "/",
		}).String()
	}
	info.URL, err = url.Parse(urlArg)
	if err != nil {
		return err
	}

	// First check front= SOCKS arg, then --front option.
	front, ok := conn.Req.Args.Get("front")
	if ok {
	} else if options.Front != "" {
		front = options.Front
		ok = true
	}
	if ok {
		info.Host = info.URL.Host
		info.URL.Host = front
	}

	// First check http-proxy= SOCKS arg, then --http-proxy option.
	httpProxy, ok := conn.Req.Args.Get("http-proxy")
	if ok {
		info.HTTPProxyURL, err = url.Parse(httpProxy)
		if err != nil {
			return err
		}
	} else if options.HTTPProxyURL != nil {
		info.HTTPProxyURL = options.HTTPProxyURL
	}

	return copyLoop(conn, &info)
}

func acceptLoop(ln *pt.SocksListener) error {
	defer ln.Close()
	for {
		conn, err := ln.AcceptSocks()
		if err != nil {
			log.Printf("error in AcceptSocks: %s", err)
			if e, ok := err.(net.Error); ok && !e.Temporary() {
				return err
			}
			continue
		}
		go func() {
			err := handler(conn)
			if err != nil {
				log.Printf("error in handling request: %s", err)
			}
		}()
	}
	return nil
}

func main() {
	var helperAddr string
	var httpProxy string
	var logFilename string
	var err error

	flag.StringVar(&options.Front, "front", "", "front domain name if no front= SOCKS arg")
	flag.StringVar(&helperAddr, "helper", "", "address of HTTP helper (browser extension)")
	flag.StringVar(&httpProxy, "http-proxy", "", "HTTP proxy URL (default from http-proxy= SOCKS arg or HTTP_PROXY environment variable)")
	flag.StringVar(&logFilename, "log", "", "name of log file")
	flag.StringVar(&options.URL, "url", "", "URL to request if no url= SOCKS arg")
	flag.Parse()

	if logFilename != "" {
		f, err := os.OpenFile(logFilename, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
		if err != nil {
			log.Fatalf("error opening log file: %s", err)
		}
		defer f.Close()
		log.SetOutput(f)
	}

	if helperAddr != "" && httpProxy != "" {
		log.Fatalf("--helper and --http-proxy can't be used together")
	}

	if helperAddr != "" {
		options.HelperAddr, err = net.ResolveTCPAddr("tcp", helperAddr)
		if err != nil {
			log.Fatalf("can't resolve helper address: %s", err)
		}
		log.Printf("using helper on %s", options.HelperAddr)
	}

	if httpProxy != "" {
		options.HTTPProxyURL, err = url.Parse(httpProxy)
		if err != nil {
			log.Fatalf("can't parse HTTP proxy URL: %s", err)
		}
	}

	ptInfo, err = pt.ClientSetup([]string{ptMethodName})
	if err != nil {
		log.Fatalf("error in ClientSetup: %s", err)
	}

	listeners := make([]net.Listener, 0)
	for _, methodName := range ptInfo.MethodNames {
		switch methodName {
		case ptMethodName:
			ln, err := pt.ListenSocks("tcp", "127.0.0.1:0")
			if err != nil {
				pt.CmethodError(methodName, err.Error())
				break
			}
			go acceptLoop(ln)
			pt.Cmethod(methodName, ln.Version(), ln.Addr())
			log.Printf("listening on %s", ln.Addr())
			listeners = append(listeners, ln)
		default:
			pt.CmethodError(methodName, "no such method")
		}
	}
	pt.CmethodsDone()

	var numHandlers int = 0
	var sig os.Signal
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// Wait for first signal.
	sig = nil
	for sig == nil {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case sig = <-sigChan:
		}
	}
	for _, ln := range listeners {
		ln.Close()
	}

	log.Printf("got signal %s", sig)
	if sig == syscall.SIGTERM {
		log.Printf("done")
		return
	}

	// Wait for second signal or no more handlers.
	sig = nil
	for sig == nil && numHandlers != 0 {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case sig = <-sigChan:
		}
	}

	log.Printf("got second signal %s", sig)
	log.Printf("done")
}
<commit_msg>Comment typo.<commit_after>// meek-client is the client transport plugin for the meek pluggable transport.
//
// Sample usage in torrc:
// 	Bridge meek 0.0.2.0:1
// 	ClientTransportPlugin meek exec ./meek-client --url=https://meek-reflect.appspot.com/ --front=www.google.com --log meek-client.log
// The transport ignores the bridge address 0.0.2.0:1 and instead connects to
// the URL given by --url. When --front is given, the domain in the URL is
// replaced by the front domain for the purpose of the DNS lookup, TCP
// connection, and TLS SNI, but the HTTP Host header in the request will be the
// one in --url. (For example, in the configuration above, the connection will
// appear on the outside to be going to www.google.com, but it will actually be
// dispatched to meek-reflect.appspot.com by the Google frontend server.)
//
// Most user configuration can happen either through SOCKS args (i.e., args on a
// Bridge line) or through command line options. SOCKS args take precedence
// per-connection over command line options. For example, this configuration
// using SOCKS args:
// 	Bridge meek 0.0.2.0:1 url=https://meek-reflect.appspot.com/ front=www.google.com
// 	ClientTransportPlugin meek exec ./meek-client
// is the same as this one using command line options:
// 	Bridge meek 0.0.2.0:1
// 	ClientTransportPlugin meek exec ./meek-client --url=https://meek-reflect.appspot.com/ --front=www.google.com
// The advantage of SOCKS args is that multiple Bridge lines can have different
// configurations, but it requires a newer tor.
//
// The --helper option prevents this program from doing any network operations
// itself. Rather, it will send all requests through a browser extension that
// makes HTTP requests.
package main

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"syscall"
	"time"
)

import "git.torproject.org/pluggable-transports/goptlib.git"

const (
	ptMethodName = "meek"
	// A session ID is a randomly generated string that identifies a
	// long-lived session. We split a TCP stream across multiple HTTP
	// requests, and those with the same session ID belong to the same
	// stream.
	sessionIdLength = 32
	// The size of the largest chunk of data we will read from the SOCKS
	// port before forwarding it in a request, and the maximum size of a
	// body we are willing to handle in a reply.
	maxPayloadLength = 0x10000
	// We must poll the server to see if it has anything to send; there is
	// no way for the server to push data back to us until we send an HTTP
	// request. When a timer expires, we send a request even if it has an
	// empty body. The interval starts at this value and then grows.
	initPollInterval = 100 * time.Millisecond
	// Maximum polling interval.
	maxPollInterval = 5 * time.Second
	// Geometric increase in the polling interval each time we fail to read
	// data.
	pollIntervalMultiplier = 1.5
	// Safety limits on interaction with the HTTP helper.
	maxHelperResponseLength = 10000000
	helperReadTimeout       = 60 * time.Second
	helperWriteTimeout      = 2 * time.Second
)

var ptInfo pt.ClientInfo

// Store for command line options.
var options struct {
	URL          string
	Front        string
	HTTPProxyURL *url.URL
	HelperAddr   *net.TCPAddr
}

// When a connection handler starts, +1 is written to this channel; when it
// ends, -1 is written.
var handlerChan = make(chan int)

// RequestInfo encapsulates all the configuration used for a request–response
// roundtrip, including variables that may come from SOCKS args or from the
// command line.
type RequestInfo struct {
	// What to put in the X-Session-ID header.
	SessionID string
	// The URL to request.
	URL *url.URL
	// The Host header to put in the HTTP request (optional and may be
	// different from the host name in URL).
	Host string
	// URL of an HTTP proxy to use. If nil, the default net/http library's
	// behavior is used, which is to check the HTTP_PROXY and http_proxy
	// environment for a proxy URL.
	HTTPProxyURL *url.URL
}

// Do an HTTP roundtrip using the payload data in buf and the request metadata
// in info.
func roundTripWithHTTP(buf []byte, info *RequestInfo) (*http.Response, error) {
	tr := http.DefaultTransport
	if info.HTTPProxyURL != nil {
		tr = &http.Transport{
			Proxy: http.ProxyURL(info.HTTPProxyURL),
		}
	}
	req, err := http.NewRequest("POST", info.URL.String(), bytes.NewReader(buf))
	if err != nil {
		return nil, err
	}
	if info.Host != "" {
		req.Host = info.Host
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	req.Header.Set("X-Session-Id", info.SessionID)
	return tr.RoundTrip(req)
}

// Send the data in buf to the remote URL, wait for a reply, and feed the reply
// body back into conn.
func sendRecv(buf []byte, conn net.Conn, info *RequestInfo) (int64, error) {
	roundTrip := roundTripWithHTTP
	if options.HelperAddr != nil {
		roundTrip = roundTripWithHelper
	}
	resp, err := roundTrip(buf, info)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return 0, errors.New(fmt.Sprintf("status code was %d, not %d", resp.StatusCode, http.StatusOK))
	}

	return io.Copy(conn, io.LimitReader(resp.Body, maxPayloadLength))
}

// Repeatedly read from conn, issue HTTP requests, and write the responses back
// to conn.
func copyLoop(conn net.Conn, info *RequestInfo) error {
	var interval time.Duration

	ch := make(chan []byte)

	// Read from the Conn and send byte slices on the channel.
	go func() {
		var buf [maxPayloadLength]byte
		r := bufio.NewReader(conn)
		for {
			n, err := r.Read(buf[:])
			b := make([]byte, n)
			copy(b, buf[:n])
			// log.Printf("read from local: %q", b)
			ch <- b
			if err != nil {
				log.Printf("error reading from local: %s", err)
				break
			}
		}
		close(ch)
	}()

	interval = initPollInterval
loop:
	for {
		var buf []byte
		var ok bool

		// log.Printf("waiting up to %.2f s", interval.Seconds())
		// start := time.Now()
		select {
		case buf, ok = <-ch:
			if !ok {
				break loop
			}
			// log.Printf("read %d bytes from local after %.2f s", len(buf), time.Since(start).Seconds())
		case <-time.After(interval):
			// log.Printf("read nothing from local after %.2f s", time.Since(start).Seconds())
			buf = nil
		}

		nw, err := sendRecv(buf, conn, info)
		if err != nil {
			return err
		}
		/*
			if nw > 0 {
				log.Printf("got %d bytes from remote", nw)
			} else {
				log.Printf("got nothing from remote")
			}
		*/

		if nw > 0 || len(buf) > 0 {
			// If we sent or received anything, poll again
			// immediately.
			interval = 0
		} else if interval == 0 {
			// The first time we don't send or receive anything,
			// wait a while.
			interval = initPollInterval
		} else {
			// After that, wait a little longer.
			interval = time.Duration(float64(interval) * pollIntervalMultiplier)
		}
		if interval > maxPollInterval {
			interval = maxPollInterval
		}
	}

	return nil
}

func genSessionId() string {
	buf := make([]byte, sessionIdLength)
	_, err := rand.Read(buf)
	if err != nil {
		panic(err.Error())
	}
	return base64.StdEncoding.EncodeToString(buf)
}

// Callback for new SOCKS requests.
func handler(conn *pt.SocksConn) error {
	handlerChan <- 1
	defer func() {
		handlerChan <- -1
	}()

	defer conn.Close()
	err := conn.Grant(&net.TCPAddr{IP: net.ParseIP("0.0.0.0"), Port: 0})
	if err != nil {
		return err
	}

	var info RequestInfo
	info.SessionID = genSessionId()

	// First check url= SOCKS arg, then --url option, then SOCKS target.
	urlArg, ok := conn.Req.Args.Get("url")
	if ok {
	} else if options.URL != "" {
		urlArg = options.URL
	} else {
		urlArg = (&url.URL{
			Scheme: "http",
			Host:   conn.Req.Target,
			Path:   "/",
		}).String()
	}
	info.URL, err = url.Parse(urlArg)
	if err != nil {
		return err
	}

	// First check front= SOCKS arg, then --front option.
	front, ok := conn.Req.Args.Get("front")
	if ok {
	} else if options.Front != "" {
		front = options.Front
		ok = true
	}
	if ok {
		info.Host = info.URL.Host
		info.URL.Host = front
	}

	// First check http-proxy= SOCKS arg, then --http-proxy option.
	httpProxy, ok := conn.Req.Args.Get("http-proxy")
	if ok {
		info.HTTPProxyURL, err = url.Parse(httpProxy)
		if err != nil {
			return err
		}
	} else if options.HTTPProxyURL != nil {
		info.HTTPProxyURL = options.HTTPProxyURL
	}

	return copyLoop(conn, &info)
}

func acceptLoop(ln *pt.SocksListener) error {
	defer ln.Close()
	for {
		conn, err := ln.AcceptSocks()
		if err != nil {
			log.Printf("error in AcceptSocks: %s", err)
			if e, ok := err.(net.Error); ok && !e.Temporary() {
				return err
			}
			continue
		}
		go func() {
			err := handler(conn)
			if err != nil {
				log.Printf("error in handling request: %s", err)
			}
		}()
	}
	return nil
}

func main() {
	var helperAddr string
	var httpProxy string
	var logFilename string
	var err error

	flag.StringVar(&options.Front, "front", "", "front domain name if no front= SOCKS arg")
	flag.StringVar(&helperAddr, "helper", "", "address of HTTP helper (browser extension)")
	flag.StringVar(&httpProxy, "http-proxy", "", "HTTP proxy URL (default from http-proxy= SOCKS arg or HTTP_PROXY environment variable)")
	flag.StringVar(&logFilename, "log", "", "name of log file")
	flag.StringVar(&options.URL, "url", "", "URL to request if no url= SOCKS arg")
	flag.Parse()

	if logFilename != "" {
		f, err := os.OpenFile(logFilename, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
		if err != nil {
			log.Fatalf("error opening log file: %s", err)
		}
		defer f.Close()
		log.SetOutput(f)
	}

	if helperAddr != "" && httpProxy != "" {
		log.Fatalf("--helper and --http-proxy can't be used together")
	}

	if helperAddr != "" {
		options.HelperAddr, err = net.ResolveTCPAddr("tcp", helperAddr)
		if err != nil {
			log.Fatalf("can't resolve helper address: %s", err)
		}
		log.Printf("using helper on %s", options.HelperAddr)
	}

	if httpProxy != "" {
		options.HTTPProxyURL, err = url.Parse(httpProxy)
		if err != nil {
			log.Fatalf("can't parse HTTP proxy URL: %s", err)
		}
	}

	ptInfo, err = pt.ClientSetup([]string{ptMethodName})
	if err != nil {
		log.Fatalf("error in ClientSetup: %s", err)
	}

	listeners := make([]net.Listener, 0)
	for _, methodName := range ptInfo.MethodNames {
		switch methodName {
		case ptMethodName:
			ln, err := pt.ListenSocks("tcp", "127.0.0.1:0")
			if err != nil {
				pt.CmethodError(methodName, err.Error())
				break
			}
			go acceptLoop(ln)
			pt.Cmethod(methodName, ln.Version(), ln.Addr())
			log.Printf("listening on %s", ln.Addr())
			listeners = append(listeners, ln)
		default:
			pt.CmethodError(methodName, "no such method")
		}
	}
	pt.CmethodsDone()

	var numHandlers int = 0
	var sig os.Signal
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// Wait for first signal.
	sig = nil
	for sig == nil {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case sig = <-sigChan:
		}
	}
	for _, ln := range listeners {
		ln.Close()
	}

	log.Printf("got signal %s", sig)
	if sig == syscall.SIGTERM {
		log.Printf("done")
		return
	}

	// Wait for second signal or no more handlers.
	sig = nil
	for sig == nil && numHandlers != 0 {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case sig = <-sigChan:
		}
	}

	log.Printf("got second signal %s", sig)
	log.Printf("done")
}
<|endoftext|>