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

package eval

import (
	"fmt"
	"io/ioutil"
	"math/big"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"testing"

	"neugram.io/ng/eval/environ"
	"neugram.io/ng/eval/shell"
	"neugram.io/ng/format"
	"neugram.io/ng/parser"
	"neugram.io/ng/stmt"
)

var exprTests = []struct {
	stmt string
	want interface{}
}{
	{"2+3*(x+y-2)", 23},
	{"func() int { return 7 }()", 7},
	{
		`func() int {
			if x > 2 && x < 500 {
				return z+1
			} else {
				return z-1
			}
		}()`,
		8,
	},
	{
		`func() int64 {
			x := 9
			x++
			if x > 5 {
				x = -x
			}
			return int64(x)
		}()`,
		int64(-10),
	},
	{
		`func() int {
			f := func() bool {
				x++
				return true
			}
			if x == 4 || f() {
				x += 4
			}
			if x == 1 && f() {
				x *= 4
			}
			return x
		}()`,
		8,
	},
	{
		`func() int {
			v := 2
			for i := 1; i < 4; i++ {
				v *= i
			}
			return v
		}()`,
		12,
	},
}

func mkBasicProgram() (*Program, error) {
	p := New("basic")
	if _, err := p.Eval(mustParse("x := 4"), nil); err != nil {
		return nil, err
	}
	if _, err := p.Eval(mustParse("y := 5"), nil); err != nil {
		return nil, err
	}
	if _, err := p.Eval(mustParse("z := 7"), nil); err != nil {
		return nil, err
	}
	return p, nil
}

func TestExprs(t *testing.T) {
	for _, test := range exprTests {
		p, err := mkBasicProgram()
		if err != nil {
			t.Fatalf("mkBasicProgram: %v", err)
		}
		s := mustParse(test.stmt)
		res, err := p.Eval(s, nil)
		if err != nil {
			t.Errorf("Eval(%s) error: %v", format.Stmt(s), err)
		}
		if len(res) != 1 {
			t.Errorf("Eval(%s) want *big.Int, got multi-valued (%d) result: %v", format.Stmt(s), len(res), res)
			continue
		}
		fmt.Printf("Returning Eval: %#+v\n", res)
		switch want := test.want.(type) {
		case *big.Int:
			got, ok := res[0].Interface().(*big.Int)
			if !ok {
				t.Errorf("Eval(%s) want *big.Int, got: %s (%T)", format.Stmt(s), got, res[0].Interface())
				continue
			}
			if want.Cmp(got) != 0 {
				t.Errorf("Eval(%s)=%v, want %v", format.Stmt(s), got, want)
			}
		case *big.Float:
			got, ok := res[0].Interface().(*big.Float)
			if !ok {
				t.Errorf("Eval(%s) want *big.Float, got: %s (%T)", format.Stmt(s), got, res[0].Interface())
				continue
			}
			if want.Cmp(got) != 0 {
				t.Errorf("Eval(%s)=%v, want %v", format.Stmt(s), got, want)
			}
		default:
			got := res[0].Interface()
			if got != want {
				t.Errorf("Eval(%s)=%v, want %v", format.Stmt(s), got, want)
			}
		}
	}
}

func mustParse(src string) stmt.Stmt {
	expr, err := parser.ParseStmt([]byte(src))
	if err != nil {
		panic(fmt.Sprintf("mustParse(%q): %v", src, err))
	}
	return expr
}

var errRE = regexp.MustCompile(`ERROR: (.*)`)

func TestPrograms(t *testing.T) {
	files, err := filepath.Glob("testdata/*.ng")
	if err != nil {
		t.Fatal(err)
	}
	if len(files) == 0 {
		t.Fatal("cannot find testdata")
	}
	origStdout := os.Stdout
	origStderr := os.Stderr
	defer func() {
		os.Stdout = origStdout
		os.Stderr = origStderr
	}()

	shell.Env = environ.New()
	for _, s := range os.Environ() {
		i := strings.Index(s, "=")
		shell.Env.Set(s[:i], s[i+1:])
	}
	shell.Alias = environ.New()

	for _, file := range files {
		out, err := ioutil.TempFile("", filepath.Base(file)+".stdout.")
		if err != nil {
			t.Fatal(err)
		}
		t.Logf("executing %s", file) // TODO: switch to subtests

		os.Stdout = out
		os.Stderr = out
		err = EvalFile(file)
		os.Stdout = origStdout
		os.Stderr = origStderr
		if err2 := out.Close(); err2 != nil {
			t.Fatal(err2)
		}

		if strings.HasSuffix(file, "_panic.ng") {
			if _, isPanic := err.(Panic); !isPanic {
				t.Errorf("%s: expect panic, got: %v", file, err)
				continue
			}
		} else if strings.HasSuffix(file, "_error.ng") {
			// TODO: support multiple errors
			// TODO: check line numbers
			b, ferr := ioutil.ReadFile(file)
			if ferr != nil {
				t.Error(ferr)
				continue
			}
			match := errRE.FindStringSubmatch(string(b))
			if match == nil {
				t.Errorf("%s: file has _error suffix but no ERROR directive", file)
				continue
			}
			wantStr := match[1]
			if !strings.Contains(err.Error(), wantStr) {
				t.Errorf("%s: want %q, got: %v", file, wantStr, err.Error())
			}
			continue
		} else if err != nil {
			t.Errorf("%s:%v", file, err)
			continue
		}

		b, err := ioutil.ReadFile(out.Name())
		if err != nil {
			t.Fatalf("%s: %v", file, err)
		}
		output := string(b)
		if !strings.HasSuffix(file, "_panic.ng") && !strings.HasSuffix(output, "OK\n") {
			t.Logf("Testing program %q, output:\n%s", file, output)
			t.Errorf("%s missing OK", file)
		}
	}
}
<commit_msg>eval: use subtests for TestPrograms<commit_after>// Copyright 2016 The Neugram Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package eval

import (
	"fmt"
	"io/ioutil"
	"math/big"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"testing"

	"neugram.io/ng/eval/environ"
	"neugram.io/ng/eval/shell"
	"neugram.io/ng/format"
	"neugram.io/ng/parser"
	"neugram.io/ng/stmt"
)

var exprTests = []struct {
	stmt string
	want interface{}
}{
	{"2+3*(x+y-2)", 23},
	{"func() int { return 7 }()", 7},
	{
		`func() int {
			if x > 2 && x < 500 {
				return z+1
			} else {
				return z-1
			}
		}()`,
		8,
	},
	{
		`func() int64 {
			x := 9
			x++
			if x > 5 {
				x = -x
			}
			return int64(x)
		}()`,
		int64(-10),
	},
	{
		`func() int {
			f := func() bool {
				x++
				return true
			}
			if x == 4 || f() {
				x += 4
			}
			if x == 1 && f() {
				x *= 4
			}
			return x
		}()`,
		8,
	},
	{
		`func() int {
			v := 2
			for i := 1; i < 4; i++ {
				v *= i
			}
			return v
		}()`,
		12,
	},
}

func mkBasicProgram() (*Program, error) {
	p := New("basic")
	if _, err := p.Eval(mustParse("x := 4"), nil); err != nil {
		return nil, err
	}
	if _, err := p.Eval(mustParse("y := 5"), nil); err != nil {
		return nil, err
	}
	if _, err := p.Eval(mustParse("z := 7"), nil); err != nil {
		return nil, err
	}
	return p, nil
}

func TestExprs(t *testing.T) {
	for _, test := range exprTests {
		p, err := mkBasicProgram()
		if err != nil {
			t.Fatalf("mkBasicProgram: %v", err)
		}
		s := mustParse(test.stmt)
		res, err := p.Eval(s, nil)
		if err != nil {
			t.Errorf("Eval(%s) error: %v", format.Stmt(s), err)
		}
		if len(res) != 1 {
			t.Errorf("Eval(%s) want *big.Int, got multi-valued (%d) result: %v", format.Stmt(s), len(res), res)
			continue
		}
		fmt.Printf("Returning Eval: %#+v\n", res)
		switch want := test.want.(type) {
		case *big.Int:
			got, ok := res[0].Interface().(*big.Int)
			if !ok {
				t.Errorf("Eval(%s) want *big.Int, got: %s (%T)", format.Stmt(s), got, res[0].Interface())
				continue
			}
			if want.Cmp(got) != 0 {
				t.Errorf("Eval(%s)=%v, want %v", format.Stmt(s), got, want)
			}
		case *big.Float:
			got, ok := res[0].Interface().(*big.Float)
			if !ok {
				t.Errorf("Eval(%s) want *big.Float, got: %s (%T)", format.Stmt(s), got, res[0].Interface())
				continue
			}
			if want.Cmp(got) != 0 {
				t.Errorf("Eval(%s)=%v, want %v", format.Stmt(s), got, want)
			}
		default:
			got := res[0].Interface()
			if got != want {
				t.Errorf("Eval(%s)=%v, want %v", format.Stmt(s), got, want)
			}
		}
	}
}

func mustParse(src string) stmt.Stmt {
	expr, err := parser.ParseStmt([]byte(src))
	if err != nil {
		panic(fmt.Sprintf("mustParse(%q): %v", src, err))
	}
	return expr
}

var errRE = regexp.MustCompile(`ERROR: (.*)`)

func TestPrograms(t *testing.T) {
	files, err := filepath.Glob("testdata/*.ng")
	if err != nil {
		t.Fatal(err)
	}
	if len(files) == 0 {
		t.Fatal("cannot find testdata")
	}
	origStdout := os.Stdout
	origStderr := os.Stderr
	defer func() {
		os.Stdout = origStdout
		os.Stderr = origStderr
	}()

	shell.Env = environ.New()
	for _, s := range os.Environ() {
		i := strings.Index(s, "=")
		shell.Env.Set(s[:i], s[i+1:])
	}
	shell.Alias = environ.New()

	for _, file := range files {
		// For the file "testdata/name.ng", name the subtest "name".
		test := file[len("testdata/") : len(file)-3]
		t.Run(test, func(t *testing.T) {
			out, err := ioutil.TempFile("", test+".stdout.")
			if err != nil {
				t.Fatal(err)
			}

			os.Stdout = out
			os.Stderr = out
			err = EvalFile(file)
			os.Stdout = origStdout
			os.Stderr = origStderr
			if err2 := out.Close(); err2 != nil {
				t.Fatal(err2)
			}

			if strings.HasSuffix(test, "_panic") {
				if _, isPanic := err.(Panic); !isPanic {
					t.Fatalf("expected panic, got: %v", err)
				}
			} else if strings.HasSuffix(test, "_error") {
				// TODO: support multiple errors
				// TODO: check line numbers
				b, ferr := ioutil.ReadFile(file)
				if ferr != nil {
					t.Fatal(ferr)
				}
				match := errRE.FindStringSubmatch(string(b))
				if match == nil {
					t.Fatal("test has _error suffix but no ERROR directive")
				}
				wantStr := match[1]
				if !strings.Contains(err.Error(), wantStr) {
					t.Fatalf("want %q, got: %v", wantStr, err.Error())
				}
				return
			} else if err != nil {
				t.Fatal(err)
			}

			b, err := ioutil.ReadFile(out.Name())
			if err != nil {
				t.Fatal(err)
			}
			output := string(b)
			if !strings.HasSuffix(test, "_panic") && !strings.HasSuffix(output, "OK\n") {
				t.Logf("output:\n%s", output)
				t.Error("test missing OK")
			}
		})
	}
}
<|endoftext|>


<commit_before>// run

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

// Test that the implementation catches nil ptr indirection
// in a large address space.

package main

import "unsafe"

// Having a big address space means that indexing
// at a 256 MB offset from a nil pointer might not
// cause a memory access fault. This test checks
// that Go is doing the correct explicit checks to catch
// these nil pointer accesses, not just relying on the hardware.
var dummy [256 << 20]byte // give us a big address space

func main() {
	// the test only tests what we intend to test
	// if dummy starts in the first 256 MB of memory.
	// otherwise there might not be anything mapped
	// at the address that might be accidentally
	// dereferenced below.
	if uintptr(unsafe.Pointer(&dummy)) > 256<<20 {
		panic("dummy too far out")
	}

	shouldPanic(p1)
	shouldPanic(p2)
	shouldPanic(p3)
	shouldPanic(p4)
	shouldPanic(p5)
	shouldPanic(p6)
	shouldPanic(p7)
	shouldPanic(p8)
	shouldPanic(p9)
	shouldPanic(p10)
	shouldPanic(p11)
	shouldPanic(p12)
}

func shouldPanic(f func()) {
	defer func() {
		if recover() == nil {
			panic("memory reference did not panic")
		}
	}()
	f()
}

func p1() {
	// Array index.
	var p *[1 << 30]byte = nil
	println(p[256<<20]) // very likely to be inside dummy, but should panic
}

var xb byte

func p2() {
	var p *[1 << 30]byte = nil
	xb = 123

	// Array index.
	println(p[uintptr(unsafe.Pointer(&xb))]) // should panic
}

func p3() {
	// Array to slice.
	var p *[1 << 30]byte = nil
	var x []byte = p[0:] // should panic
	_ = x
}

var q *[1 << 30]byte

func p4() {
	// Array to slice.
	var x []byte
	var y = &x
	*y = q[0:] // should crash (uses arraytoslice runtime routine)
}

func fb([]byte) {
	panic("unreachable")
}

func p5() {
	// Array to slice.
	var p *[1 << 30]byte = nil
	fb(p[0:]) // should crash
}

func p6() {
	// Array to slice.
	var p *[1 << 30]byte = nil
	var _ []byte = p[10 : len(p)-10] // should crash
}

type T struct {
	x [256 << 20]byte
	i int
}

func f() *T {
	return nil
}

var y *T
var x = &y

func p7() {
	// Struct field access with large offset.
	println(f().i) // should crash
}

func p8() {
	// Struct field access with large offset.
	println((*x).i) // should crash
}

func p9() {
	// Struct field access with large offset.
	var t *T
	println(&t.i) // should crash
}

func p10() {
	// Struct field access with large offset.
	var t *T
	println(t.i) // should crash
}

type T1 struct {
	T
}

type T2 struct {
	*T1
}

func p11() {
	t := &T2{}
	p := &t.i
	println(*p)
}

// ADDR(DOT(IND(p))) needs a check also
func p12() {
	var p *T = nil
	println(*(&((*p).i)))
}
<commit_msg>test/nilptr: add more tests<commit_after>// run

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

// Test that the implementation catches nil ptr indirection
// in a large address space.

package main

import "unsafe"

// Having a big address space means that indexing
// at a 256 MB offset from a nil pointer might not
// cause a memory access fault. This test checks
// that Go is doing the correct explicit checks to catch
// these nil pointer accesses, not just relying on the hardware.
var dummy [256 << 20]byte // give us a big address space

func main() {
	// the test only tests what we intend to test
	// if dummy starts in the first 256 MB of memory.
	// otherwise there might not be anything mapped
	// at the address that might be accidentally
	// dereferenced below.
	if uintptr(unsafe.Pointer(&dummy)) > 256<<20 {
		panic("dummy too far out")
	}

	shouldPanic(p1)
	shouldPanic(p2)
	shouldPanic(p3)
	shouldPanic(p4)
	shouldPanic(p5)
	shouldPanic(p6)
	shouldPanic(p7)
	shouldPanic(p8)
	shouldPanic(p9)
	shouldPanic(p10)
	shouldPanic(p11)
	shouldPanic(p12)
	shouldPanic(p13)
	shouldPanic(p14)
	shouldPanic(p15)
	shouldPanic(p16)
}

func shouldPanic(f func()) {
	defer func() {
		if recover() == nil {
			panic("memory reference did not panic")
		}
	}()
	f()
}

func p1() {
	// Array index.
	var p *[1 << 30]byte = nil
	println(p[256<<20]) // very likely to be inside dummy, but should panic
}

var xb byte

func p2() {
	var p *[1 << 30]byte = nil
	xb = 123

	// Array index.
	println(p[uintptr(unsafe.Pointer(&xb))]) // should panic
}

func p3() {
	// Array to slice.
	var p *[1 << 30]byte = nil
	var x []byte = p[0:] // should panic
	_ = x
}

var q *[1 << 30]byte

func p4() {
	// Array to slice.
	var x []byte
	var y = &x
	*y = q[0:] // should crash (uses arraytoslice runtime routine)
}

func fb([]byte) {
	panic("unreachable")
}

func p5() {
	// Array to slice.
	var p *[1 << 30]byte = nil
	fb(p[0:]) // should crash
}

func p6() {
	// Array to slice.
	var p *[1 << 30]byte = nil
	var _ []byte = p[10 : len(p)-10] // should crash
}

type T struct {
	x [256 << 20]byte
	i int
}

func f() *T {
	return nil
}

var y *T
var x = &y

func p7() {
	// Struct field access with large offset.
	println(f().i) // should crash
}

func p8() {
	// Struct field access with large offset.
	println((*x).i) // should crash
}

func p9() {
	// Struct field access with large offset.
	var t *T
	println(&t.i) // should crash
}

func p10() {
	// Struct field access with large offset.
	var t *T
	println(t.i) // should crash
}

type T1 struct {
	T
}

type T2 struct {
	*T1
}

func p11() {
	t := &T2{}
	p := &t.i
	println(*p)
}

// ADDR(DOT(IND(p))) needs a check also
func p12() {
	var p *T = nil
	println(*(&((*p).i)))
}

// Tests suggested in golang.org/issue/6080.

func p13() {
	var x *[10]int
	y := x[:]
	_ = y
}

func p14() {
	println((*[1]int)(nil)[:])
}

func p15() {
	for i := range (*[1]int)(nil)[:] {
		_ = i
	}
}

func p16() {
	for i, v := range (*[1]int)(nil)[:] {
		_ = i + v
	}
}
<|endoftext|>


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

// Binary wrap is a test helper program for //elisp:binary_test, which see.
package main

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

	"github.com/google/go-cmp/cmp"
	"github.com/phst/runfiles"
)

func main() {
	log.Println("Args:", os.Args)
	log.Println("Environment:", os.Environ())
	var manifestFile string
	flag.StringVar(&manifestFile, "manifest", "", "")
	flag.Parse()
	if manifestFile == "" {
		log.Fatal("--manifest is empty")
	}
	workspaceDir, err := runfiles.Path("phst_rules_elisp")
	if err != nil {
		log.Fatal(err)
	}
	gotArgs := flag.Args()
	wantArgs := []string{
		"--quick", "--batch",
		"--directory=" + workspaceDir,
		"--option",
		"elisp/binary.cc",
		" \t\n\r\f äα𝐴🐈'\\\"",
		"/:/tmp/output.dat",
	}
	if diff := cmp.Diff(gotArgs, wantArgs); diff != "" {
		log.Fatalf("positional arguments: -git +want:\n%s", diff)
	}
	jsonData, err := ioutil.ReadFile(manifestFile)
	if err != nil {
		log.Fatalf("can’t read manifest: %s", err)
	}
	var gotManifest map[string]interface{}
	if err := json.Unmarshal(jsonData, &gotManifest); err != nil {
		log.Fatalf("can’t decode manifest: %s", err)
	}
	wantManifest := map[string]interface{}{
		"root":        "RUNFILES_ROOT",
		"tags":        []interface{}{"local", "mytag"},
		"loadPath":    []interface{}{"phst_rules_elisp"},
		"inputFiles":  []interface{}{"phst_rules_elisp/elisp/binary.cc", "phst_rules_elisp/elisp/binary.h"},
		"outputFiles": []interface{}{"/tmp/output.dat"},
	}
	if diff := cmp.Diff(gotManifest, wantManifest); diff != "" {
		log.Fatalf("manifest: -got +want:\n%s", diff)
	}
}
<commit_msg>Fix typo<commit_after>// Copyright 2020, 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Binary wrap is a test helper program for //elisp:binary_test, which see.
package main

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

	"github.com/google/go-cmp/cmp"
	"github.com/phst/runfiles"
)

func main() {
	log.Println("Args:", os.Args)
	log.Println("Environment:", os.Environ())
	var manifestFile string
	flag.StringVar(&manifestFile, "manifest", "", "")
	flag.Parse()
	if manifestFile == "" {
		log.Fatal("--manifest is empty")
	}
	workspaceDir, err := runfiles.Path("phst_rules_elisp")
	if err != nil {
		log.Fatal(err)
	}
	gotArgs := flag.Args()
	wantArgs := []string{
		"--quick", "--batch",
		"--directory=" + workspaceDir,
		"--option",
		"elisp/binary.cc",
		" \t\n\r\f äα𝐴🐈'\\\"",
		"/:/tmp/output.dat",
	}
	if diff := cmp.Diff(gotArgs, wantArgs); diff != "" {
		log.Fatalf("positional arguments: -got +want:\n%s", diff)
	}
	jsonData, err := ioutil.ReadFile(manifestFile)
	if err != nil {
		log.Fatalf("can’t read manifest: %s", err)
	}
	var gotManifest map[string]interface{}
	if err := json.Unmarshal(jsonData, &gotManifest); err != nil {
		log.Fatalf("can’t decode manifest: %s", err)
	}
	wantManifest := map[string]interface{}{
		"root":        "RUNFILES_ROOT",
		"tags":        []interface{}{"local", "mytag"},
		"loadPath":    []interface{}{"phst_rules_elisp"},
		"inputFiles":  []interface{}{"phst_rules_elisp/elisp/binary.cc", "phst_rules_elisp/elisp/binary.h"},
		"outputFiles": []interface{}{"/tmp/output.dat"},
	}
	if diff := cmp.Diff(gotManifest, wantManifest); diff != "" {
		log.Fatalf("manifest: -got +want:\n%s", diff)
	}
}
<|endoftext|>


<commit_before>package rabbitgo

import (
  "errors"
  "github.com/streadway/amqp"
  log "github.com/koding/logging"
)

type Consumer struct {
	conn        *Connection
	ch          *amqp.Channel
  e           *Exchange
  q           *Queue
	deliveries  <-chan amqp.Delivery
	handler     func(amqp.Delivery)
  cc          *ConsumerConfig
  bc          *BindingConfig
	// A notifiyng channel for publishings
	// will be used for sync. between close channel and consume handler
	done        chan error
}

type ConsumerConfig struct {
	Tag string
	AutoAck    bool
	Exclusive  bool
	NoLocal    bool
	NoWait     bool
	Args       amqp.Table
}

type BindingConfig struct {
	RoutingKey string
	NoWait bool
	Args amqp.Table
}

func (c *Consumer) Deliveries() <-chan amqp.Delivery {
	return c.deliveries
}

// NewConsumer is a constructor for consumer creation
// Accepts Exchange, Queue, BindingOptions and ConsumerOptions
func (c *Connection) NewConsumer(e *Exchange, q *Queue, bc *BindingConfig, cc *ConsumerConfig) (*Consumer, error) {
  if c.ch == nil {
    return nil, errors.New("No channel found. Are you using NewConnection?")
  }
	consumer := &Consumer {
		ch:   c.ch,
		done: make(chan error),
    cc:   cc,
    bc:   bc,
		e:    e,
    q:    q,
	}
	if err := consumer.connect(); err != nil {
		return nil, err
	}
	return consumer, nil
}

// connect internally declares the exchanges and queues
func (c *Consumer) connect() error {
	err := c.ch.ExchangeDeclare(
		c.e.Name,       // name of the exchange
		c.e.Type,       // type
		c.e.Durable,    // durable
		c.e.AutoDelete, // delete when complete
		c.e.Internal,   // internal
		c.e.NoWait,     // noWait
		c.e.Args,       // arguments
	)
  if err != nil {
		return err
	}
	_, err = c.ch.QueueDeclare(
		c.q.Name,       // name of the queue
		c.q.Durable,    // durable
		c.q.AutoDelete, // delete when usused
		c.q.Exclusive,  // exclusive
		c.q.NoWait,     // noWait
		c.q.Args,       // arguments
	)
	if err != nil {
		return err
	}
	err = c.ch.QueueBind(
		// bind to real queue
		c.q.Name,        // name of the queue
		c.bc.RoutingKey, // bindingKey
		c.e.Name,        // sourceExchange
		c.bc.NoWait,     // noWait
		c.bc.Args,       // arguments
	)
  if err != nil {
		return err
	}
	return nil
}

// Consume accepts a handler function for every message streamed from RabbitMq
// will be called within this handler func
func (c *Consumer) Consume(handler func(delivery amqp.Delivery)) error {
	deliveries, err := c.ch.Consume(
		c.q.Name,       // name
		c.cc.Tag,       // consumerTag,
		c.cc.AutoAck,   // autoAck
		c.cc.Exclusive, // exclusive
		c.cc.NoLocal,   // noLocal
		c.cc.NoWait,    // noWait
		c.cc.Args,      // arguments
	)
	if err != nil {
		return err
	}

	// should we stop streaming, in order not to consume from server?
	c.deliveries = deliveries
	c.handler = handler

	log.Info("handle: deliveries channel starting")

	// handle all consumer errors, if required re-connect
	// there are problems with reconnection logic for now
	for delivery := range c.deliveries {
		handler(delivery)
	}

	log.Info("handle: deliveries channel closed")
	c.done <- nil
	return nil
}

// QOS controls how many messages the server will try to keep on the network for
// consumers before receiving delivery acks.  The intent of Qos is to make sure
// the network buffers stay full between the server and client.
func (c *Consumer) QOS(messageCount int) error {
	return c.ch.Qos(messageCount, 0, false)
}

// ConsumeMessage accepts a handler function and only consumes one message
// stream from RabbitMq
func (c *Consumer) Get(handler func(delivery amqp.Delivery)) error {
	message, ok, err := c.ch.Get(c.q.Name, c.cc.AutoAck)
	if err != nil {
		return err
	}
	c.handler = handler
	if ok {
		log.Debug("Message received")
		handler(message)
	} else {
		log.Debug("No message received")
	}
	// TODO maybe we should return ok too?
	return nil
}
<commit_msg>ConsumerRPC allows you to reply just by returning the message and content type from the handler.<commit_after>package rabbitgo

import (
  "errors"
  "github.com/streadway/amqp"
  log "github.com/koding/logging"
)

type Consumer struct {
	conn        *Connection
	ch          *amqp.Channel
  e           *Exchange
  q           *Queue
	deliveries  <-chan amqp.Delivery
  handler     func(amqp.Delivery)
	handlerRPC  func(amqp.Delivery)([]byte, string)
  cc          *ConsumerConfig
  bc          *BindingConfig
	// A notifiyng channel for publishings
	// will be used for sync. between close channel and consume handler
	done        chan error
}

type ConsumerConfig struct {
	Tag string
	AutoAck    bool
	Exclusive  bool
	NoLocal    bool
	NoWait     bool
	Args       amqp.Table
}

type BindingConfig struct {
	RoutingKey string
	NoWait bool
	Args amqp.Table
}

func (c *Consumer) Deliveries() <-chan amqp.Delivery {
	return c.deliveries
}

// NewConsumer is a constructor for consumer creation
// Accepts Exchange, Queue, BindingOptions and ConsumerOptions
func (c *Connection) NewConsumer(e *Exchange, q *Queue, bc *BindingConfig, cc *ConsumerConfig) (*Consumer, error) {
  if c.ch == nil {
    return nil, errors.New("No channel found. Are you using NewConnection?")
  }
	consumer := &Consumer {
		ch:   c.ch,
		done: make(chan error),
    cc:   cc,
    bc:   bc,
		e:    e,
    q:    q,
	}
	if err := consumer.connect(); err != nil {
		return nil, err
	}
	return consumer, nil
}

// connect internally declares the exchanges and queues
func (c *Consumer) connect() error {
	err := c.ch.ExchangeDeclare(
		c.e.Name,       // name of the exchange
		c.e.Type,       // type
		c.e.Durable,    // durable
		c.e.AutoDelete, // delete when complete
		c.e.Internal,   // internal
		c.e.NoWait,     // noWait
		c.e.Args,       // arguments
	)
  if err != nil {
		return err
	}
	_, err = c.ch.QueueDeclare(
		c.q.Name,       // name of the queue
		c.q.Durable,    // durable
		c.q.AutoDelete, // delete when usused
		c.q.Exclusive,  // exclusive
		c.q.NoWait,     // noWait
		c.q.Args,       // arguments
	)
	if err != nil {
		return err
	}
	err = c.ch.QueueBind(
		// bind to real queue
		c.q.Name,        // name of the queue
		c.bc.RoutingKey, // bindingKey
		c.e.Name,        // sourceExchange
		c.bc.NoWait,     // noWait
		c.bc.Args,       // arguments
	)
  if err != nil {
		return err
	}
	return nil
}

// Consume accepts a handler function for every message streamed from RabbitMq
// will be called within this handler func
func (c *Consumer) Consume(handler func(delivery amqp.Delivery)) error {
	deliveries, err := c.ch.Consume(
		c.q.Name,       // name
		c.cc.Tag,       // consumerTag,
		c.cc.AutoAck,   // autoAck
		c.cc.Exclusive, // exclusive
		c.cc.NoLocal,   // noLocal
		c.cc.NoWait,    // noWait
		c.cc.Args,      // arguments
	)
	if err != nil {
		return err
	}

	// should we stop streaming, in order not to consume from server?
	c.deliveries = deliveries
	c.handler = handler

	log.Info("handle: deliveries channel starting")

	// handle all consumer errors, if required re-connect
	// there are problems with reconnection logic for now
	for delivery := range c.deliveries {
		handler(delivery)
	}

	log.Info("handle: deliveries channel closed")
	c.done <- nil
	return nil
}

// ConsumeRPC accepts a handler function for every message streamed from RabbitMq
// will be called within this handler func.
// It returns the message to send and the content type.
func (c *Consumer) ConsumeRPC(handler func(delivery amqp.Delivery)([]byte, string)) error {
	deliveries, err := c.ch.Consume(
		c.q.Name,       // name
		c.cc.Tag,       // consumerTag,
		c.cc.AutoAck,   // autoAck
		c.cc.Exclusive, // exclusive
		c.cc.NoLocal,   // noLocal
		c.cc.NoWait,    // noWait
		c.cc.Args,      // arguments
	)
	if err != nil {
		return err
	}

	// should we stop streaming, in order not to consume from server?
	c.deliveries = deliveries
	c.handlerRPC = handler

	log.Info("handle: deliveries channel starting")

	// handle all consumer errors, if required re-connect
	// there are problems with reconnection logic for now
	for delivery := range c.deliveries {
		body, contentType := handler(delivery)
    replyTo := delivery.ReplyTo
    if replyTo != "" {
      err := c.ch.Publish(
    		"",
    		replyTo,
    		false,
    		false,
    		amqp.Publishing{
    			ContentType:   contentType,
    			CorrelationId: delivery.CorrelationId,
    			Body:          body,
    		},
    	)
      if err != nil {
        // TODO: What if err != nil?
      }
    }
	}

	log.Info("handle: deliveries channel closed")
	c.done <- nil
	return nil
}

// QOS controls how many messages the server will try to keep on the network for
// consumers before receiving delivery acks.  The intent of Qos is to make sure
// the network buffers stay full between the server and client.
func (c *Consumer) QOS(messageCount int) error {
	return c.ch.Qos(messageCount, 0, false)
}

// ConsumeMessage accepts a handler function and only consumes one message
// stream from RabbitMq
func (c *Consumer) Get(handler func(delivery amqp.Delivery)) error {
	message, ok, err := c.ch.Get(c.q.Name, c.cc.AutoAck)
	if err != nil {
		return err
	}
	c.handler = handler
	if ok {
		log.Debug("Message received")
		handler(message)
	} else {
		log.Debug("No message received")
	}
	// TODO maybe we should return ok too?
	return nil
}
<|endoftext|>


<commit_before>package spamc

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/textproto"
	"os"
	"strconv"
	"strings"
	"time"
)

const clientProtocolVersion = "1.5"

var serverProtocolVersions = []string{"1.0", "1.1"}

// mapping of the error codes to the error messages.
var errorMessages = map[int]string{
	64: "Command line usage error",               // EX_USAGE
	65: "Data format error",                      // EX_DATA_ERR
	66: "Cannot open input",                      // EX_NO_INPUT
	67: "Addressee unknown",                      // EX_NO_USER
	68: "Host name unknown",                      // EX_NO_HOST
	69: "Service unavailable",                    // EX_UNAVAILABLE
	70: "Internal software error",                // EX_SOFTWARE
	71: "System error",                           // EX_OSERR
	72: "Critical OS file missing",               // EX_OSFILE
	73: "Can't create (user) output file",        // EX_CANTCREAT
	74: "Input/output error",                     // EX_IOERR
	75: "Temp failure; user is invited to retry", // EX_TEMPFAIL
	76: "Remote error in protocol",               // EX_PROTOCOL
	77: "Permission denied",                      // EX_NOPERM
	78: "Configuration error",                    // EX_CONFIG
	79: "Read timeout",                           // EX_TIMEOUT
}

// Temporary hack to write tests.
var testConnHook net.Conn

// send a command to spamd.
func (c *Client) send(
	ctx context.Context,
	cmd string,
	message io.ReadSeeker,
	headers Header,
) (io.ReadCloser, error) {

	var conn net.Conn
	if testConnHook != nil {
		conn = testConnHook
	} else {
		var err error
		conn, err = c.dial(ctx)
		if err != nil {
			return nil, fmt.Errorf("could not dial to %v: %v", c.host, err)
		}
	}

	if err := c.write(conn, cmd, message, headers); err != nil {
		return nil, err
	}

	return conn, nil
}

// isASCII checks for non-ascii bytes
func (c Client) isASCII(m []byte) bool {
	for _, c := range m {
		if c > 127 {
			return false
		}
	}
	return true
}

// isMessageTerminated will check to see if the message is \r\n terminated
func (c Client) isMessageTerminated(message io.ReadSeeker) (bool, error) {
	// make sure to seek back to the current position or
	// we mess up the internal message pointer
	current, _ := message.Seek(0, io.SeekCurrent)
	defer message.Seek(current, io.SeekStart) // nolint: errcheck

	_, err := message.Seek(-2, io.SeekEnd)

	if err != nil {
		return false, err
	}
	buf := make([]byte, 2)
	_, err = message.Read(buf)
	if err != nil {
		return false, err
	}
	if bytes.Equal(buf, []byte("\r\n")) {
		// and we are done, its terminated
		return true, nil
	}
	return false, nil
}

// rearrangeMessageReader will rearrange the message reader to make sure SA will understand the message
func (c Client) rearrangeMessageReader(message io.ReadSeeker, size int64) (io.Reader, int64, error) {
	var headers bytes.Buffer
	var rest bytes.Buffer

	r := bufio.NewReader(message)

	var returnError error

	// we assume that there is no newline between headers and body, so increment size
	size += 2

	for {
		chars, err := r.Peek(2)
		if err != nil {
			returnError = err
			break
		}
		if bytes.Equal(chars, []byte("\r\n")) {
			// we were wrong there was a newline between headers and body
			// decrement size
			size -= 2
			// read and discard it, and we are done
			_, err = r.Read(make([]byte, 2))
			if err != nil {
				returnError = err
			}
			break
		}
		if chars[0] == ' ' || chars[0] == '\t' {
			// if we have a previous key append value, and continue
			// if not we are not reading headers
			if headers.Len() > 0 {
				line, err := r.ReadSlice('\n')
				if err != nil {
					rest.Write(line)
					returnError = err
					break
				}
				headers.Write(line)
				continue
			} else {
				// this line start with whitespace but there is no
				// previous header to append, so it's not a header
				returnError = errors.New("multiline header part found without first header")
				break
			}
		}

		// it should be a header read it
		line, err := r.ReadSlice('\n')

		if err != nil {
			rest.Write(line)
			returnError = err
			break
		}
		// Check for a key, we must have one for it to be a header
		i := bytes.IndexByte(line, ':')
		if i < 0 {
			rest.Write(line)
			returnError = errors.New("malformed MIME header line: [" + string(line) + "]")
			break
		}
		if !c.isASCII(line[0:i]) {
			rest.Write(line)
			returnError = errors.New("non ascii characters found in header key: [" + string(line[0:i]) + "]")
			break
		}
		headers.Write(line)
	}

	messageTerminated, err := c.isMessageTerminated(message)
	if err != nil {
		returnError = err
	}

	end := ""
	if !messageTerminated {
		size += 2
		end = "\r\n"
	}

	resultReader := io.MultiReader(
		bufio.NewReader(&headers),
		strings.NewReader("\r\n"),
		bufio.NewReader(&rest),
		r,
		strings.NewReader(end),
	)

	// we might have hit EOF, reset that
	if returnError == io.EOF {
		returnError = nil
	}

	return resultReader, size, returnError
}

// write the command to the connection.
func (c *Client) write(
	conn net.Conn,
	cmd string,
	message io.ReadSeeker,
	headers Header,
) error {

	if strings.TrimSpace(cmd) == "" {
		return errors.New("empty command")
	}

	if headers == nil {
		headers = make(Header)
	}
	if _, ok := headers[HeaderUser]; !ok && c.DefaultUser != "" {
		headers.Add(HeaderUser, c.DefaultUser)
	}

	buf := bytes.NewBufferString("")
	w := bufio.NewWriter(buf)
	tp := textproto.NewWriter(w)

	// Make sure to seek to start of message
	if _, err := message.Seek(0, io.SeekStart); err != nil {
		return err
	}

	size := int64(-1)
	switch v := message.(type) {
	case *strings.Reader:
		size = v.Size()
	case *bytes.Reader:
		size = v.Size()
	case *os.File:
		stat, err := v.Stat()
		if err != nil {
			return err
		}
		size = stat.Size()
	}

	// make sure to delete a present content-length (and use it if
	// we could not set size above
	if _, ok := headers[HeaderContentLength]; ok {
		if size == -1 {
			headerContentLength, err := strconv.ParseInt(headers.Get(HeaderContentLength), 10, 64)
			if err != nil {
				return err
			}
			size = headerContentLength
		}
	}

	// Make the sure message always ends in \r\n, if it doesn't SA will just
	// keep reading from the connection and it will block until it timeouts.

	// messageReader contains the new message messageSize and we suppress warnings, _
	messageReader, messageSize, _ := c.rearrangeMessageReader(message, size)

	err := tp.PrintfLine("%v SPAMC/%v", cmd, clientProtocolVersion)
	if err != nil {
		return err
	}

	// Always add Content-length header.
	if messageSize >= 0 {
		err = tp.PrintfLine("Content-length: %v", messageSize)
		if err != nil {
			return err
		}
	}

	// Write headers.
	for k, vals := range headers {
		for _, v := range vals {
			err := tp.PrintfLine("%v: %v", k, v)
			if err != nil {
				return err
			}
		}
	}
	err = tp.PrintfLine("")
	if err != nil {
		return err
	}

	err = tp.W.Flush()
	if err != nil {
		return err
	}

	// Write to spamd.
	if _, err := io.Copy(conn, io.MultiReader(buf, messageReader)); err != nil {
		conn.Close() // nolint: errcheck
		return fmt.Errorf("could not send to spamd: %v", err)
	}

	return nil
}

func (c *Client) dial(ctx context.Context) (net.Conn, error) {
	conn, err := c.dialer.DialContext(ctx, "tcp", c.host)
	if err != nil {
		if conn != nil {
			conn.Close() // nolint: errcheck
		}
		return nil, fmt.Errorf("could not connect to spamd: %v", err)
	}

	// Set connection timeout
	err = conn.SetDeadline(time.Now().Add(c.dialer.Timeout))
	if err != nil {
		conn.Close() // nolint: errcheck
		return nil, fmt.Errorf("connection to spamd timed out: %v", err)
	}

	return conn, nil
}

// The spamd protocol is a HTTP-esque protocol; a response's first line is the
// response code:
//
//     SPAMD/1.1 0 EX_OK\r\n
//
// Next, it can set some headers:
//
//     Content-length: <size>\r\n
//
// After a blank line we get the response body, which is different for the
// various commands.
//
// A non-0 (or EX_OK) status code is considered an error.
func readResponse(read io.Reader) (headers Header, body string, err error) {
	data := bufio.NewReader(read)
	tp := textproto.NewReader(data)

	// We can't use textproto's ReadCodeLine() here, as SA's response is not
	// quite compatible.
	if err := parseCodeLine(tp, false); err != nil {
		return nil, "", err
	}

	tpHeader, err := tp.ReadMIMEHeader()
	if err != nil {
		return nil, "", fmt.Errorf("could not read headers: %v", err)
	}

	headers = Header(tpHeader)

	body, err = readBody(tp)
	if err != nil {
		return nil, "", fmt.Errorf("could not read body: %v", err)
	}

	return headers, body, nil
}

func parseCodeLine(tp *textproto.Reader, isPing bool) error {
	line, err := tp.ReadLine()
	if err != nil {
		return err
	}

	if len(line) < 11 {
		return fmt.Errorf("short response: %v", line)
	}
	if !strings.HasPrefix(line, "SPAMD/") {
		return fmt.Errorf("unrecognised response: %v", line)
	}

	version := line[6:9]

	// The PING command is special as it will return the *client* version,
	// rather than the server version.
	if isPing {
		if version != clientProtocolVersion {
			return fmt.Errorf("unexpected version: %v; we expected %v",
				version, clientProtocolVersion)
		}
	} else {
		// in some errors it uses version 1.0, so accept both 1.0 and 1.1.
		//     spamd/1.0 76 bad header line: asdasd
		if !supportedVersion(version) {
			return fmt.Errorf("unknown server protocol version %v; we only understand versions %v",
				version, serverProtocolVersions)
		}
	}

	s := strings.Split(line[10:], " ")
	code, err := strconv.Atoi(s[0])
	if err != nil {
		return fmt.Errorf("could not parse return code: %v", err)
	}
	if code != 0 {
		text := strings.Join(s[1:], " ")
		if msg, ok := errorMessages[code]; ok {
			return fmt.Errorf("spamd returned code %v: %v: %v", code, msg, text)
		}
		return fmt.Errorf("spamd returned code %v: %v", code, text)
	}

	return nil
}

func supportedVersion(v string) bool {
	for i := range serverProtocolVersions {
		if serverProtocolVersions[i] == v {
			return true
		}
	}
	return false
}

func readBody(tp *textproto.Reader) (string, error) {
	body := ""
loop:
	for {
		line, err := tp.ReadLine()
		switch err {
		case nil:
			// Do nothing
		case io.EOF:
			break loop
		default:
			return "", err
		}

		body += line + "\r\n"
	}

	return body, nil
}

// Parse the Spam: response header:
//    Spam <yes|no> ; <score> / <base-score>
// example:
//    Spam: yes ; 6.66 / 5.0
func parseSpamHeader(respHeaders Header) (bool, float64, float64, error) {
	spam, ok := respHeaders["Spam"]
	if !ok || len(spam) == 0 {
		return false, 0, 0, errors.New("header missing")
	}

	if len(spam[0]) == 0 {
		return false, 0, 0, errors.New("header empty")
	}

	s := strings.Split(spam[0], ";")
	if len(s) != 2 {
		return false, 0, 0, fmt.Errorf("unexpected data: %v", spam[0])
	}

	isSpam := false
	switch strings.ToLower(strings.TrimSpace(s[0])) {
	case "true", "yes":
		isSpam = true
	case "false", "no":
		isSpam = false
	default:
		return false, 0, 0, fmt.Errorf("unknown spam status: %v", s[0])
	}

	split := strings.Split(s[1], "/")
	if len(split) != 2 {
		return false, 0, 0, fmt.Errorf("unexpected data: %v", s[1])
	}
	score, err := strconv.ParseFloat(strings.TrimSpace(split[0]), 64)
	if err != nil {
		return false, 0, 0, fmt.Errorf("could not parse spam score: %v", err)
	}
	baseScore, err := strconv.ParseFloat(strings.TrimSpace(split[1]), 64)
	if err != nil {
		return false, 0, 0, fmt.Errorf("could not parse base spam score: %v", err)
	}

	return isSpam, score, baseScore, nil
}
<commit_msg>Bugfixed, unfortunately HeaderContentLength is not Canonical<commit_after>package spamc

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/textproto"
	"os"
	"strconv"
	"strings"
	"time"
)

const clientProtocolVersion = "1.5"

var serverProtocolVersions = []string{"1.0", "1.1"}

// mapping of the error codes to the error messages.
var errorMessages = map[int]string{
	64: "Command line usage error",               // EX_USAGE
	65: "Data format error",                      // EX_DATA_ERR
	66: "Cannot open input",                      // EX_NO_INPUT
	67: "Addressee unknown",                      // EX_NO_USER
	68: "Host name unknown",                      // EX_NO_HOST
	69: "Service unavailable",                    // EX_UNAVAILABLE
	70: "Internal software error",                // EX_SOFTWARE
	71: "System error",                           // EX_OSERR
	72: "Critical OS file missing",               // EX_OSFILE
	73: "Can't create (user) output file",        // EX_CANTCREAT
	74: "Input/output error",                     // EX_IOERR
	75: "Temp failure; user is invited to retry", // EX_TEMPFAIL
	76: "Remote error in protocol",               // EX_PROTOCOL
	77: "Permission denied",                      // EX_NOPERM
	78: "Configuration error",                    // EX_CONFIG
	79: "Read timeout",                           // EX_TIMEOUT
}

// Temporary hack to write tests.
var testConnHook net.Conn

// send a command to spamd.
func (c *Client) send(
	ctx context.Context,
	cmd string,
	message io.ReadSeeker,
	headers Header,
) (io.ReadCloser, error) {

	var conn net.Conn
	if testConnHook != nil {
		conn = testConnHook
	} else {
		var err error
		conn, err = c.dial(ctx)
		if err != nil {
			return nil, fmt.Errorf("could not dial to %v: %v", c.host, err)
		}
	}

	if err := c.write(conn, cmd, message, headers); err != nil {
		return nil, err
	}

	return conn, nil
}

// isASCII checks for non-ascii bytes
func (c Client) isASCII(m []byte) bool {
	for _, c := range m {
		if c > 127 {
			return false
		}
	}
	return true
}

// isMessageTerminated will check to see if the message is \r\n terminated
func (c Client) isMessageTerminated(message io.ReadSeeker) (bool, error) {
	// make sure to seek back to the current position or
	// we mess up the internal message pointer
	current, _ := message.Seek(0, io.SeekCurrent)
	defer message.Seek(current, io.SeekStart) // nolint: errcheck

	_, err := message.Seek(-2, io.SeekEnd)

	if err != nil {
		return false, err
	}
	buf := make([]byte, 2)
	_, err = message.Read(buf)
	if err != nil {
		return false, err
	}
	if bytes.Equal(buf, []byte("\r\n")) {
		// and we are done, its terminated
		return true, nil
	}
	return false, nil
}

// rearrangeMessageReader will rearrange the message reader to make sure SA will understand the message
func (c Client) rearrangeMessageReader(message io.ReadSeeker, size int64) (io.Reader, int64, error) {
	var headers bytes.Buffer
	var rest bytes.Buffer

	r := bufio.NewReader(message)

	var returnError error

	// we assume that there is no newline between headers and body, so increment size
	size += 2

	for {
		chars, err := r.Peek(2)
		if err != nil {
			returnError = err
			break
		}
		if bytes.Equal(chars, []byte("\r\n")) {
			// we were wrong there was a newline between headers and body
			// decrement size
			size -= 2
			// read and discard it, and we are done
			_, err = r.Read(make([]byte, 2))
			if err != nil {
				returnError = err
			}
			break
		}
		if chars[0] == ' ' || chars[0] == '\t' {
			// if we have a previous key append value, and continue
			// if not we are not reading headers
			if headers.Len() > 0 {
				line, err := r.ReadSlice('\n')
				if err != nil {
					rest.Write(line)
					returnError = err
					break
				}
				headers.Write(line)
				continue
			} else {
				// this line start with whitespace but there is no
				// previous header to append, so it's not a header
				returnError = errors.New("multiline header part found without first header")
				break
			}
		}

		// it should be a header read it
		line, err := r.ReadSlice('\n')

		if err != nil {
			rest.Write(line)
			returnError = err
			break
		}
		// Check for a key, we must have one for it to be a header
		i := bytes.IndexByte(line, ':')
		if i < 0 {
			rest.Write(line)
			returnError = errors.New("malformed MIME header line: [" + string(line) + "]")
			break
		}
		if !c.isASCII(line[0:i]) {
			rest.Write(line)
			returnError = errors.New("non ascii characters found in header key: [" + string(line[0:i]) + "]")
			break
		}
		headers.Write(line)
	}

	messageTerminated, err := c.isMessageTerminated(message)
	if err != nil {
		returnError = err
	}

	end := ""
	if !messageTerminated {
		size += 2
		end = "\r\n"
	}

	resultReader := io.MultiReader(
		bufio.NewReader(&headers),
		strings.NewReader("\r\n"),
		bufio.NewReader(&rest),
		r,
		strings.NewReader(end),
	)

	// we might have hit EOF, reset that
	if returnError == io.EOF {
		returnError = nil
	}

	return resultReader, size, returnError
}

// write the command to the connection.
func (c *Client) write(
	conn net.Conn,
	cmd string,
	message io.ReadSeeker,
	headers Header,
) error {

	if strings.TrimSpace(cmd) == "" {
		return errors.New("empty command")
	}

	if headers == nil {
		headers = make(Header)
	}
	if _, ok := headers[HeaderUser]; !ok && c.DefaultUser != "" {
		headers.Add(HeaderUser, c.DefaultUser)
	}

	buf := bytes.NewBufferString("")
	w := bufio.NewWriter(buf)
	tp := textproto.NewWriter(w)

	// Make sure to seek to start of message
	if _, err := message.Seek(0, io.SeekStart); err != nil {
		return err
	}

	size := int64(-1)
	switch v := message.(type) {
	case *strings.Reader:
		size = v.Size()
	case *bytes.Reader:
		size = v.Size()
	case *os.File:
		stat, err := v.Stat()
		if err != nil {
			return err
		}
		size = stat.Size()
	}

	// make sure to delete a present content-length (and use it if
	// we could not set size above
	if _, ok := headers[HeaderContentLength]; ok {
		if size == -1 {
			headerContentLength, err := strconv.ParseInt(headers.Get(HeaderContentLength), 10, 64)
			if err != nil {
				return err
			}
			size = headerContentLength
		}
		// delete the Content-length header from headers is set, prevents double sending, can not use
		// headers.Delete here .. HeaderContentLength is not Canonical
		delete(headers, HeaderContentLength)
	}

	// Make the sure message always ends in \r\n, if it doesn't SA will just
	// keep reading from the connection and it will block until it timeouts.

	// messageReader contains the new message messageSize and we suppress warnings, _
	messageReader, messageSize, _ := c.rearrangeMessageReader(message, size)

	err := tp.PrintfLine("%v SPAMC/%v", cmd, clientProtocolVersion)
	if err != nil {
		return err
	}

	// Always add Content-length header.
	if messageSize >= 0 {
		err = tp.PrintfLine("Content-length: %v", messageSize)
		if err != nil {
			return err
		}
	}

	// Write headers.
	for k, vals := range headers {
		for _, v := range vals {
			err := tp.PrintfLine("%v: %v", k, v)
			if err != nil {
				return err
			}
		}
	}
	err = tp.PrintfLine("")
	if err != nil {
		return err
	}

	err = tp.W.Flush()
	if err != nil {
		return err
	}

	// Write to spamd.
	if _, err := io.Copy(conn, io.MultiReader(buf, messageReader)); err != nil {
		conn.Close() // nolint: errcheck
		return fmt.Errorf("could not send to spamd: %v", err)
	}

	return nil
}

func (c *Client) dial(ctx context.Context) (net.Conn, error) {
	conn, err := c.dialer.DialContext(ctx, "tcp", c.host)
	if err != nil {
		if conn != nil {
			conn.Close() // nolint: errcheck
		}
		return nil, fmt.Errorf("could not connect to spamd: %v", err)
	}

	// Set connection timeout
	err = conn.SetDeadline(time.Now().Add(c.dialer.Timeout))
	if err != nil {
		conn.Close() // nolint: errcheck
		return nil, fmt.Errorf("connection to spamd timed out: %v", err)
	}

	return conn, nil
}

// The spamd protocol is a HTTP-esque protocol; a response's first line is the
// response code:
//
//     SPAMD/1.1 0 EX_OK\r\n
//
// Next, it can set some headers:
//
//     Content-length: <size>\r\n
//
// After a blank line we get the response body, which is different for the
// various commands.
//
// A non-0 (or EX_OK) status code is considered an error.
func readResponse(read io.Reader) (headers Header, body string, err error) {
	data := bufio.NewReader(read)
	tp := textproto.NewReader(data)

	// We can't use textproto's ReadCodeLine() here, as SA's response is not
	// quite compatible.
	if err := parseCodeLine(tp, false); err != nil {
		return nil, "", err
	}

	tpHeader, err := tp.ReadMIMEHeader()
	if err != nil {
		return nil, "", fmt.Errorf("could not read headers: %v", err)
	}

	headers = Header(tpHeader)

	body, err = readBody(tp)
	if err != nil {
		return nil, "", fmt.Errorf("could not read body: %v", err)
	}

	return headers, body, nil
}

func parseCodeLine(tp *textproto.Reader, isPing bool) error {
	line, err := tp.ReadLine()
	if err != nil {
		return err
	}

	if len(line) < 11 {
		return fmt.Errorf("short response: %v", line)
	}
	if !strings.HasPrefix(line, "SPAMD/") {
		return fmt.Errorf("unrecognised response: %v", line)
	}

	version := line[6:9]

	// The PING command is special as it will return the *client* version,
	// rather than the server version.
	if isPing {
		if version != clientProtocolVersion {
			return fmt.Errorf("unexpected version: %v; we expected %v",
				version, clientProtocolVersion)
		}
	} else {
		// in some errors it uses version 1.0, so accept both 1.0 and 1.1.
		//     spamd/1.0 76 bad header line: asdasd
		if !supportedVersion(version) {
			return fmt.Errorf("unknown server protocol version %v; we only understand versions %v",
				version, serverProtocolVersions)
		}
	}

	s := strings.Split(line[10:], " ")
	code, err := strconv.Atoi(s[0])
	if err != nil {
		return fmt.Errorf("could not parse return code: %v", err)
	}
	if code != 0 {
		text := strings.Join(s[1:], " ")
		if msg, ok := errorMessages[code]; ok {
			return fmt.Errorf("spamd returned code %v: %v: %v", code, msg, text)
		}
		return fmt.Errorf("spamd returned code %v: %v", code, text)
	}

	return nil
}

func supportedVersion(v string) bool {
	for i := range serverProtocolVersions {
		if serverProtocolVersions[i] == v {
			return true
		}
	}
	return false
}

func readBody(tp *textproto.Reader) (string, error) {
	body := ""
loop:
	for {
		line, err := tp.ReadLine()
		switch err {
		case nil:
			// Do nothing
		case io.EOF:
			break loop
		default:
			return "", err
		}

		body += line + "\r\n"
	}

	return body, nil
}

// Parse the Spam: response header:
//    Spam <yes|no> ; <score> / <base-score>
// example:
//    Spam: yes ; 6.66 / 5.0
func parseSpamHeader(respHeaders Header) (bool, float64, float64, error) {
	spam, ok := respHeaders["Spam"]
	if !ok || len(spam) == 0 {
		return false, 0, 0, errors.New("header missing")
	}

	if len(spam[0]) == 0 {
		return false, 0, 0, errors.New("header empty")
	}

	s := strings.Split(spam[0], ";")
	if len(s) != 2 {
		return false, 0, 0, fmt.Errorf("unexpected data: %v", spam[0])
	}

	isSpam := false
	switch strings.ToLower(strings.TrimSpace(s[0])) {
	case "true", "yes":
		isSpam = true
	case "false", "no":
		isSpam = false
	default:
		return false, 0, 0, fmt.Errorf("unknown spam status: %v", s[0])
	}

	split := strings.Split(s[1], "/")
	if len(split) != 2 {
		return false, 0, 0, fmt.Errorf("unexpected data: %v", s[1])
	}
	score, err := strconv.ParseFloat(strings.TrimSpace(split[0]), 64)
	if err != nil {
		return false, 0, 0, fmt.Errorf("could not parse spam score: %v", err)
	}
	baseScore, err := strconv.ParseFloat(strings.TrimSpace(split[1]), 64)
	if err != nil {
		return false, 0, 0, fmt.Errorf("could not parse base spam score: %v", err)
	}

	return isSpam, score, baseScore, nil
}
<|endoftext|>


<commit_before>package ss

import (
	"fmt"
	"io"
	"net"
	"os/user"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"

	"github.com/olekukonko/tablewriter"
)

type TransportProtocol int

const (
	TCP TransportProtocol = iota
	TCP6
)

var (
	stringToProtocol = map[string]TransportProtocol{
		"tcp":  TCP,
		"tcp6": TCP6,
	}
	protocolToString = map[TransportProtocol]string{
		TCP:  "tcp",
		TCP6: "tcp6",
	}
)

// Process describes OS processes.
type Process struct {
	Protocol string

	Program string
	PID     int

	LocalIP   string
	LocalPort string

	RemoteIP   string
	RemotePort string

	State string

	User user.User
}

var processMembers = []string{
	"PROTOCOL",
	"PROGRAM",
	"PID",
	"LOCAL_ADDR",
	"REMOTE_ADDR",
	"USER",
}

// String describes Process type.
func (p Process) String() string {
	return fmt.Sprintf("Protocol: '%s' | Program: '%s' | PID: '%d' | LocalAddr: '%s%s' | RemoteAddr: '%s%s' | State: '%s' | User: '%+v'",
		p.Protocol,
		p.Program,
		p.PID,
		p.LocalIP,
		p.LocalPort,
		p.RemoteIP,
		p.RemotePort,
		p.State,
		p.User,
	)
}

// UnsafeKillPIDs lists PIDs to be protected.
// (https://en.wikipedia.org/wiki/List_of_TCP_and_UDP_port_numbers)
var UnsafeKillPIDs = map[int]struct{}{
	0:  struct{}{},
	22: struct{}{},
}

// UnsafeKillNames lists program names to be protected.
// (https://en.wikipedia.org/wiki/List_of_TCP_and_UDP_port_numbers)
var UnsafeKillNames = map[string]struct{}{
	"/bin/sh":      struct{}{},
	"/bin/bash":    struct{}{},
	"/usr/bin/ssh": struct{}{},
	"systemd":      struct{}{},
}

func filterOut(pidToKill *map[int]string) {
	m := *pidToKill
	for k, v := range m {
		if _, ok := UnsafeKillPIDs[k]; ok {
			delete(m, k)
		} else {
			for pname := range UnsafeKillNames {
				if strings.HasSuffix(v, pname) {
					delete(m, k)
				}
			}
		}
	}
}

// Kill kills all processes in arguments.
func Kill(w io.Writer, force bool, ps ...Process) {
	defer func() {
		if err := recover(); err != nil {
			fmt.Fprintln(w, "Kill:", err)
		}
	}()

	pidToKill := make(map[int]string)
	for _, p := range ps {
		if _, ok := pidToKill[p.PID]; !ok {
			pidToKill[p.PID] = p.Program
		}
	}
	if !force {
		filterOut(&pidToKill)
	}
	pids := []int{}
	for pid := range pidToKill {
		pids = append(pids, pid)
	}
	sort.Ints(pids)

	for _, pid := range pids {
		fmt.Fprintf(w, "syscall.Kill: %s [PID: %d]\n", pidToKill[pid], pid)
		if err := syscall.Kill(pid, syscall.SIGTERM); err != nil {
			fmt.Fprintf(w, "error when sending syscall.SIGTERM (%v):", err)
		}
		if err := syscall.Kill(pid, syscall.SIGKILL); err != nil {
			fmt.Fprintf(w, "error when sending syscall.SIGKILL (%v):", err)
		}
	}
}

// parseLittleEndianIpv4 parses hexadecimal ipv4 IP addresses.
// For example, it converts '0101007F:0035' into string.
// It assumes that the system has little endian order.
func parseLittleEndianIpv4(s string) (string, string, error) {
	arr := strings.Split(s, ":")
	if len(arr) != 2 {
		return "", "", fmt.Errorf("cannot parse ipv4 %s", s)
	}
	if len(arr[0]) != 8 {
		return "", "", fmt.Errorf("cannot parse ipv4 ip %s", arr[0])
	}
	if len(arr[1]) != 4 {
		return "", "", fmt.Errorf("cannot parse ipv4 port %s", arr[1])
	}

	d0, err := strconv.ParseInt(arr[0][6:8], 16, 32)
	if err != nil {
		return "", "", err
	}
	d1, err := strconv.ParseInt(arr[0][4:6], 16, 32)
	if err != nil {
		return "", "", err
	}
	d2, err := strconv.ParseInt(arr[0][2:4], 16, 32)
	if err != nil {
		return "", "", err
	}
	d3, err := strconv.ParseInt(arr[0][0:2], 16, 32)
	if err != nil {
		return "", "", err
	}
	ip := fmt.Sprintf("%d.%d.%d.%d", d0, d1, d2, d3)

	p0, err := strconv.ParseInt(arr[1], 16, 32)
	if err != nil {
		return "", "", err
	}
	port := fmt.Sprintf("%d", p0)

	return ip, ":" + port, nil
}

// parseLittleEndianIpv6 parses hexadecimal ipv6 IP addresses.
// For example, it converts '4506012691A700C165EB1DE1F912918C:8BDA'
// into string.
// It assumes that the system has little endian order.
func parseLittleEndianIpv6(s string) (string, string, error) {
	arr := strings.Split(s, ":")
	if len(arr) != 2 {
		return "", "", fmt.Errorf("cannot parse ipv6 %s", s)
	}
	if len(arr[0]) != 32 {
		return "", "", fmt.Errorf("cannot parse ipv6 ip %s", arr[0])
	}
	if len(arr[1]) != 4 {
		return "", "", fmt.Errorf("cannot parse ipv6 port %s", arr[1])
	}

	// 32 characters, reverse by 2 characters
	ip := ""
	for i := 32; i > 0; i -= 2 {
		ip += fmt.Sprintf("%s", arr[0][i-2:i])
		if (len(ip)+1)%5 == 0 && len(ip) != 39 {
			ip += ":"
		}
	}

	p0, err := strconv.ParseInt(arr[1], 16, 32)
	if err != nil {
		return "", "", err
	}
	port := fmt.Sprintf("%d", p0)

	return ip, ":" + port, nil
}

// List lists all processes. filter is used to return Processes
// that matches the member values in filter struct.
func List(filter *Process, opts ...TransportProtocol) ([]Process, error) {
	rs := []Process{}
	for _, opt := range opts {

		ps, err := list(opt)
		if err != nil {
			return nil, err
		}

		if filter == nil {
			rs = append(rs, ps...)
			continue
		}

		pCh, done := make(chan Process), make(chan struct{})
		fv := *filter

		for _, p := range ps {

			go func(ft, process Process) {

				if !ft.Match(process) {
					done <- struct{}{}
					return
				}
				pCh <- process

			}(fv, p)

		}

		cn := 0
		ns := []Process{}
		for cn != len(ps) {
			select {
			case p := <-pCh:
				ns = append(ns, p)
				cn++
			case <-done:
				cn++
			}
		}

		close(pCh)
		close(done)

		rs = append(rs, ns...)
	}
	return rs, nil
}

// Match returns true if the Process matches the filter.
func (filter Process) Match(p Process) bool {
	if filter.Protocol != "" {
		if p.Protocol != filter.Protocol {
			return false
		}
	}
	// matches the suffix
	if filter.Program != "" {
		if !strings.HasSuffix(p.Program, filter.Program) {
			return false
		}
	}
	if filter.PID != 0 {
		if p.PID != filter.PID {
			return false
		}
	}
	if filter.LocalIP != "" {
		if p.LocalIP != filter.LocalIP {
			return false
		}
	}
	if filter.LocalPort != "" {
		p0 := p.LocalPort
		if !strings.HasPrefix(p0, ":") {
			p0 = ":" + p0
		}
		p1 := filter.LocalPort
		if !strings.HasPrefix(p1, ":") {
			p1 = ":" + p1
		}
		if p0 != p1 {
			return false
		}
	}
	if filter.RemoteIP != "" {
		if p.RemoteIP != filter.RemoteIP {
			return false
		}
	}
	if filter.RemotePort != "" {
		p0 := p.RemotePort
		if !strings.HasPrefix(p0, ":") {
			p0 = ":" + p0
		}
		p1 := filter.RemotePort
		if !strings.HasPrefix(p1, ":") {
			p1 = ":" + p1
		}
		if p0 != p1 {
			return false
		}
	}
	if filter.State != "" {
		if p.State != filter.State {
			return false
		}
	}
	// currently only support user name
	if filter.User.Username != "" {
		if p.User.Username != filter.User.Username {
			return false
		}
	}
	return true
}

// WriteToTable writes slice of Processes to ASCII table.
func WriteToTable(w io.Writer, ps ...Process) {
	table := tablewriter.NewWriter(w)
	table.SetHeader(processMembers)

	rows := make([][]string, len(ps))
	for i, p := range ps {
		sl := make([]string, len(processMembers))
		sl[0] = p.Protocol
		sl[1] = p.Program
		sl[2] = strconv.Itoa(p.PID)
		sl[3] = p.LocalIP + p.LocalPort
		sl[4] = p.RemoteIP + p.RemotePort
		sl[5] = p.User.Name
		rows[i] = sl
	}

	by(
		rows,
		makeAscendingFunc(0), // PROTOCOL
		makeAscendingFunc(1), // PROGRAM
		makeAscendingFunc(2), // PID
		makeAscendingFunc(3), // LOCAL_ADDR
		makeAscendingFunc(5), // USER
	).Sort(rows)

	for _, row := range rows {
		table.Append(row)
	}

	table.Render()
}

// ListPorts lists all ports that are being used.
func ListPorts(filter *Process, opts ...TransportProtocol) map[string]struct{} {
	rm := make(map[string]struct{})
	ps, _ := List(filter, opts...)
	for _, p := range ps {
		rm[p.LocalPort] = struct{}{}
		rm[p.RemotePort] = struct{}{}
	}
	return rm
}

type Ports struct {
	mu        sync.Mutex // guards the following
	beingUsed map[string]struct{}
}

func NewPorts() *Ports {
	p := &Ports{}
	p.beingUsed = make(map[string]struct{})
	return p
}

func (p *Ports) Exist(port string) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	_, ok := p.beingUsed[port]
	return ok
}

func (p *Ports) Add(ports ...string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, port := range ports {
		p.beingUsed[port] = struct{}{}
	}
}

func (p *Ports) Free(ports ...string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, port := range ports {
		delete(p.beingUsed, port)
	}
}

// Refresh refreshes ports that are being used by the OS.
// You would use as follows:
//
//	var (
//		globalPorts     = NewPorts()
//		refreshInterval = 10 * time.Second
//	)
//	func init() {
//		globalPorts.Refresh()
//		go func() {
//			for {
//				select {
//				case <-time.After(refreshInterval):
//					globalPorts.Refresh()
//				}
//			}
//		}()
//	}
//
func (p *Ports) Refresh() {
	p.mu.Lock()
	p.beingUsed = ListPorts(nil, TCP, TCP6)
	p.mu.Unlock()
}

func getFreePort(opts ...TransportProtocol) (string, error) {
	fp := ""
	var errMsg error
	for _, opt := range opts {
		protocolStr := ""
		switch opt {
		case TCP:
			protocolStr = "tcp"
		case TCP6:
			protocolStr = "tcp6"
		}
		addr, err := net.ResolveTCPAddr(protocolStr, "localhost:0")
		if err != nil {
			errMsg = err
			continue
		}
		l, err := net.ListenTCP(protocolStr, addr)
		if err != nil {
			errMsg = err
			continue
		}
		pd := l.Addr().(*net.TCPAddr).Port
		l.Close()
		fp = ":" + strconv.Itoa(pd)
		break
	}
	return fp, errMsg
}

// GetFreePorts returns multiple free ports from OS.
func GetFreePorts(num int, opts ...TransportProtocol) ([]string, error) {
	try := 0
	rm := make(map[string]struct{})
	for len(rm) != num {
		if try > 150 {
			return nil, fmt.Errorf("too many tries to find free ports")
		}
		try++
		p, err := getFreePort(opts...)
		if err != nil {
			return nil, err
		}
		rm[p] = struct{}{}
	}
	rs := []string{}
	for p := range rm {
		rs = append(rs, p)
	}
	sort.Strings(rs)
	return rs, nil
}
<commit_msg>ss: add sbin/init to whitelist<commit_after>package ss

import (
	"fmt"
	"io"
	"net"
	"os/user"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"

	"github.com/olekukonko/tablewriter"
)

type TransportProtocol int

const (
	TCP TransportProtocol = iota
	TCP6
)

var (
	stringToProtocol = map[string]TransportProtocol{
		"tcp":  TCP,
		"tcp6": TCP6,
	}
	protocolToString = map[TransportProtocol]string{
		TCP:  "tcp",
		TCP6: "tcp6",
	}
)

// Process describes OS processes.
type Process struct {
	Protocol string

	Program string
	PID     int

	LocalIP   string
	LocalPort string

	RemoteIP   string
	RemotePort string

	State string

	User user.User
}

var processMembers = []string{
	"PROTOCOL",
	"PROGRAM",
	"PID",
	"LOCAL_ADDR",
	"REMOTE_ADDR",
	"USER",
}

// String describes Process type.
func (p Process) String() string {
	return fmt.Sprintf("Protocol: '%s' | Program: '%s' | PID: '%d' | LocalAddr: '%s%s' | RemoteAddr: '%s%s' | State: '%s' | User: '%+v'",
		p.Protocol,
		p.Program,
		p.PID,
		p.LocalIP,
		p.LocalPort,
		p.RemoteIP,
		p.RemotePort,
		p.State,
		p.User,
	)
}

// UnsafeKillPIDs lists PIDs to be protected.
// (https://en.wikipedia.org/wiki/List_of_TCP_and_UDP_port_numbers)
var UnsafeKillPIDs = map[int]struct{}{
	0:  struct{}{},
	22: struct{}{},
}

// UnsafeKillNames lists program names to be protected.
// (https://en.wikipedia.org/wiki/List_of_TCP_and_UDP_port_numbers)
var UnsafeKillNames = map[string]struct{}{
	"/bin/sh":      struct{}{},
	"/bin/bash":    struct{}{},
	"sbin/init":    struct{}{},
	"/usr/bin/ssh": struct{}{},
	"systemd":      struct{}{},
}

func filterOut(pidToKill *map[int]string) {
	m := *pidToKill
	for k, v := range m {
		if _, ok := UnsafeKillPIDs[k]; ok {
			delete(m, k)
		} else {
			for pname := range UnsafeKillNames {
				if strings.HasSuffix(v, pname) {
					delete(m, k)
				}
			}
		}
	}
}

// Kill kills all processes in arguments.
func Kill(w io.Writer, force bool, ps ...Process) {
	defer func() {
		if err := recover(); err != nil {
			fmt.Fprintln(w, "Kill:", err)
		}
	}()

	pidToKill := make(map[int]string)
	for _, p := range ps {
		if _, ok := pidToKill[p.PID]; !ok {
			pidToKill[p.PID] = p.Program
		}
	}
	if !force {
		filterOut(&pidToKill)
	}
	pids := []int{}
	for pid := range pidToKill {
		pids = append(pids, pid)
	}
	sort.Ints(pids)

	for _, pid := range pids {
		fmt.Fprintf(w, "syscall.Kill: %s [PID: %d]\n", pidToKill[pid], pid)
		if err := syscall.Kill(pid, syscall.SIGTERM); err != nil {
			fmt.Fprintf(w, "error when sending syscall.SIGTERM (%v):", err)
		}
		if err := syscall.Kill(pid, syscall.SIGKILL); err != nil {
			fmt.Fprintf(w, "error when sending syscall.SIGKILL (%v):", err)
		}
	}
}

// parseLittleEndianIpv4 parses hexadecimal ipv4 IP addresses.
// For example, it converts '0101007F:0035' into string.
// It assumes that the system has little endian order.
func parseLittleEndianIpv4(s string) (string, string, error) {
	arr := strings.Split(s, ":")
	if len(arr) != 2 {
		return "", "", fmt.Errorf("cannot parse ipv4 %s", s)
	}
	if len(arr[0]) != 8 {
		return "", "", fmt.Errorf("cannot parse ipv4 ip %s", arr[0])
	}
	if len(arr[1]) != 4 {
		return "", "", fmt.Errorf("cannot parse ipv4 port %s", arr[1])
	}

	d0, err := strconv.ParseInt(arr[0][6:8], 16, 32)
	if err != nil {
		return "", "", err
	}
	d1, err := strconv.ParseInt(arr[0][4:6], 16, 32)
	if err != nil {
		return "", "", err
	}
	d2, err := strconv.ParseInt(arr[0][2:4], 16, 32)
	if err != nil {
		return "", "", err
	}
	d3, err := strconv.ParseInt(arr[0][0:2], 16, 32)
	if err != nil {
		return "", "", err
	}
	ip := fmt.Sprintf("%d.%d.%d.%d", d0, d1, d2, d3)

	p0, err := strconv.ParseInt(arr[1], 16, 32)
	if err != nil {
		return "", "", err
	}
	port := fmt.Sprintf("%d", p0)

	return ip, ":" + port, nil
}

// parseLittleEndianIpv6 parses hexadecimal ipv6 IP addresses.
// For example, it converts '4506012691A700C165EB1DE1F912918C:8BDA'
// into string.
// It assumes that the system has little endian order.
func parseLittleEndianIpv6(s string) (string, string, error) {
	arr := strings.Split(s, ":")
	if len(arr) != 2 {
		return "", "", fmt.Errorf("cannot parse ipv6 %s", s)
	}
	if len(arr[0]) != 32 {
		return "", "", fmt.Errorf("cannot parse ipv6 ip %s", arr[0])
	}
	if len(arr[1]) != 4 {
		return "", "", fmt.Errorf("cannot parse ipv6 port %s", arr[1])
	}

	// 32 characters, reverse by 2 characters
	ip := ""
	for i := 32; i > 0; i -= 2 {
		ip += fmt.Sprintf("%s", arr[0][i-2:i])
		if (len(ip)+1)%5 == 0 && len(ip) != 39 {
			ip += ":"
		}
	}

	p0, err := strconv.ParseInt(arr[1], 16, 32)
	if err != nil {
		return "", "", err
	}
	port := fmt.Sprintf("%d", p0)

	return ip, ":" + port, nil
}

// List lists all processes. filter is used to return Processes
// that matches the member values in filter struct.
func List(filter *Process, opts ...TransportProtocol) ([]Process, error) {
	rs := []Process{}
	for _, opt := range opts {

		ps, err := list(opt)
		if err != nil {
			return nil, err
		}

		if filter == nil {
			rs = append(rs, ps...)
			continue
		}

		pCh, done := make(chan Process), make(chan struct{})
		fv := *filter

		for _, p := range ps {

			go func(ft, process Process) {

				if !ft.Match(process) {
					done <- struct{}{}
					return
				}
				pCh <- process

			}(fv, p)

		}

		cn := 0
		ns := []Process{}
		for cn != len(ps) {
			select {
			case p := <-pCh:
				ns = append(ns, p)
				cn++
			case <-done:
				cn++
			}
		}

		close(pCh)
		close(done)

		rs = append(rs, ns...)
	}
	return rs, nil
}

// Match returns true if the Process matches the filter.
func (filter Process) Match(p Process) bool {
	if filter.Protocol != "" {
		if p.Protocol != filter.Protocol {
			return false
		}
	}
	// matches the suffix
	if filter.Program != "" {
		if !strings.HasSuffix(p.Program, filter.Program) {
			return false
		}
	}
	if filter.PID != 0 {
		if p.PID != filter.PID {
			return false
		}
	}
	if filter.LocalIP != "" {
		if p.LocalIP != filter.LocalIP {
			return false
		}
	}
	if filter.LocalPort != "" {
		p0 := p.LocalPort
		if !strings.HasPrefix(p0, ":") {
			p0 = ":" + p0
		}
		p1 := filter.LocalPort
		if !strings.HasPrefix(p1, ":") {
			p1 = ":" + p1
		}
		if p0 != p1 {
			return false
		}
	}
	if filter.RemoteIP != "" {
		if p.RemoteIP != filter.RemoteIP {
			return false
		}
	}
	if filter.RemotePort != "" {
		p0 := p.RemotePort
		if !strings.HasPrefix(p0, ":") {
			p0 = ":" + p0
		}
		p1 := filter.RemotePort
		if !strings.HasPrefix(p1, ":") {
			p1 = ":" + p1
		}
		if p0 != p1 {
			return false
		}
	}
	if filter.State != "" {
		if p.State != filter.State {
			return false
		}
	}
	// currently only support user name
	if filter.User.Username != "" {
		if p.User.Username != filter.User.Username {
			return false
		}
	}
	return true
}

// WriteToTable writes slice of Processes to ASCII table.
func WriteToTable(w io.Writer, ps ...Process) {
	table := tablewriter.NewWriter(w)
	table.SetHeader(processMembers)

	rows := make([][]string, len(ps))
	for i, p := range ps {
		sl := make([]string, len(processMembers))
		sl[0] = p.Protocol
		sl[1] = p.Program
		sl[2] = strconv.Itoa(p.PID)
		sl[3] = p.LocalIP + p.LocalPort
		sl[4] = p.RemoteIP + p.RemotePort
		sl[5] = p.User.Name
		rows[i] = sl
	}

	by(
		rows,
		makeAscendingFunc(0), // PROTOCOL
		makeAscendingFunc(1), // PROGRAM
		makeAscendingFunc(2), // PID
		makeAscendingFunc(3), // LOCAL_ADDR
		makeAscendingFunc(5), // USER
	).Sort(rows)

	for _, row := range rows {
		table.Append(row)
	}

	table.Render()
}

// ListPorts lists all ports that are being used.
func ListPorts(filter *Process, opts ...TransportProtocol) map[string]struct{} {
	rm := make(map[string]struct{})
	ps, _ := List(filter, opts...)
	for _, p := range ps {
		rm[p.LocalPort] = struct{}{}
		rm[p.RemotePort] = struct{}{}
	}
	return rm
}

type Ports struct {
	mu        sync.Mutex // guards the following
	beingUsed map[string]struct{}
}

func NewPorts() *Ports {
	p := &Ports{}
	p.beingUsed = make(map[string]struct{})
	return p
}

func (p *Ports) Exist(port string) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	_, ok := p.beingUsed[port]
	return ok
}

func (p *Ports) Add(ports ...string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, port := range ports {
		p.beingUsed[port] = struct{}{}
	}
}

func (p *Ports) Free(ports ...string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, port := range ports {
		delete(p.beingUsed, port)
	}
}

// Refresh refreshes ports that are being used by the OS.
// You would use as follows:
//
//	var (
//		globalPorts     = NewPorts()
//		refreshInterval = 10 * time.Second
//	)
//	func init() {
//		globalPorts.Refresh()
//		go func() {
//			for {
//				select {
//				case <-time.After(refreshInterval):
//					globalPorts.Refresh()
//				}
//			}
//		}()
//	}
//
func (p *Ports) Refresh() {
	p.mu.Lock()
	p.beingUsed = ListPorts(nil, TCP, TCP6)
	p.mu.Unlock()
}

func getFreePort(opts ...TransportProtocol) (string, error) {
	fp := ""
	var errMsg error
	for _, opt := range opts {
		protocolStr := ""
		switch opt {
		case TCP:
			protocolStr = "tcp"
		case TCP6:
			protocolStr = "tcp6"
		}
		addr, err := net.ResolveTCPAddr(protocolStr, "localhost:0")
		if err != nil {
			errMsg = err
			continue
		}
		l, err := net.ListenTCP(protocolStr, addr)
		if err != nil {
			errMsg = err
			continue
		}
		pd := l.Addr().(*net.TCPAddr).Port
		l.Close()
		fp = ":" + strconv.Itoa(pd)
		break
	}
	return fp, errMsg
}

// GetFreePorts returns multiple free ports from OS.
func GetFreePorts(num int, opts ...TransportProtocol) ([]string, error) {
	try := 0
	rm := make(map[string]struct{})
	for len(rm) != num {
		if try > 150 {
			return nil, fmt.Errorf("too many tries to find free ports")
		}
		try++
		p, err := getFreePort(opts...)
		if err != nil {
			return nil, err
		}
		rm[p] = struct{}{}
	}
	rs := []string{}
	for p := range rm {
		rs = append(rs, p)
	}
	sort.Strings(rs)
	return rs, nil
}
<|endoftext|>


<commit_before>package kuiperbelt

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"sync/atomic"
	"time"
)

// macopy may be embedded into structs which must not be copied
// after the first use.
// See https://github.com/golang/go/issues/8005#issuecomment-190753527
// for details.
type macopy struct{}

func (*macopy) Lock() {}

type Stats struct {
	connections        int64
	totalConnections   int64
	totalMessages      int64
	connectErrors      int64
	messageErrors      int64
	closingConnections int64
	noCopy             macopy
}

func NewStats() *Stats {
	return &Stats{}
}

func (s *Stats) Connections() int64 {
	return atomic.LoadInt64(&s.connections)
}

func (s *Stats) TotalConnections() int64 {
	return atomic.LoadInt64(&s.totalConnections)
}

func (s *Stats) TotalMessages() int64 {
	return atomic.LoadInt64(&s.totalMessages)
}

func (s *Stats) ConnectErrors() int64 {
	return atomic.LoadInt64(&s.connectErrors)
}

func (s *Stats) MessageErrors() int64 {
	return atomic.LoadInt64(&s.messageErrors)
}

func (s *Stats) ClosingConnections() int64 {
	return atomic.LoadInt64(&s.closingConnections)
}

func (s *Stats) Dump(w io.Writer) error {
	return json.NewEncoder(w).Encode(struct {
		Connections        int64 `json:"connections"`
		TotalConnections   int64 `json:"total_connections"`
		TotalMessages      int64 `json:"total_messages"`
		ConnectErrors      int64 `json:"connect_errors"`
		MessageErrors      int64 `json:"message_errors"`
		ClosingConnections int64 `json:"closing_connections"`
	}{
		Connections:        s.Connections(),
		TotalConnections:   s.TotalConnections(),
		TotalMessages:      s.TotalMessages(),
		ConnectErrors:      s.ConnectErrors(),
		MessageErrors:      s.MessageErrors(),
		ClosingConnections: s.ClosingConnections(),
	})
}

func (s *Stats) DumpText(w io.Writer) error {
	now := time.Now().Unix()
	buf := new(bytes.Buffer)
	buf.Grow(512)
	fmt.Fprintf(buf, "kuiperbelt.connections\t%d\t%d\n", s.Connections(), now)
	fmt.Fprintf(buf, "kuiperbelt.total_connections\t%d\t%d\n", s.TotalConnections(), now)
	fmt.Fprintf(buf, "kuiperbelt.total_messages\t%d\t%d\n", s.TotalMessages(), now)
	fmt.Fprintf(buf, "kuiperbelt.connect_errors\t%d\t%d\n", s.ConnectErrors(), now)
	fmt.Fprintf(buf, "kuiperbelt.message_errors\t%d\t%d\n", s.MessageErrors(), now)
	fmt.Fprintf(buf, "kuiperbelt.closing_connections\t%d\t%d\n", s.ClosingConnections(), now)
	_, err := buf.WriteTo(w)
	return err
}

func (s *Stats) ConnectEvent() {
	atomic.AddInt64(&s.totalConnections, 1)
	atomic.AddInt64(&s.connections, 1)
}

func (s *Stats) ConnectErrorEvent() {
	atomic.AddInt64(&s.connectErrors, 1)
}

func (s *Stats) DisconnectEvent() {
	atomic.AddInt64(&s.connections, -1)
}

func (s *Stats) MessageEvent() {
	atomic.AddInt64(&s.totalMessages, 1)
}

func (s *Stats) MessageErrorEvent() {
	atomic.AddInt64(&s.messageErrors, 1)
}

func (s *Stats) ClosingEvent() {
	atomic.AddInt64(&s.closingConnections, 1)
}

func (s *Stats) ClosedEvent() {
	atomic.AddInt64(&s.closingConnections, -1)
}
<commit_msg>rename stats key in format=text<commit_after>package kuiperbelt

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"sync/atomic"
	"time"
)

// macopy may be embedded into structs which must not be copied
// after the first use.
// See https://github.com/golang/go/issues/8005#issuecomment-190753527
// for details.
type macopy struct{}

func (*macopy) Lock() {}

type Stats struct {
	connections        int64
	totalConnections   int64
	totalMessages      int64
	connectErrors      int64
	messageErrors      int64
	closingConnections int64
	noCopy             macopy
}

func NewStats() *Stats {
	return &Stats{}
}

func (s *Stats) Connections() int64 {
	return atomic.LoadInt64(&s.connections)
}

func (s *Stats) TotalConnections() int64 {
	return atomic.LoadInt64(&s.totalConnections)
}

func (s *Stats) TotalMessages() int64 {
	return atomic.LoadInt64(&s.totalMessages)
}

func (s *Stats) ConnectErrors() int64 {
	return atomic.LoadInt64(&s.connectErrors)
}

func (s *Stats) MessageErrors() int64 {
	return atomic.LoadInt64(&s.messageErrors)
}

func (s *Stats) ClosingConnections() int64 {
	return atomic.LoadInt64(&s.closingConnections)
}

func (s *Stats) Dump(w io.Writer) error {
	return json.NewEncoder(w).Encode(struct {
		Connections        int64 `json:"connections"`
		TotalConnections   int64 `json:"total_connections"`
		TotalMessages      int64 `json:"total_messages"`
		ConnectErrors      int64 `json:"connect_errors"`
		MessageErrors      int64 `json:"message_errors"`
		ClosingConnections int64 `json:"closing_connections"`
	}{
		Connections:        s.Connections(),
		TotalConnections:   s.TotalConnections(),
		TotalMessages:      s.TotalMessages(),
		ConnectErrors:      s.ConnectErrors(),
		MessageErrors:      s.MessageErrors(),
		ClosingConnections: s.ClosingConnections(),
	})
}

func (s *Stats) DumpText(w io.Writer) error {
	now := time.Now().Unix()
	buf := new(bytes.Buffer)
	buf.Grow(512)
	fmt.Fprintf(buf, "kuiperbelt.conn.current\t%d\t%d\n", s.Connections(), now)
	fmt.Fprintf(buf, "kuiperbelt.conn.total\t%d\t%d\n", s.TotalConnections(), now)
	fmt.Fprintf(buf, "kuiperbelt.conn.errors\t%d\t%d\n", s.ConnectErrors(), now)
	fmt.Fprintf(buf, "kuiperbelt.conn.closing\t%d\t%d\n", s.ClosingConnections(), now)
	fmt.Fprintf(buf, "kuiperbelt.messages.total\t%d\t%d\n", s.TotalMessages(), now)
	fmt.Fprintf(buf, "kuiperbelt.messages.errors\t%d\t%d\n", s.MessageErrors(), now)
	_, err := buf.WriteTo(w)
	return err
}

func (s *Stats) ConnectEvent() {
	atomic.AddInt64(&s.totalConnections, 1)
	atomic.AddInt64(&s.connections, 1)
}

func (s *Stats) ConnectErrorEvent() {
	atomic.AddInt64(&s.connectErrors, 1)
}

func (s *Stats) DisconnectEvent() {
	atomic.AddInt64(&s.connections, -1)
}

func (s *Stats) MessageEvent() {
	atomic.AddInt64(&s.totalMessages, 1)
}

func (s *Stats) MessageErrorEvent() {
	atomic.AddInt64(&s.messageErrors, 1)
}

func (s *Stats) ClosingEvent() {
	atomic.AddInt64(&s.closingConnections, 1)
}

func (s *Stats) ClosedEvent() {
	atomic.AddInt64(&s.closingConnections, -1)
}
<|endoftext|>


<commit_before>package onecache

import (
	"bytes"
	"encoding/gob"
	"strconv"
	"time"
)

//Helper method to check if an item is expired.
//Current usecase for this is for garbage collection
func (i *Item) IsExpired() bool {
	return time.Now().After(i.ExpiresAt)
}

type Marshaller interface {
	UnMarshallBytes(data []byte) (*Item, error)
	MarshalBytes(i *Item) ([]byte, error)
}

func NewBytesItemMarshaller() *BytesItemMarshaller {
	return &BytesItemMarshaller{}
}

type BytesItemMarshaller struct {

}

func (b *BytesItemMarshaller) MarshalBytes(i *Item) ([]byte, error){
	var buf bytes.Buffer

	enc := gob.NewEncoder(&buf)

	if err := enc.Encode(i); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func (b *BytesItemMarshaller) UnMarshallBytes(data []byte) (*Item, error) {
	i := new(Item)

	err := gob.NewDecoder(bytes.NewBuffer(data)).Decode(i)

	return i, err
}

func Increment(val interface{}, steps int) (interface{}, error) {

	var ret interface{}

	switch val.(type) {

	case int:
		ret = val.(int) + steps

	case int32:
		ret = val.(int32) + int32(steps)

	case int64:
		ret = val.(int64) + int64(steps)

	case uint:
		ret = val.(uint) + uint(steps)

	case uint8:
		ret = val.(uint8) + uint8(steps)

	case uint16:
		ret = val.(uint16) + uint16(steps)

	case uint32:
		ret = val.(uint32) + uint32(steps)

	case uint64:
		ret = val.(uint64) + uint64(steps)

	case string:

		num, err := strconv.Atoi(val.(string))

		if err != nil {
			return -0, err
		}

		num += steps

		ret = strconv.Itoa(num)

	default:
		return -0, ErrCacheDataCannotBeIncreasedOrDecreased
	}

	return ret, nil

}

func Decrement(val interface{}, steps int) (interface{}, error) {
	return Increment(val, steps*-1)
}
<commit_msg>Named serializer properly and opened it up to work with "any" type.. This is allow client code make use of it<commit_after>package onecache

import (
	"bytes"
	"encoding/gob"
	"strconv"
	"time"
)

//Helper method to check if an item is expired.
//Current usecase for this is for garbage collection
func (i *Item) IsExpired() bool {
	return time.Now().After(i.ExpiresAt)
}

type Serializer interface {
	Serialize(i interface{}) ([]byte, error)
	DeSerialize(data []byte, i interface{}) error
}

func NewCacheSerializer() *CacheSerializer {
	return &CacheSerializer{}
}

//Helper to serialize and deserialize types
type CacheSerializer struct {
}

//Convert a given type into a byte array
//Caveat -> Types you create might have to be registered with the encoding/gob package
func (b *CacheSerializer) Serialize(i interface{}) ([]byte, error) {
	var buf bytes.Buffer

	enc := gob.NewEncoder(&buf)

	if err := enc.Encode(i); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

//Writes a byte array into a type.
func (b *CacheSerializer) DeSerialize(data []byte, i interface{}) error {
	return gob.NewDecoder(bytes.NewBuffer(data)).Decode(i)
}

func Increment(val interface{}, steps int) (interface{}, error) {

	var ret interface{}

	switch val.(type) {

	case int:
		ret = val.(int) + steps

	case int32:
		ret = val.(int32) + int32(steps)

	case int64:
		ret = val.(int64) + int64(steps)

	case uint:
		ret = val.(uint) + uint(steps)

	case uint8:
		ret = val.(uint8) + uint8(steps)

	case uint16:
		ret = val.(uint16) + uint16(steps)

	case uint32:
		ret = val.(uint32) + uint32(steps)

	case uint64:
		ret = val.(uint64) + uint64(steps)

	case string:

		num, err := strconv.Atoi(val.(string))

		if err != nil {
			return -0, err
		}

		num += steps

		ret = strconv.Itoa(num)

	default:
		return -0, ErrCacheDataCannotBeIncreasedOrDecreased
	}

	return ret, nil

}

func Decrement(val interface{}, steps int) (interface{}, error) {
	return Increment(val, steps*-1)
}
<|endoftext|>