<commit_before>package filterjson

import (
	"context"
	"strings"
	"testing"
	"time"

	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tsaikd/gogstash/config"
	"github.com/tsaikd/gogstash/config/goglog"
	"github.com/tsaikd/gogstash/config/logevent"
)

func init() {
	goglog.Logger.SetLevel(logrus.DebugLevel)
	config.RegistFilterHandler(ModuleName, InitHandler)
}

func Test_filter_json_module(t *testing.T) {
	assert := assert.New(t)
	assert.NotNil(assert)
	require := require.New(t)
	require.NotNil(require)

	ctx := context.Background()
	conf, err := config.LoadFromYAML([]byte(strings.TrimSpace(`
debugch: true
filter:
  - type: json
    message: message
    timestamp: time
    timeformat: "2006-01-02T15:04:05Z"
	`)))
	require.NoError(err)
	require.NoError(conf.Start(ctx))

	timestamp, err := time.Parse("2006-01-02T15:04:05Z", "2016-12-04T09:09:41.193Z")
	require.NoError(err)

	expectedEvent := logevent.LogEvent{
		Timestamp: timestamp,
		Message:   "Test",
		Extra: map[string]interface{}{
			"host": "Hostname",
		},
		Tags: []string{"foo"},
	}

	conf.TestInputEvent(logevent.LogEvent{
		Timestamp: time.Now(),
		Message:   "{ \"message\": \"Test\", \"host\": \"Hostname\", \"time\":\"2016-12-04T09:09:41.193Z\", \"tags\": [ \"foo\" ] }",
	})

	if event, err := conf.TestGetOutputEvent(300 * time.Millisecond); assert.NoError(err) {
		require.Equal(expectedEvent, event)
	}
}

func Test_filter_json_module_source(t *testing.T) {
	assert := assert.New(t)
	assert.NotNil(assert)
	require := require.New(t)
	require.NotNil(require)

	ctx := context.Background()
	conf, err := config.LoadFromYAML([]byte(strings.TrimSpace(`
debugch: true
filter:
  - type: json
    message: message
    timestamp: time
    source: myfield
    timeformat: "2006-01-02T15:04:05Z"
	`)))
	require.NoError(err)
	require.NoError(conf.Start(ctx))

	timestamp, err := time.Parse("2006-01-02T15:04:05Z", "2016-12-04T09:09:41.193Z")
	require.NoError(err)
	fieldvalue := `{ "message": "Test", "host":"Hostname", "time":"2016-12-04T09:09:41.193Z", "tags": [ "foo" ] }`

	expectedEvent := logevent.LogEvent{
		Timestamp: timestamp,
		Message:   "Test",
		Extra: map[string]interface{}{
			"host":    "Hostname",
			"myfield": fieldvalue,
		},
		Tags: []string{"foo"},
	}

	conf.TestInputEvent(logevent.LogEvent{
		Timestamp: time.Now(),
		Extra: map[string]interface{}{
			"myfield": fieldvalue,
		},
	})

	if event, err := conf.TestGetOutputEvent(300 * time.Millisecond); assert.NoError(err) {
		require.Equal(expectedEvent, event)
	}
}
<commit_msg>Add test cases<commit_after>package filterjson

import (
	"context"
	"strings"
	"testing"
	"time"

	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tsaikd/gogstash/config"
	"github.com/tsaikd/gogstash/config/goglog"
	"github.com/tsaikd/gogstash/config/logevent"
)

func init() {
	goglog.Logger.SetLevel(logrus.DebugLevel)
	config.RegistFilterHandler(ModuleName, InitHandler)
}

func Test_filter_json_module(t *testing.T) {
	assert := assert.New(t)
	assert.NotNil(assert)
	require := require.New(t)
	require.NotNil(require)

	ctx := context.Background()
	conf, err := config.LoadFromYAML([]byte(strings.TrimSpace(`
debugch: true
filter:
  - type: json
    message: message
    timestamp: time
    timeformat: "2006-01-02T15:04:05Z"
    remove_field: ["myfield"]
	`)))
	require.NoError(err)
	require.NoError(conf.Start(ctx))

	timestamp, err := time.Parse("2006-01-02T15:04:05Z", "2016-12-04T09:09:41.193Z")
	require.NoError(err)

	expectedEvent := logevent.LogEvent{
		Timestamp: timestamp,
		Message:   "Test",
		Extra: map[string]interface{}{
			"host": "Hostname",
		},
		Tags: []string{"foo"},
	}

	conf.TestInputEvent(logevent.LogEvent{
		Timestamp: time.Now(),
		Message:   "{ \"message\": \"Test\", \"host\": \"Hostname\", \"time\":\"2016-12-04T09:09:41.193Z\", \"tags\": [ \"foo\" ] }",
		Extra: map[string]interface{}{
			"myfield": "test",
		},
	})

	if event, err := conf.TestGetOutputEvent(300 * time.Millisecond); assert.NoError(err) {
		require.Equal(expectedEvent, event)
	}

	// not ok
	expectedEvent = logevent.LogEvent{
		Timestamp: timestamp,
		Message:   "Test",
		Extra: map[string]interface{}{
			"myfield": "test",
		},
		Tags: []string{"gogstash_filter_json_error"},
	}

	conf.TestInputEvent(logevent.LogEvent{
		Timestamp: timestamp,
		Message:   "Test",
		Extra: map[string]interface{}{
			"myfield": "test",
		},
	})

	if event, err := conf.TestGetOutputEvent(300 * time.Millisecond); assert.NoError(err) {
		require.Equal(expectedEvent, event)
	}
}

func Test_filter_json_module_source(t *testing.T) {
	assert := assert.New(t)
	assert.NotNil(assert)
	require := require.New(t)
	require.NotNil(require)

	ctx := context.Background()
	conf, err := config.LoadFromYAML([]byte(strings.TrimSpace(`
debugch: true
filter:
  - type: json
    message: message
    timestamp: time
    source: myfield
    timeformat: "2006-01-02T15:04:05Z"
	`)))
	require.NoError(err)
	require.NoError(conf.Start(ctx))

	timestamp, err := time.Parse("2006-01-02T15:04:05Z", "2016-12-04T09:09:41.193Z")
	require.NoError(err)
	fieldvalue := `{ "message": "Test", "host":"Hostname", "time":"2016-12-04T09:09:41.193Z", "tags": [ "foo" ] }`

	expectedEvent := logevent.LogEvent{
		Timestamp: timestamp,
		Message:   "Test",
		Extra: map[string]interface{}{
			"host":    "Hostname",
			"myfield": fieldvalue,
		},
		Tags: []string{"foo"},
	}

	conf.TestInputEvent(logevent.LogEvent{
		Timestamp: time.Now(),
		Extra: map[string]interface{}{
			"myfield": fieldvalue,
		},
	})

	if event, err := conf.TestGetOutputEvent(300 * time.Millisecond); assert.NoError(err) {
		require.Equal(expectedEvent, event)
	}
}
<|endoftext|>


<commit_before>// 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.

// Package norm contains types and functions for normalizing Unicode strings.
package norm

import "utf8"

// A Form denotes a canonical representation of Unicode code points.
// The Unicode-defined normalization and equivalence forms are:
//
//   NFC   Unicode Normalization Form C
//   NFD   Unicode Normalization Form D
//   NFKC  Unicode Normalization Form KC
//   NFKD  Unicode Normalization Form KD
//
// For a Form f, this documentation uses the notation f(x) to mean
// the bytes or string x converted to the given form.
// A position n in x is called a boundary if conversion to the form can
// proceed independently on both sides:
//   f(x) == append(f(x[0:n]), f(x[n:])...)
//
// References: http://unicode.org/reports/tr15/ and
// http://unicode.org/notes/tn5/.
type Form int

const (
	NFC Form = iota
	NFD
	NFKC
	NFKD
)

// Bytes returns f(b). May return b if f(b) = b.
func (f Form) Bytes(b []byte) []byte {
	n := f.QuickSpan(b)
	if n == len(b) {
		return b
	}
	out := make([]byte, n, len(b))
	copy(out, b[0:n])
	return f.Append(out, b[n:]...)
}

// String returns f(s).
func (f Form) String(s string) string {
	n := f.QuickSpanString(s)
	if n == len(s) {
		return s
	}
	out := make([]byte, 0, len(s))
	copy(out, s[0:n])
	return string(f.AppendString(out, s[n:]))
}

// IsNormal returns true if b == f(b).
func (f Form) IsNormal(b []byte) bool {
	rb := reorderBuffer{}
	rb.init(f, b)
	bp := quickSpan(&rb, 0)
	if bp == len(b) {
		return true
	}
	for bp < len(b) {
		decomposeSegment(&rb, bp)
		if rb.f.composing {
			rb.compose()
		}
		for i := 0; i < rb.nrune; i++ {
			info := rb.rune[i]
			if bp+int(info.size) > len(b) {
				return false
			}
			p := info.pos
			pe := p + info.size
			for ; p < pe; p++ {
				if b[bp] != rb.byte[p] {
					return false
				}
				bp++
			}
		}
		rb.reset()
		bp = quickSpan(&rb, bp)
	}
	return true
}

// IsNormalString returns true if s == f(s).
func (f Form) IsNormalString(s string) bool {
	rb := reorderBuffer{}
	rb.initString(f, s)
	bp := quickSpan(&rb, 0)
	if bp == len(s) {
		return true
	}
	for bp < len(s) {
		decomposeSegment(&rb, bp)
		if rb.f.composing {
			rb.compose()
		}
		for i := 0; i < rb.nrune; i++ {
			info := rb.rune[i]
			if bp+int(info.size) > len(s) {
				return false
			}
			p := info.pos
			pe := p + info.size
			for ; p < pe; p++ {
				if s[bp] != rb.byte[p] {
					return false
				}
				bp++
			}
		}
		rb.reset()
		bp = quickSpan(&rb, bp)
	}
	return true
}

// patchTail fixes a case where a rune may be incorrectly normalized
// if it is followed by illegal continuation bytes. It returns the
// patched buffer and the number of trailing continuation bytes that
// have been dropped.
func patchTail(rb *reorderBuffer, buf []byte) ([]byte, int) {
	info, p := lastRuneStart(&rb.f, buf)
	if p == -1 || info.size == 0 {
		return buf, 0
	}
	end := p + int(info.size)
	extra := len(buf) - end
	if extra > 0 {
		buf = decomposeToLastBoundary(rb, buf[:end])
		if rb.f.composing {
			rb.compose()
		}
		return rb.flush(buf), extra
	}
	return buf, 0
}

func appendQuick(rb *reorderBuffer, dst []byte, i int) ([]byte, int) {
	if rb.nsrc == i {
		return dst, i
	}
	end := quickSpan(rb, i)
	return rb.src.appendSlice(dst, i, end), end
}

// Append returns f(append(out, b...)).
// The buffer out must be nil, empty, or equal to f(out).
func (f Form) Append(out []byte, src ...byte) []byte {
	if len(src) == 0 {
		return out
	}
	rb := reorderBuffer{}
	rb.init(f, src)
	return doAppend(&rb, out)
}

func doAppend(rb *reorderBuffer, out []byte) []byte {
	src, n := rb.src, rb.nsrc
	doMerge := len(out) > 0
	p := 0
	if p = src.skipNonStarter(); p > 0 {
		// Move leading non-starters to destination.
		out = src.appendSlice(out, 0, p)
		buf, ndropped := patchTail(rb, out)
		if ndropped > 0 {
			out = src.appendSlice(buf, p-ndropped, p)
			doMerge = false // no need to merge, ends with illegal UTF-8
		} else {
			out = decomposeToLastBoundary(rb, buf) // force decomposition
		}
	}
	fd := &rb.f
	if doMerge {
		var info runeInfo
		if p < n {
			info = fd.info(src, p)
			if p == 0 && !fd.boundaryBefore(fd, info) {
				out = decomposeToLastBoundary(rb, out)
			}
		}
		if info.size == 0 || fd.boundaryBefore(fd, info) {
			if fd.composing {
				rb.compose()
			}
			out = rb.flush(out)
			if info.size == 0 {
				// Append incomplete UTF-8 encoding.
				return src.appendSlice(out, p, n)
			}
		}
	}
	if rb.nrune == 0 {
		out, p = appendQuick(rb, out, p)
	}
	for p < n {
		p = decomposeSegment(rb, p)
		if fd.composing {
			rb.compose()
		}
		out = rb.flush(out)
		out, p = appendQuick(rb, out, p)
	}
	return out
}

// AppendString returns f(append(out, []byte(s))).
// The buffer out must be nil, empty, or equal to f(out).
func (f Form) AppendString(out []byte, src string) []byte {
	if len(src) == 0 {
		return out
	}
	rb := reorderBuffer{}
	rb.initString(f, src)
	return doAppend(&rb, out)
}

// QuickSpan returns a boundary n such that b[0:n] == f(b[0:n]).
// It is not guaranteed to return the largest such n.
func (f Form) QuickSpan(b []byte) int {
	rb := reorderBuffer{}
	rb.init(f, b)
	return quickSpan(&rb, 0)
}

func quickSpan(rb *reorderBuffer, i int) int {
	var lastCC uint8
	var nc int
	lastSegStart := i
	src, n := rb.src, rb.nsrc
	for i < n {
		if j := src.skipASCII(i); i != j {
			i = j
			lastSegStart = i - 1
			lastCC = 0
			nc = 0
			continue
		}
		info := rb.f.info(src, i)
		if info.size == 0 {
			// include incomplete runes
			return n
		}
		cc := info.ccc
		if rb.f.composing {
			if !info.flags.isYesC() {
				break
			}
		} else {
			if !info.flags.isYesD() {
				break
			}
		}
		if cc == 0 {
			lastSegStart = i
			nc = 0
		} else {
			if nc >= maxCombiningChars {
				lastSegStart = i
				lastCC = cc
				nc = 1
			} else {
				if lastCC > cc {
					return lastSegStart
				}
				nc++
			}
		}
		lastCC = cc
		i += int(info.size)
	}
	if i == n {
		return n
	}
	if rb.f.composing {
		return lastSegStart
	}
	return i
}

// QuickSpanString returns a boundary n such that b[0:n] == f(s[0:n]).
// It is not guaranteed to return the largest such n.
func (f Form) QuickSpanString(s string) int {
	rb := reorderBuffer{}
	rb.initString(f, s)
	return quickSpan(&rb, 0)
}

// FirstBoundary returns the position i of the first boundary in b
// or -1 if b contains no boundary.
func (f Form) FirstBoundary(b []byte) int {
	rb := reorderBuffer{}
	rb.init(f, b)
	return firstBoundary(&rb)
}

func firstBoundary(rb *reorderBuffer) int {
	src, nsrc := rb.src, rb.nsrc
	i := src.skipNonStarter()
	if i >= nsrc {
		return -1
	}
	fd := &rb.f
	info := fd.info(src, i)
	for n := 0; info.size != 0 && !fd.boundaryBefore(fd, info); {
		i += int(info.size)
		if n++; n >= maxCombiningChars {
			return i
		}
		if i >= nsrc {
			if !fd.boundaryAfter(fd, info) {
				return -1
			}
			return nsrc
		}
		info = fd.info(src, i)
	}
	if info.size == 0 {
		return -1
	}
	return i
}

// FirstBoundaryInString returns the position i of the first boundary in s
// or -1 if s contains no boundary.
func (f Form) FirstBoundaryInString(s string) int {
	rb := reorderBuffer{}
	rb.initString(f, s)
	return firstBoundary(&rb)
}

// LastBoundary returns the position i of the last boundary in b
// or -1 if b contains no boundary.
func (f Form) LastBoundary(b []byte) int {
	return lastBoundary(formTable[f], b)
}

func lastBoundary(fd *formInfo, b []byte) int {
	i := len(b)
	info, p := lastRuneStart(fd, b)
	if p == -1 {
		return -1
	}
	if info.size == 0 { // ends with incomplete rune
		if p == 0 { // starts wtih incomplete rune
			return -1
		}
		i = p
		info, p = lastRuneStart(fd, b[:i])
		if p == -1 { // incomplete UTF-8 encoding or non-starter bytes without a starter
			return i
		}
	}
	if p+int(info.size) != i { // trailing non-starter bytes: illegal UTF-8
		return i
	}
	if fd.boundaryAfter(fd, info) {
		return i
	}
	i = p
	for n := 0; i >= 0 && !fd.boundaryBefore(fd, info); {
		info, p = lastRuneStart(fd, b[:i])
		if n++; n >= maxCombiningChars {
			return len(b)
		}
		if p+int(info.size) != i {
			if p == -1 { // no boundary found
				return -1
			}
			return i // boundary after an illegal UTF-8 encoding
		}
		i = p
	}
	return i
}

// LastBoundaryInString returns the position i of the last boundary in s
// or -1 if s contains no boundary.
func (f Form) LastBoundaryInString(s string) int {
	panic("not implemented")
}

// decomposeSegment scans the first segment in src into rb.
// It returns the number of bytes consumed from src.
// TODO(mpvl): consider inserting U+034f (Combining Grapheme Joiner)
// when we detect a sequence of 30+ non-starter chars.
func decomposeSegment(rb *reorderBuffer, sp int) int {
	// Force one character to be consumed.
	info := rb.f.info(rb.src, sp)
	if info.size == 0 {
		return 0
	}
	for rb.insert(rb.src, sp, info) {
		sp += int(info.size)
		if sp >= rb.nsrc {
			break
		}
		info = rb.f.info(rb.src, sp)
		bound := rb.f.boundaryBefore(&rb.f, info)
		if bound || info.size == 0 {
			break
		}
	}
	return sp
}

// lastRuneStart returns the runeInfo and position of the last
// rune in buf or the zero runeInfo and -1 if no rune was found.
func lastRuneStart(fd *formInfo, buf []byte) (runeInfo, int) {
	p := len(buf) - 1
	for ; p >= 0 && !utf8.RuneStart(buf[p]); p-- {
	}
	if p < 0 {
		return runeInfo{0, 0, 0, 0}, -1
	}
	return fd.info(inputBytes(buf), p), p
}

// decomposeToLastBoundary finds an open segment at the end of the buffer
// and scans it into rb. Returns the buffer minus the last segment.
func decomposeToLastBoundary(rb *reorderBuffer, buf []byte) []byte {
	fd := &rb.f
	info, i := lastRuneStart(fd, buf)
	if int(info.size) != len(buf)-i {
		// illegal trailing continuation bytes
		return buf
	}
	if rb.f.boundaryAfter(fd, info) {
		return buf
	}
	var add [maxBackRunes]runeInfo // stores runeInfo in reverse order
	add[0] = info
	padd := 1
	n := 1
	p := len(buf) - int(info.size)
	for ; p >= 0 && !rb.f.boundaryBefore(fd, info); p -= int(info.size) {
		info, i = lastRuneStart(fd, buf[:p])
		if int(info.size) != p-i {
			break
		}
		// Check that decomposition doesn't result in overflow.
		if info.flags.hasDecomposition() {
			dcomp := rb.f.decompose(inputBytes(buf), p-int(info.size))
			for i := 0; i < len(dcomp); {
				inf := rb.f.info(inputBytes(dcomp), i)
				i += int(inf.size)
				n++
			}
		} else {
			n++
		}
		if n > maxBackRunes {
			break
		}
		add[padd] = info
		padd++
	}
	pp := p
	for padd--; padd >= 0; padd-- {
		info = add[padd]
		rb.insert(inputBytes(buf), pp, info)
		pp += int(info.size)
	}
	return buf[:p]
}
<commit_msg>exp/norm: LastBoundary is used in preparation for an append operation. It seems therefore unlikely that there is a good use for its string version LastBoundaryInString. Yet, the implemenation of this method would complicate things a bit as it would require the introduction for another interface and some duplication of code. Removing it seems a better choice.<commit_after>// 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.

// Package norm contains types and functions for normalizing Unicode strings.
package norm

import "utf8"

// A Form denotes a canonical representation of Unicode code points.
// The Unicode-defined normalization and equivalence forms are:
//
//   NFC   Unicode Normalization Form C
//   NFD   Unicode Normalization Form D
//   NFKC  Unicode Normalization Form KC
//   NFKD  Unicode Normalization Form KD
//
// For a Form f, this documentation uses the notation f(x) to mean
// the bytes or string x converted to the given form.
// A position n in x is called a boundary if conversion to the form can
// proceed independently on both sides:
//   f(x) == append(f(x[0:n]), f(x[n:])...)
//
// References: http://unicode.org/reports/tr15/ and
// http://unicode.org/notes/tn5/.
type Form int

const (
	NFC Form = iota
	NFD
	NFKC
	NFKD
)

// Bytes returns f(b). May return b if f(b) = b.
func (f Form) Bytes(b []byte) []byte {
	n := f.QuickSpan(b)
	if n == len(b) {
		return b
	}
	out := make([]byte, n, len(b))
	copy(out, b[0:n])
	return f.Append(out, b[n:]...)
}

// String returns f(s).
func (f Form) String(s string) string {
	n := f.QuickSpanString(s)
	if n == len(s) {
		return s
	}
	out := make([]byte, 0, len(s))
	copy(out, s[0:n])
	return string(f.AppendString(out, s[n:]))
}

// IsNormal returns true if b == f(b).
func (f Form) IsNormal(b []byte) bool {
	rb := reorderBuffer{}
	rb.init(f, b)
	bp := quickSpan(&rb, 0)
	if bp == len(b) {
		return true
	}
	for bp < len(b) {
		decomposeSegment(&rb, bp)
		if rb.f.composing {
			rb.compose()
		}
		for i := 0; i < rb.nrune; i++ {
			info := rb.rune[i]
			if bp+int(info.size) > len(b) {
				return false
			}
			p := info.pos
			pe := p + info.size
			for ; p < pe; p++ {
				if b[bp] != rb.byte[p] {
					return false
				}
				bp++
			}
		}
		rb.reset()
		bp = quickSpan(&rb, bp)
	}
	return true
}

// IsNormalString returns true if s == f(s).
func (f Form) IsNormalString(s string) bool {
	rb := reorderBuffer{}
	rb.initString(f, s)
	bp := quickSpan(&rb, 0)
	if bp == len(s) {
		return true
	}
	for bp < len(s) {
		decomposeSegment(&rb, bp)
		if rb.f.composing {
			rb.compose()
		}
		for i := 0; i < rb.nrune; i++ {
			info := rb.rune[i]
			if bp+int(info.size) > len(s) {
				return false
			}
			p := info.pos
			pe := p + info.size
			for ; p < pe; p++ {
				if s[bp] != rb.byte[p] {
					return false
				}
				bp++
			}
		}
		rb.reset()
		bp = quickSpan(&rb, bp)
	}
	return true
}

// patchTail fixes a case where a rune may be incorrectly normalized
// if it is followed by illegal continuation bytes. It returns the
// patched buffer and the number of trailing continuation bytes that
// have been dropped.
func patchTail(rb *reorderBuffer, buf []byte) ([]byte, int) {
	info, p := lastRuneStart(&rb.f, buf)
	if p == -1 || info.size == 0 {
		return buf, 0
	}
	end := p + int(info.size)
	extra := len(buf) - end
	if extra > 0 {
		buf = decomposeToLastBoundary(rb, buf[:end])
		if rb.f.composing {
			rb.compose()
		}
		return rb.flush(buf), extra
	}
	return buf, 0
}

func appendQuick(rb *reorderBuffer, dst []byte, i int) ([]byte, int) {
	if rb.nsrc == i {
		return dst, i
	}
	end := quickSpan(rb, i)
	return rb.src.appendSlice(dst, i, end), end
}

// Append returns f(append(out, b...)).
// The buffer out must be nil, empty, or equal to f(out).
func (f Form) Append(out []byte, src ...byte) []byte {
	if len(src) == 0 {
		return out
	}
	rb := reorderBuffer{}
	rb.init(f, src)
	return doAppend(&rb, out)
}

func doAppend(rb *reorderBuffer, out []byte) []byte {
	src, n := rb.src, rb.nsrc
	doMerge := len(out) > 0
	p := 0
	if p = src.skipNonStarter(); p > 0 {
		// Move leading non-starters to destination.
		out = src.appendSlice(out, 0, p)
		buf, ndropped := patchTail(rb, out)
		if ndropped > 0 {
			out = src.appendSlice(buf, p-ndropped, p)
			doMerge = false // no need to merge, ends with illegal UTF-8
		} else {
			out = decomposeToLastBoundary(rb, buf) // force decomposition
		}
	}
	fd := &rb.f
	if doMerge {
		var info runeInfo
		if p < n {
			info = fd.info(src, p)
			if p == 0 && !fd.boundaryBefore(fd, info) {
				out = decomposeToLastBoundary(rb, out)
			}
		}
		if info.size == 0 || fd.boundaryBefore(fd, info) {
			if fd.composing {
				rb.compose()
			}
			out = rb.flush(out)
			if info.size == 0 {
				// Append incomplete UTF-8 encoding.
				return src.appendSlice(out, p, n)
			}
		}
	}
	if rb.nrune == 0 {
		out, p = appendQuick(rb, out, p)
	}
	for p < n {
		p = decomposeSegment(rb, p)
		if fd.composing {
			rb.compose()
		}
		out = rb.flush(out)
		out, p = appendQuick(rb, out, p)
	}
	return out
}

// AppendString returns f(append(out, []byte(s))).
// The buffer out must be nil, empty, or equal to f(out).
func (f Form) AppendString(out []byte, src string) []byte {
	if len(src) == 0 {
		return out
	}
	rb := reorderBuffer{}
	rb.initString(f, src)
	return doAppend(&rb, out)
}

// QuickSpan returns a boundary n such that b[0:n] == f(b[0:n]).
// It is not guaranteed to return the largest such n.
func (f Form) QuickSpan(b []byte) int {
	rb := reorderBuffer{}
	rb.init(f, b)
	return quickSpan(&rb, 0)
}

func quickSpan(rb *reorderBuffer, i int) int {
	var lastCC uint8
	var nc int
	lastSegStart := i
	src, n := rb.src, rb.nsrc
	for i < n {
		if j := src.skipASCII(i); i != j {
			i = j
			lastSegStart = i - 1
			lastCC = 0
			nc = 0
			continue
		}
		info := rb.f.info(src, i)
		if info.size == 0 {
			// include incomplete runes
			return n
		}
		cc := info.ccc
		if rb.f.composing {
			if !info.flags.isYesC() {
				break
			}
		} else {
			if !info.flags.isYesD() {
				break
			}
		}
		if cc == 0 {
			lastSegStart = i
			nc = 0
		} else {
			if nc >= maxCombiningChars {
				lastSegStart = i
				lastCC = cc
				nc = 1
			} else {
				if lastCC > cc {
					return lastSegStart
				}
				nc++
			}
		}
		lastCC = cc
		i += int(info.size)
	}
	if i == n {
		return n
	}
	if rb.f.composing {
		return lastSegStart
	}
	return i
}

// QuickSpanString returns a boundary n such that b[0:n] == f(s[0:n]).
// It is not guaranteed to return the largest such n.
func (f Form) QuickSpanString(s string) int {
	rb := reorderBuffer{}
	rb.initString(f, s)
	return quickSpan(&rb, 0)
}

// FirstBoundary returns the position i of the first boundary in b
// or -1 if b contains no boundary.
func (f Form) FirstBoundary(b []byte) int {
	rb := reorderBuffer{}
	rb.init(f, b)
	return firstBoundary(&rb)
}

func firstBoundary(rb *reorderBuffer) int {
	src, nsrc := rb.src, rb.nsrc
	i := src.skipNonStarter()
	if i >= nsrc {
		return -1
	}
	fd := &rb.f
	info := fd.info(src, i)
	for n := 0; info.size != 0 && !fd.boundaryBefore(fd, info); {
		i += int(info.size)
		if n++; n >= maxCombiningChars {
			return i
		}
		if i >= nsrc {
			if !fd.boundaryAfter(fd, info) {
				return -1
			}
			return nsrc
		}
		info = fd.info(src, i)
	}
	if info.size == 0 {
		return -1
	}
	return i
}

// FirstBoundaryInString returns the position i of the first boundary in s
// or -1 if s contains no boundary.
func (f Form) FirstBoundaryInString(s string) int {
	rb := reorderBuffer{}
	rb.initString(f, s)
	return firstBoundary(&rb)
}

// LastBoundary returns the position i of the last boundary in b
// or -1 if b contains no boundary.
func (f Form) LastBoundary(b []byte) int {
	return lastBoundary(formTable[f], b)
}

func lastBoundary(fd *formInfo, b []byte) int {
	i := len(b)
	info, p := lastRuneStart(fd, b)
	if p == -1 {
		return -1
	}
	if info.size == 0 { // ends with incomplete rune
		if p == 0 { // starts wtih incomplete rune
			return -1
		}
		i = p
		info, p = lastRuneStart(fd, b[:i])
		if p == -1 { // incomplete UTF-8 encoding or non-starter bytes without a starter
			return i
		}
	}
	if p+int(info.size) != i { // trailing non-starter bytes: illegal UTF-8
		return i
	}
	if fd.boundaryAfter(fd, info) {
		return i
	}
	i = p
	for n := 0; i >= 0 && !fd.boundaryBefore(fd, info); {
		info, p = lastRuneStart(fd, b[:i])
		if n++; n >= maxCombiningChars {
			return len(b)
		}
		if p+int(info.size) != i {
			if p == -1 { // no boundary found
				return -1
			}
			return i // boundary after an illegal UTF-8 encoding
		}
		i = p
	}
	return i
}

// decomposeSegment scans the first segment in src into rb.
// It returns the number of bytes consumed from src.
// TODO(mpvl): consider inserting U+034f (Combining Grapheme Joiner)
// when we detect a sequence of 30+ non-starter chars.
func decomposeSegment(rb *reorderBuffer, sp int) int {
	// Force one character to be consumed.
	info := rb.f.info(rb.src, sp)
	if info.size == 0 {
		return 0
	}
	for rb.insert(rb.src, sp, info) {
		sp += int(info.size)
		if sp >= rb.nsrc {
			break
		}
		info = rb.f.info(rb.src, sp)
		bound := rb.f.boundaryBefore(&rb.f, info)
		if bound || info.size == 0 {
			break
		}
	}
	return sp
}

// lastRuneStart returns the runeInfo and position of the last
// rune in buf or the zero runeInfo and -1 if no rune was found.
func lastRuneStart(fd *formInfo, buf []byte) (runeInfo, int) {
	p := len(buf) - 1
	for ; p >= 0 && !utf8.RuneStart(buf[p]); p-- {
	}
	if p < 0 {
		return runeInfo{0, 0, 0, 0}, -1
	}
	return fd.info(inputBytes(buf), p), p
}

// decomposeToLastBoundary finds an open segment at the end of the buffer
// and scans it into rb. Returns the buffer minus the last segment.
func decomposeToLastBoundary(rb *reorderBuffer, buf []byte) []byte {
	fd := &rb.f
	info, i := lastRuneStart(fd, buf)
	if int(info.size) != len(buf)-i {
		// illegal trailing continuation bytes
		return buf
	}
	if rb.f.boundaryAfter(fd, info) {
		return buf
	}
	var add [maxBackRunes]runeInfo // stores runeInfo in reverse order
	add[0] = info
	padd := 1
	n := 1
	p := len(buf) - int(info.size)
	for ; p >= 0 && !rb.f.boundaryBefore(fd, info); p -= int(info.size) {
		info, i = lastRuneStart(fd, buf[:p])
		if int(info.size) != p-i {
			break
		}
		// Check that decomposition doesn't result in overflow.
		if info.flags.hasDecomposition() {
			dcomp := rb.f.decompose(inputBytes(buf), p-int(info.size))
			for i := 0; i < len(dcomp); {
				inf := rb.f.info(inputBytes(dcomp), i)
				i += int(inf.size)
				n++
			}
		} else {
			n++
		}
		if n > maxBackRunes {
			break
		}
		add[padd] = info
		padd++
	}
	pp := p
	for padd--; padd >= 0; padd-- {
		info = add[padd]
		rb.insert(inputBytes(buf), pp, info)
		pp += int(info.size)
	}
	return buf[:p]
}
<|endoftext|>


<commit_before>// Copyright 2012 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.

package net_test

import (
	"bytes"
	"net"
	"os"
	"runtime"
	"testing"
	"time"
)

func TestUDPConnSpecificMethods(t *testing.T) {
	la, err := net.ResolveUDPAddr("udp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("net.ResolveUDPAddr failed: %v", err)
	}
	c, err := net.ListenUDP("udp4", la)
	if err != nil {
		t.Fatalf("net.ListenUDP failed: %v", err)
	}
	c.File()
	c.LocalAddr()
	c.RemoteAddr()
	c.SetDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetWriteDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetReadBuffer(2048)
	c.SetWriteBuffer(2048)
	defer c.Close()

	wb := []byte("UDPCONN TEST")
	if _, err := c.WriteToUDP(wb, c.LocalAddr().(*net.UDPAddr)); err != nil {
		t.Fatalf("net.UDPConn.WriteToUDP failed: %v", err)
	}
	rb := make([]byte, 128)
	if _, _, err := c.ReadFromUDP(rb); err != nil {
		t.Fatalf("net.UDPConn.ReadFromUDP failed: %v", err)
	}
}

func TestIPConnSpecificMethods(t *testing.T) {
	switch runtime.GOOS {
	case "plan9":
		t.Logf("skipping read test on %q", runtime.GOOS)
		return
	}
	if os.Getuid() != 0 {
		t.Logf("skipping test; must be root")
		return
	}

	la, err := net.ResolveIPAddr("ip4", "127.0.0.1")
	if err != nil {
		t.Fatalf("net.ResolveIPAddr failed: %v", err)
	}
	c, err := net.ListenIP("ip4:icmp", la)
	if err != nil {
		t.Fatalf("net.ListenIP failed: %v", err)
	}
	c.File()
	c.LocalAddr()
	c.RemoteAddr()
	c.SetDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetWriteDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetReadBuffer(2048)
	c.SetWriteBuffer(2048)
	defer c.Close()

	id := os.Getpid() & 0xffff
	wb := newICMPEchoRequest(id, 1, 128, []byte("IPCONN TEST "))
	if _, err := c.WriteToIP(wb, c.LocalAddr().(*net.IPAddr)); err != nil {
		t.Fatalf("net.IPConn.WriteToIP failed: %v", err)
	}
	rb := make([]byte, 20+128)
	if _, _, err := c.ReadFromIP(rb); err != nil {
		t.Fatalf("net.IPConn.ReadFromIP failed: %v", err)
	}
}

// TODO: Find out the use case of ListenUnixgram, I have no idea.
func TestUnixConnSpecificMethods(t *testing.T) {
	switch runtime.GOOS {
	case "plan9", "windows":
		t.Logf("skipping test on %q", runtime.GOOS)
		return
	}

	p1, p2 := "/tmp/gotest.net1", "/tmp/gotest.net2"
	os.Remove(p1)
	os.Remove(p2)

	a1, err := net.ResolveUnixAddr("unixgram", p1)
	if err != nil {
		t.Fatalf("net.ResolveUnixAddr failed: %v", err)
	}
	c1, err := net.DialUnix("unixgram", a1, nil)
	if err != nil {
		t.Fatalf("net.DialUnix failed: %v", err)
	}
	c1.File()
	c1.LocalAddr()
	c1.RemoteAddr()
	c1.SetDeadline(time.Now().Add(100 * time.Millisecond))
	c1.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	c1.SetWriteDeadline(time.Now().Add(100 * time.Millisecond))
	c1.SetReadBuffer(2048)
	c1.SetWriteBuffer(2048)
	defer c1.Close()
	defer os.Remove(p1)

	a2, err := net.ResolveUnixAddr("unixgram", p2)
	if err != nil {
		t.Fatalf("net.ResolveUnixAddr failed: %v", err)
	}
	c2, err := net.DialUnix("unixgram", a2, nil)
	if err != nil {
		t.Fatalf("net.DialUnix failed: %v", err)
	}
	c2.File()
	c2.LocalAddr()
	c2.RemoteAddr()
	c2.SetDeadline(time.Now().Add(100 * time.Millisecond))
	c2.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	c2.SetWriteDeadline(time.Now().Add(100 * time.Millisecond))
	c2.SetReadBuffer(2048)
	c2.SetWriteBuffer(2048)
	defer c2.Close()
	defer os.Remove(p2)

	wb := []byte("UNIXCONN TEST")
	if _, _, err := c1.WriteMsgUnix(wb, nil, a2); err != nil {
		t.Fatalf("net.UnixConn.WriteMsgUnix failed: %v", err)
	}
	rb2 := make([]byte, 128)
	if _, _, _, _, err := c2.ReadMsgUnix(rb2, nil); err != nil {
		t.Fatalf("net.UnixConn.ReadMsgUnix failed: %v", err)
	}
	if _, err := c2.WriteToUnix(wb, a1); err != nil {
		t.Fatalf("net.UnixConn.WriteToUnix failed: %v", err)
	}
	rb1 := make([]byte, 128)
	if _, _, err := c1.ReadFromUnix(rb1); err != nil {
		t.Fatalf("net.UnixConn.ReadFromUnix failed: %v", err)
	}
}

func newICMPEchoRequest(id, seqnum, msglen int, filler []byte) []byte {
	b := newICMPInfoMessage(id, seqnum, msglen, filler)
	b[0] = 8
	// calculate ICMP checksum
	cklen := len(b)
	s := uint32(0)
	for i := 0; i < cklen-1; i += 2 {
		s += uint32(b[i+1])<<8 | uint32(b[i])
	}
	if cklen&1 == 1 {
		s += uint32(b[cklen-1])
	}
	s = (s >> 16) + (s & 0xffff)
	s = s + (s >> 16)
	// place checksum back in header; using ^= avoids the
	// assumption the checksum bytes are zero
	b[2] ^= byte(^s & 0xff)
	b[3] ^= byte(^s >> 8)
	return b
}

func newICMPInfoMessage(id, seqnum, msglen int, filler []byte) []byte {
	b := make([]byte, msglen)
	copy(b[8:], bytes.Repeat(filler, (msglen-8)/len(filler)+1))
	b[0] = 0                   // type
	b[1] = 0                   // code
	b[2] = 0                   // checksum
	b[3] = 0                   // checksum
	b[4] = byte(id >> 8)       // identifier
	b[5] = byte(id & 0xff)     // identifier
	b[6] = byte(seqnum >> 8)   // sequence number
	b[7] = byte(seqnum & 0xff) // sequence number
	return b
}
<commit_msg>net: add more tests for protocol specific methods<commit_after>// Copyright 2012 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.

package net_test

import (
	"bytes"
	"net"
	"os"
	"runtime"
	"testing"
	"time"
)

var condErrorf = func() func(*testing.T, string, ...interface{}) {
	// A few APIs are not implemented yet on both Plan 9 and Windows.
	switch runtime.GOOS {
	case "plan9", "windows":
		return (*testing.T).Logf
	}
	return (*testing.T).Errorf
}()

func TestTCPListenerSpecificMethods(t *testing.T) {
	la, err := net.ResolveTCPAddr("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("net.ResolveTCPAddr failed: %v", err)
	}
	ln, err := net.ListenTCP("tcp4", la)
	if err != nil {
		t.Fatalf("net.ListenTCP failed: %v", err)
	}
	ln.Addr()
	ln.SetDeadline(time.Now().Add(30 * time.Nanosecond))
	defer ln.Close()

	if c, err := ln.Accept(); err != nil {
		if !err.(net.Error).Timeout() {
			t.Errorf("net.TCPListener.Accept failed: %v", err)
			return
		}
	} else {
		c.Close()
	}
	if c, err := ln.AcceptTCP(); err != nil {
		if !err.(net.Error).Timeout() {
			t.Errorf("net.TCPListener.AcceptTCP failed: %v", err)
			return
		}
	} else {
		c.Close()
	}

	if f, err := ln.File(); err != nil {
		condErrorf(t, "net.TCPListener.File failed: %v", err)
		return
	} else {
		f.Close()
	}
}

func TestTCPConnSpecificMethods(t *testing.T) {
	la, err := net.ResolveTCPAddr("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("net.ResolveTCPAddr failed: %v", err)
	}
	ln, err := net.ListenTCP("tcp4", la)
	if err != nil {
		t.Fatalf("net.ListenTCP failed: %v", err)
	}
	ln.Addr()
	defer ln.Close()

	done := make(chan int)
	go transponder(t, ln, done)

	ra, err := net.ResolveTCPAddr("tcp4", ln.Addr().String())
	if err != nil {
		t.Errorf("net.ResolveTCPAddr failed: %v", err)
		return
	}
	c, err := net.DialTCP("tcp4", nil, ra)
	if err != nil {
		t.Errorf("net.DialTCP failed: %v", err)
		return
	}
	c.SetKeepAlive(false)
	c.SetLinger(0)
	c.SetNoDelay(false)
	c.LocalAddr()
	c.RemoteAddr()
	c.SetDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetWriteDeadline(time.Now().Add(100 * time.Millisecond))
	defer c.Close()

	if _, err := c.Write([]byte("TCPCONN TEST")); err != nil {
		t.Errorf("net.TCPConn.Write failed: %v", err)
		return
	}
	rb := make([]byte, 128)
	if _, err := c.Read(rb); err != nil {
		t.Errorf("net.TCPConn.Read failed: %v", err)
		return
	}

	<-done
}

func TestUDPConnSpecificMethods(t *testing.T) {
	la, err := net.ResolveUDPAddr("udp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("net.ResolveUDPAddr failed: %v", err)
	}
	c, err := net.ListenUDP("udp4", la)
	if err != nil {
		t.Fatalf("net.ListenUDP failed: %v", err)
	}
	c.LocalAddr()
	c.RemoteAddr()
	c.SetDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetWriteDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetReadBuffer(2048)
	c.SetWriteBuffer(2048)
	defer c.Close()

	wb := []byte("UDPCONN TEST")
	rb := make([]byte, 128)
	if _, err := c.WriteToUDP(wb, c.LocalAddr().(*net.UDPAddr)); err != nil {
		t.Errorf("net.UDPConn.WriteToUDP failed: %v", err)
		return
	}
	if _, _, err := c.ReadFromUDP(rb); err != nil {
		t.Errorf("net.UDPConn.ReadFromUDP failed: %v", err)
		return
	}
	if _, _, err := c.WriteMsgUDP(wb, nil, c.LocalAddr().(*net.UDPAddr)); err != nil {
		condErrorf(t, "net.UDPConn.WriteMsgUDP failed: %v", err)
		return
	}
	if _, _, _, _, err := c.ReadMsgUDP(rb, nil); err != nil {
		condErrorf(t, "net.UDPConn.ReadMsgUDP failed: %v", err)
		return
	}

	if f, err := c.File(); err != nil {
		condErrorf(t, "net.UDPConn.File failed: %v", err)
		return
	} else {
		f.Close()
	}
}

func TestIPConnSpecificMethods(t *testing.T) {
	switch runtime.GOOS {
	case "plan9":
		t.Logf("skipping read test on %q", runtime.GOOS)
		return
	}
	if os.Getuid() != 0 {
		t.Logf("skipping test; must be root")
		return
	}

	la, err := net.ResolveIPAddr("ip4", "127.0.0.1")
	if err != nil {
		t.Fatalf("net.ResolveIPAddr failed: %v", err)
	}
	c, err := net.ListenIP("ip4:icmp", la)
	if err != nil {
		t.Fatalf("net.ListenIP failed: %v", err)
	}
	c.LocalAddr()
	c.RemoteAddr()
	c.SetDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetWriteDeadline(time.Now().Add(100 * time.Millisecond))
	c.SetReadBuffer(2048)
	c.SetWriteBuffer(2048)
	defer c.Close()

	id := os.Getpid() & 0xffff
	wb := newICMPEchoRequest(id, 1, 128, []byte("IPCONN TEST "))
	rb := make([]byte, 20+128)
	if _, err := c.WriteToIP(wb, c.LocalAddr().(*net.IPAddr)); err != nil {
		t.Errorf("net.IPConn.WriteToIP failed: %v", err)
		return
	}
	if _, _, err := c.ReadFromIP(rb); err != nil {
		t.Errorf("net.IPConn.ReadFromIP failed: %v", err)
		return
	}
	if _, _, err := c.WriteMsgIP(wb, nil, c.LocalAddr().(*net.IPAddr)); err != nil {
		condErrorf(t, "net.UDPConn.WriteMsgIP failed: %v", err)
		return
	}
	if _, _, _, _, err := c.ReadMsgIP(rb, nil); err != nil {
		condErrorf(t, "net.UDPConn.ReadMsgIP failed: %v", err)
		return
	}

	if f, err := c.File(); err != nil {
		condErrorf(t, "net.IPConn.File failed: %v", err)
		return
	} else {
		f.Close()
	}
}

func TestUnixListenerSpecificMethods(t *testing.T) {
	switch runtime.GOOS {
	case "plan9", "windows":
		t.Logf("skipping read test on %q", runtime.GOOS)
		return
	}

	p := "/tmp/gotest.net"
	os.Remove(p)
	la, err := net.ResolveUnixAddr("unix", p)
	if err != nil {
		t.Fatalf("net.ResolveUnixAddr failed: %v", err)
	}
	ln, err := net.ListenUnix("unix", la)
	if err != nil {
		t.Fatalf("net.ListenUnix failed: %v", err)
	}
	ln.Addr()
	ln.SetDeadline(time.Now().Add(30 * time.Nanosecond))
	defer ln.Close()
	defer os.Remove(p)

	if c, err := ln.Accept(); err != nil {
		if !err.(net.Error).Timeout() {
			t.Errorf("net.TCPListener.AcceptTCP failed: %v", err)
			return
		}
	} else {
		c.Close()
	}
	if c, err := ln.AcceptUnix(); err != nil {
		if !err.(net.Error).Timeout() {
			t.Errorf("net.TCPListener.AcceptTCP failed: %v", err)
			return
		}
	} else {
		c.Close()
	}

	if f, err := ln.File(); err != nil {
		t.Errorf("net.UnixListener.File failed: %v", err)
		return
	} else {
		f.Close()
	}
}

func TestUnixConnSpecificMethods(t *testing.T) {
	switch runtime.GOOS {
	case "plan9", "windows":
		t.Logf("skipping test on %q", runtime.GOOS)
		return
	}

	p1, p2 := "/tmp/gotest.net1", "/tmp/gotest.net2"
	os.Remove(p1)
	os.Remove(p2)

	a1, err := net.ResolveUnixAddr("unixgram", p1)
	if err != nil {
		t.Fatalf("net.ResolveUnixAddr failed: %v", err)
	}
	c1, err := net.DialUnix("unixgram", a1, nil)
	if err != nil {
		t.Fatalf("net.DialUnix failed: %v", err)
	}
	c1.LocalAddr()
	c1.RemoteAddr()
	c1.SetDeadline(time.Now().Add(100 * time.Millisecond))
	c1.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	c1.SetWriteDeadline(time.Now().Add(100 * time.Millisecond))
	c1.SetReadBuffer(2048)
	c1.SetWriteBuffer(2048)
	defer c1.Close()
	defer os.Remove(p1)

	a2, err := net.ResolveUnixAddr("unixgram", p2)
	if err != nil {
		t.Errorf("net.ResolveUnixAddr failed: %v", err)
		return
	}
	c2, err := net.DialUnix("unixgram", a2, nil)
	if err != nil {
		t.Errorf("net.DialUnix failed: %v", err)
		return
	}
	c2.LocalAddr()
	c2.RemoteAddr()
	c2.SetDeadline(time.Now().Add(100 * time.Millisecond))
	c2.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	c2.SetWriteDeadline(time.Now().Add(100 * time.Millisecond))
	c2.SetReadBuffer(2048)
	c2.SetWriteBuffer(2048)
	defer c2.Close()
	defer os.Remove(p2)

	wb := []byte("UNIXCONN TEST")
	rb1 := make([]byte, 128)
	rb2 := make([]byte, 128)
	if _, _, err := c1.WriteMsgUnix(wb, nil, a2); err != nil {
		t.Errorf("net.UnixConn.WriteMsgUnix failed: %v", err)
		return
	}
	if _, _, _, _, err := c2.ReadMsgUnix(rb2, nil); err != nil {
		t.Errorf("net.UnixConn.ReadMsgUnix failed: %v", err)
		return
	}
	if _, err := c2.WriteToUnix(wb, a1); err != nil {
		t.Errorf("net.UnixConn.WriteToUnix failed: %v", err)
		return
	}
	if _, _, err := c1.ReadFromUnix(rb1); err != nil {
		t.Errorf("net.UnixConn.ReadFromUnix failed: %v", err)
		return
	}

	// TODO: http://golang.org/issue/3875
	net.ListenUnixgram("unixgram", nil)

	if f, err := c1.File(); err != nil {
		t.Errorf("net.UnixConn.File failed: %v", err)
		return
	} else {
		f.Close()
	}
}

func newICMPEchoRequest(id, seqnum, msglen int, filler []byte) []byte {
	b := newICMPInfoMessage(id, seqnum, msglen, filler)
	b[0] = 8
	// calculate ICMP checksum
	cklen := len(b)
	s := uint32(0)
	for i := 0; i < cklen-1; i += 2 {
		s += uint32(b[i+1])<<8 | uint32(b[i])
	}
	if cklen&1 == 1 {
		s += uint32(b[cklen-1])
	}
	s = (s >> 16) + (s & 0xffff)
	s = s + (s >> 16)
	// place checksum back in header; using ^= avoids the
	// assumption the checksum bytes are zero
	b[2] ^= byte(^s & 0xff)
	b[3] ^= byte(^s >> 8)
	return b
}

func newICMPInfoMessage(id, seqnum, msglen int, filler []byte) []byte {
	b := make([]byte, msglen)
	copy(b[8:], bytes.Repeat(filler, (msglen-8)/len(filler)+1))
	b[0] = 0                   // type
	b[1] = 0                   // code
	b[2] = 0                   // checksum
	b[3] = 0                   // checksum
	b[4] = byte(id >> 8)       // identifier
	b[5] = byte(id & 0xff)     // identifier
	b[6] = byte(seqnum >> 8)   // sequence number
	b[7] = byte(seqnum & 0xff) // sequence number
	return b
}
<|endoftext|>


<commit_before>package bbs

import "fmt"

//This has all the structs for various commands

type BBSCommand struct {
	Command string `json:"cmd"`
}

type UserCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session"`
}

//From start to end inclusive, starting from 1.
type Range struct {
	Start int `json:"start"`
	End   int `json:"end"`
}

func (r Range) String() string {
	return fmt.Sprintf("%d-%d", r.Start, r.End)
}

func (r Range) Validate() bool {
	if r.Start > r.End {
		return false
	}
	return true
}

func (r Range) Empty() bool {
	return r.End == 0
}

// "hello" message (server -> client)
type HelloMessage struct {
	Command         string     `json:"cmd"`
	Name            string     `json:"name"`
	ProtocolVersion int        `json:"version"`
	Description     string     `json:"desc"`
	SecureURL       string     `json:"secure,omitempty"` //https URL, if any
	Options         []string   `json:"options,omitempty"`
	Access          AccessInfo `json:"access"`
	Formats         []string   `json:"format"` //formats the server accepts, the first one should be the primary one
	Lists           []string   `json:"lists"`
	ServerVersion   string     `json:"server"`
	IconURL         string     `json:"icon"`
	// for option "range"
	DefaultRange Range `json:"default_range,omitempty"`
	// for option "realtime"
	RealtimeURL string `json:"realtime"`
}

// guest commands are commands you can use without logging on (e.g. "list", "get")
// user commands require being logged in first (usually "post" and "reply")
type AccessInfo struct {
	GuestCommands []string `json:"guest,omitempty"`
	UserCommands  []string `json:"user,omitempty"`
}

// "error" message (server -> client)
type ErrorMessage struct {
	Command string `json:"cmd"`
	ReplyTo string `json:"wrt"`
	Error   string `json:"error"`
}

// session expired or invalid? use this
var SessionErrorMessage ErrorMessage = Error("session", "bad session")

// "ok" message (server -> client)
type OKMessage struct {
	Command string `json:"cmd"`
	ReplyTo string `json:"wrt"`
	Result  string `json:"result,omitempty"`
}

// "login" command (client -> server)
type LoginCommand struct {
	Command         string `json:"cmd"`
	Username        string `json:"username"`
	Password        string `json:"password"`
	ProtocolVersion int    `json:"version"`

	// for "re-logins" only
	Session string `json:"session,omitempty"`
}

// "welcome" message (server -> client)
type WelcomeMessage struct {
	Command  string `json:"cmd"`
	Username string `json:"username,omitempty"` //omit for option 'anon'
	Session  string `json:"session"`
}

// "logout" command (client -> server)
type LogoutCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session"`
}

// "register" command (client -> server)
type RegisterCommand struct {
	Command  string `json:"cmd"`
	Username string `json:"username"`
	Password string `json:"password"`
	Email    string `json:"email,omitempty"`
}

// "get" command (client -> server)
type GetCommand struct {
	Command  string `json:"cmd"`
	Session  string `json:"session,omitempty"`
	ThreadID string `json:"id"`
	Range    Range  `json:"range"`            //option: "range"
	Filter   string `json:"filter,omitempty"` //option: "filter"
	Format   string `json:"format,omitempty"`
	Token    string `json:"token,omitempty"`
}

// "list" command (client -> server)
type ListCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session,omitempty"`
	Type    string `json:"type"`
	Query   string `json:"query"` //board for "boards", tag expression for "tags" (like "Dogs+Pizza-Anime")
	Token   string `json:"token,omitempty"`
}

// "reply" command (client -> server)
type ReplyCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session,omitempty"`
	To      string `json:"to"`
	Text    string `json:"body"`
	Format  string `json:"format,omitempty"`
}

// "post" command (client -> server)
type PostCommand struct {
	Command string   `json:"cmd"`
	Session string   `json:"session,omitempty"`
	Title   string   `json:"title"`
	Text    string   `json:"body"`
	Format  string   `json:"format,omitempty"`
	Board   string   `json:"board,omitempty"` //option: "boards"
	Tags    []string `json:"tags,omitempty"`  //option: "tags"
}

// "msg" message (server -> client) [response to "get"]
type ThreadMessage struct {
	Command   string    `json:"cmd"`
	ID        string    `json:"id" bson:"_id"`
	Title     string    `json:"title,omitempty"`
	Range     Range     `json:"range,omitempty"`
	Closed    bool      `json:"closed,omitempty"`
	Filter    string    `json:"filter,omitempty"` //option: "filter"
	Board     string    `json:"board,omitempty"`  //option: "boards"
	Tags      []string  `json:"tags,omitempty"`   //option: "tags"
	Format    string    `json:"format,omitempty"`
	Messages  []Message `json:"messages"`
	More      bool      `json:"more,omitempty"`
	NextToken string    `json:"next,omitempty"`
}

func (t ThreadMessage) Size() int {
	if t.Messages != nil {
		return len(t.Messages)
	}
	return 0
}

// format for posts used in "msg"
type Message struct {
	ID                 string `json:"id"`
	Author             string `json:"user"`
	AuthorID           string `json:"user_id,omitempty"`
	Date               string `json:"date,omitempty"`
	Text               string `json:"body"`
	Signature          string `json:"sig,omitempty"`
	AuthorTitle        string `json:"user_title,omitempty"`   //option: "usertitles"
	AvatarURL          string `json:"avatar,omitempty"`       //option: "avatars"
	AvatarThumbnailURL string `json:"avatar_thumb,omitempty"` //option: "avatars"
	PictureURL         string `json:"img,omitempty"`          //option: "imageboard"
	ThumbnailURL       string `json:"thumb,omitempty"`        //option: "imageboard"
}

type TypedMessage struct {
	Command string `json:"cmd"`
	Type    string `json:"type"`
}

// "list" message where type = "thread" (server -> client)
type ListMessage struct {
	Command   string          `json:"cmd"`
	Type      string          `json:"type"`
	Query     string          `json:"query,omitempty"`
	Threads   []ThreadListing `json:"threads"`
	NextToken string          `json:"next,omitempty"`
}

// "list" message where type = "board" (server -> client)
type BoardListMessage struct {
	Command string         `json:"cmd"`
	Type    string         `json:"type"`
	Query   string         `json:"query,omitempty"`
	Boards  []BoardListing `json:"boards"`
}

type BookmarkListMessage struct {
	Command   string     `json:"cmd"`
	Type      string     `json:"type"`
	Bookmarks []Bookmark `json:"bookmarks"`
}

type Bookmark struct {
	ID    string `json:"id,omitempty"`
	Name  string `json:"name"`
	Query string `json:"query"`
}

// format for threads in "list"
type ThreadListing struct {
	ID           string   `json:"id"`
	Title        string   `json:"title"`
	Author       string   `json:"user,omitempty"`
	AuthorID     string   `json:"user_id,omitempty"`
	Date         string   `json:"date,omitempty"`
	PostCount    int      `json:"posts,omitempty"`
	UnreadPosts  int      `json:"unread_posts,omitempty"`
	Sticky       bool     `json:"sticky,omitempty"` //a sticky (aka pinned) topic
	Closed       bool     `json:"closed,omitempty"` //a closed (aka locked) topic
	Tags         []string `json:"tags,omitempty"`   //option: "tags"
	PictureURL   string   `json:"img,omitempty"`    //option: "imageboard"
	ThumbnailURL string   `json:"thumb,omitempty"`  //option: "imageboard"
}

// format for boards in "list"
type BoardListing struct {
	ID          string `json:"id"`
	Name        string `json:"name,omitempty"`
	Description string `json:"desc,omitempty"`
	ThreadCount int    `json:"threads,omitempty"`
	PostCount   int    `json:"posts,omitempty"`
	Date        string `json:"date,omitempty"`
}
<commit_msg>added total post count to `msg` msg<commit_after>package bbs

import "fmt"

//This has all the structs for various commands

type BBSCommand struct {
	Command string `json:"cmd"`
}

type UserCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session"`
}

//From start to end inclusive, starting from 1.
type Range struct {
	Start int `json:"start"`
	End   int `json:"end"`
}

func (r Range) String() string {
	return fmt.Sprintf("%d-%d", r.Start, r.End)
}

func (r Range) Validate() bool {
	if r.Start > r.End {
		return false
	}
	return true
}

func (r Range) Empty() bool {
	return r.End == 0
}

// "hello" message (server -> client)
type HelloMessage struct {
	Command         string     `json:"cmd"`
	Name            string     `json:"name"`
	ProtocolVersion int        `json:"version"`
	Description     string     `json:"desc"`
	SecureURL       string     `json:"secure,omitempty"` //https URL, if any
	Options         []string   `json:"options,omitempty"`
	Access          AccessInfo `json:"access"`
	Formats         []string   `json:"format"` //formats the server accepts, the first one should be the primary one
	Lists           []string   `json:"lists"`
	ServerVersion   string     `json:"server"`
	IconURL         string     `json:"icon"`
	// for option "range"
	DefaultRange Range `json:"default_range,omitempty"`
	// for option "realtime"
	RealtimeURL string `json:"realtime"`
}

// guest commands are commands you can use without logging on (e.g. "list", "get")
// user commands require being logged in first (usually "post" and "reply")
type AccessInfo struct {
	GuestCommands []string `json:"guest,omitempty"`
	UserCommands  []string `json:"user,omitempty"`
}

// "error" message (server -> client)
type ErrorMessage struct {
	Command string `json:"cmd"`
	ReplyTo string `json:"wrt"`
	Error   string `json:"error"`
}

// session expired or invalid? use this
var SessionErrorMessage ErrorMessage = Error("session", "bad session")

// "ok" message (server -> client)
type OKMessage struct {
	Command string `json:"cmd"`
	ReplyTo string `json:"wrt"`
	Result  string `json:"result,omitempty"`
}

// "login" command (client -> server)
type LoginCommand struct {
	Command         string `json:"cmd"`
	Username        string `json:"username"`
	Password        string `json:"password"`
	ProtocolVersion int    `json:"version"`

	// for "re-logins" only
	Session string `json:"session,omitempty"`
}

// "welcome" message (server -> client)
type WelcomeMessage struct {
	Command  string `json:"cmd"`
	Username string `json:"username,omitempty"` //omit for option 'anon'
	Session  string `json:"session"`
}

// "logout" command (client -> server)
type LogoutCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session"`
}

// "register" command (client -> server)
type RegisterCommand struct {
	Command  string `json:"cmd"`
	Username string `json:"username"`
	Password string `json:"password"`
	Email    string `json:"email,omitempty"`
}

// "get" command (client -> server)
type GetCommand struct {
	Command  string `json:"cmd"`
	Session  string `json:"session,omitempty"`
	ThreadID string `json:"id"`
	Range    Range  `json:"range"`            //option: "range"
	Filter   string `json:"filter,omitempty"` //option: "filter"
	Format   string `json:"format,omitempty"`
	Token    string `json:"token,omitempty"`
}

// "list" command (client -> server)
type ListCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session,omitempty"`
	Type    string `json:"type"`
	Query   string `json:"query"` //board for "boards", tag expression for "tags" (like "Dogs+Pizza-Anime")
	Token   string `json:"token,omitempty"`
}

// "reply" command (client -> server)
type ReplyCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session,omitempty"`
	To      string `json:"to"`
	Text    string `json:"body"`
	Format  string `json:"format,omitempty"`
}

// "post" command (client -> server)
type PostCommand struct {
	Command string   `json:"cmd"`
	Session string   `json:"session,omitempty"`
	Title   string   `json:"title"`
	Text    string   `json:"body"`
	Format  string   `json:"format,omitempty"`
	Board   string   `json:"board,omitempty"` //option: "boards"
	Tags    []string `json:"tags,omitempty"`  //option: "tags"
}

// "msg" message (server -> client) [response to "get"]
type ThreadMessage struct {
	Command   string    `json:"cmd"`
	ID        string    `json:"id" bson:"_id"`
	Title     string    `json:"title,omitempty"`
	Range     Range     `json:"range,omitempty"`
	Closed    bool      `json:"closed,omitempty"`
	Filter    string    `json:"filter,omitempty"` //option: "filter"
	Board     string    `json:"board,omitempty"`  //option: "boards"
	Tags      []string  `json:"tags,omitempty"`   //option: "tags"
	Format    string    `json:"format,omitempty"`
	Messages  []Message `json:"messages"`
	Total     int       `json:"total,omitempty"`
	More      bool      `json:"more,omitempty"`
	NextToken string    `json:"next,omitempty"`
}

func (t ThreadMessage) Size() int {
	if t.Messages != nil {
		return len(t.Messages)
	}
	return 0
}

// format for posts used in "msg"
type Message struct {
	ID                 string `json:"id"`
	Author             string `json:"user"`
	AuthorID           string `json:"user_id,omitempty"`
	Date               string `json:"date,omitempty"`
	Text               string `json:"body"`
	Signature          string `json:"sig,omitempty"`
	AuthorTitle        string `json:"user_title,omitempty"`   //option: "usertitles"
	AvatarURL          string `json:"avatar,omitempty"`       //option: "avatars"
	AvatarThumbnailURL string `json:"avatar_thumb,omitempty"` //option: "avatars"
	PictureURL         string `json:"img,omitempty"`          //option: "imageboard"
	ThumbnailURL       string `json:"thumb,omitempty"`        //option: "imageboard"
}

type TypedMessage struct {
	Command string `json:"cmd"`
	Type    string `json:"type"`
}

// "list" message where type = "thread" (server -> client)
type ListMessage struct {
	Command   string          `json:"cmd"`
	Type      string          `json:"type"`
	Query     string          `json:"query,omitempty"`
	Threads   []ThreadListing `json:"threads"`
	NextToken string          `json:"next,omitempty"`
}

// "list" message where type = "board" (server -> client)
type BoardListMessage struct {
	Command string         `json:"cmd"`
	Type    string         `json:"type"`
	Query   string         `json:"query,omitempty"`
	Boards  []BoardListing `json:"boards"`
}

type BookmarkListMessage struct {
	Command   string     `json:"cmd"`
	Type      string     `json:"type"`
	Bookmarks []Bookmark `json:"bookmarks"`
}

type Bookmark struct {
	ID    string `json:"id,omitempty"`
	Name  string `json:"name"`
	Query string `json:"query"`
}

// format for threads in "list"
type ThreadListing struct {
	ID           string   `json:"id"`
	Title        string   `json:"title"`
	Author       string   `json:"user,omitempty"`
	AuthorID     string   `json:"user_id,omitempty"`
	Date         string   `json:"date,omitempty"`
	PostCount    int      `json:"posts,omitempty"`
	UnreadPosts  int      `json:"unread_posts,omitempty"`
	Sticky       bool     `json:"sticky,omitempty"` //a sticky (aka pinned) topic
	Closed       bool     `json:"closed,omitempty"` //a closed (aka locked) topic
	Tags         []string `json:"tags,omitempty"`   //option: "tags"
	PictureURL   string   `json:"img,omitempty"`    //option: "imageboard"
	ThumbnailURL string   `json:"thumb,omitempty"`  //option: "imageboard"
}

// format for boards in "list"
type BoardListing struct {
	ID          string `json:"id"`
	Name        string `json:"name,omitempty"`
	Description string `json:"desc,omitempty"`
	ThreadCount int    `json:"threads,omitempty"`
	PostCount   int    `json:"posts,omitempty"`
	Date        string `json:"date,omitempty"`
}
<|endoftext|>


<commit_before>// Copyright 2012 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.

package runtime_test

import (
	. "runtime"
	"sync"
	"testing"
	"time"
	"unsafe"
)

// See stack.h.
const (
	StackGuard = 256
	StackLimit = 128
)

// Test stack split logic by calling functions of every frame size
// from near 0 up to and beyond the default segment size (4k).
// Each of those functions reports its SP + stack limit, and then
// the test (the caller) checks that those make sense.  By not
// doing the actual checking and reporting from the suspect functions,
// we minimize the possibility of crashes during the test itself.
//
// Exhaustive test for http://golang.org/issue/3310.
// The linker used to get a few sizes near the segment size wrong:
//
// --- FAIL: TestStackSplit (0.01 seconds)
// 	stack_test.go:22: after runtime_test.stack3812: sp=0x7f7818d5d078 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3816: sp=0x7f7818d5d078 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3820: sp=0x7f7818d5d070 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3824: sp=0x7f7818d5d070 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3828: sp=0x7f7818d5d068 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3832: sp=0x7f7818d5d068 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3836: sp=0x7f7818d5d060 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3840: sp=0x7f7818d5d060 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3844: sp=0x7f7818d5d058 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3848: sp=0x7f7818d5d058 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3852: sp=0x7f7818d5d050 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3856: sp=0x7f7818d5d050 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3860: sp=0x7f7818d5d048 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3864: sp=0x7f7818d5d048 < limit=0x7f7818d5d080
// FAIL
func TestStackSplit(t *testing.T) {
	for _, f := range splitTests {
		sp, guard := f()
		bottom := guard - StackGuard
		if sp < bottom+StackLimit {
			fun := FuncForPC(**(**uintptr)(unsafe.Pointer(&f)))
			t.Errorf("after %s: sp=%#x < limit=%#x (guard=%#x, bottom=%#x)",
				fun.Name(), sp, bottom+StackLimit, guard, bottom)
		}
	}
}

var Used byte

func use(buf []byte) {
	for _, c := range buf {
		Used += c
	}
}

// TestStackMem measures per-thread stack segment cache behavior.
// The test consumed up to 500MB in the past.
func TestStackMem(t *testing.T) {
	const (
		BatchSize      = 32
		BatchCount     = 256
		ArraySize      = 1024
		RecursionDepth = 128
	)
	if testing.Short() {
		return
	}
	defer GOMAXPROCS(GOMAXPROCS(BatchSize))
	s0 := new(MemStats)
	ReadMemStats(s0)
	for b := 0; b < BatchCount; b++ {
		c := make(chan bool, BatchSize)
		for i := 0; i < BatchSize; i++ {
			go func() {
				var f func(k int, a [ArraySize]byte)
				f = func(k int, a [ArraySize]byte) {
					if k == 0 {
						time.Sleep(time.Millisecond)
						return
					}
					f(k-1, a)
				}
				f(RecursionDepth, [ArraySize]byte{})
				c <- true
			}()
		}
		for i := 0; i < BatchSize; i++ {
			<-c
		}

		// The goroutines have signaled via c that they are ready to exit.
		// Give them a chance to exit by sleeping. If we don't wait, we
		// might not reuse them on the next batch.
		time.Sleep(10 * time.Millisecond)
	}
	s1 := new(MemStats)
	ReadMemStats(s1)
	consumed := s1.StackSys - s0.StackSys
	t.Logf("Consumed %vMB for stack mem", consumed>>20)
	estimate := uint64(8 * BatchSize * ArraySize * RecursionDepth) // 8 is to reduce flakiness.
	if consumed > estimate {
		t.Fatalf("Stack mem: want %v, got %v", estimate, consumed)
	}
	// Due to broken stack memory accounting (http://golang.org/issue/7468),
	// StackInuse can decrease during function execution, so we cast the values to int64.
	inuse := int64(s1.StackInuse) - int64(s0.StackInuse)
	t.Logf("Inuse %vMB for stack mem", inuse>>20)
	if inuse > 4<<20 {
		t.Fatalf("Stack inuse: want %v, got %v", 4<<20, inuse)
	}
}

// Test stack growing in different contexts.
func TestStackGrowth(t *testing.T) {
	t.Parallel()
	var wg sync.WaitGroup

	// in a normal goroutine
	wg.Add(1)
	go func() {
		defer wg.Done()
		growStack()
	}()
	wg.Wait()

	// in locked goroutine
	wg.Add(1)
	go func() {
		defer wg.Done()
		LockOSThread()
		growStack()
		UnlockOSThread()
	}()
	wg.Wait()

	// in finalizer
	wg.Add(1)
	go func() {
		defer wg.Done()
		done := make(chan bool)
		go func() {
			s := new(string)
			SetFinalizer(s, func(ss *string) {
				growStack()
				done <- true
			})
			s = nil
			done <- true
		}()
		<-done
		GC()
		select {
		case <-done:
		case <-time.After(4 * time.Second):
			t.Fatal("finalizer did not run")
		}
	}()
	wg.Wait()
}

// ... and in init
//func init() {
//	growStack()
//}

func growStack() {
	n := 1 << 10
	if testing.Short() {
		n = 1 << 8
	}
	for i := 0; i < n; i++ {
		x := 0
		growStackIter(&x, i)
		if x != i+1 {
			panic("stack is corrupted")
		}
	}
	GC()
}

// This function is not an anonimous func, so that the compiler can do escape
// analysis and place x on stack (and subsequently stack growth update the pointer).
func growStackIter(p *int, n int) {
	if n == 0 {
		*p = n + 1
		GC()
		return
	}
	*p = n + 1
	x := 0
	growStackIter(&x, n-1)
	if x != n {
		panic("stack is corrupted")
	}
}

func TestStackGrowthCallback(t *testing.T) {
	t.Parallel()
	var wg sync.WaitGroup

	// test stack growth at chan op
	wg.Add(1)
	go func() {
		defer wg.Done()
		c := make(chan int, 1)
		growStackWithCallback(func() {
			c <- 1
			<-c
		})
	}()

	// test stack growth at map op
	wg.Add(1)
	go func() {
		defer wg.Done()
		m := make(map[int]int)
		growStackWithCallback(func() {
			_, _ = m[1]
			m[1] = 1
		})
	}()

	// test stack growth at goroutine creation
	wg.Add(1)
	go func() {
		defer wg.Done()
		growStackWithCallback(func() {
			done := make(chan bool)
			go func() {
				done <- true
			}()
			<-done
		})
	}()

	wg.Wait()
}

func growStackWithCallback(cb func()) {
	var f func(n int)
	f = func(n int) {
		if n == 0 {
			cb()
			return
		}
		f(n - 1)
	}
	for i := 0; i < 1<<10; i++ {
		f(i)
	}
}

// TestDeferPtrs tests the adjustment of Defer's argument pointers (p aka &y)
// during a stack copy.
func set(p *int, x int) {
	*p = x
}
func TestDeferPtrs(t *testing.T) {
	var y int

	defer func() {
		if y != 42 {
			t.Errorf("defer's stack references were not adjusted appropriately")
		}
	}()
	defer set(&y, 42)
	growStack()
}
<commit_msg>runtime: increase timeout in TestStackGrowth<commit_after>// Copyright 2012 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.

package runtime_test

import (
	. "runtime"
	"sync"
	"testing"
	"time"
	"unsafe"
)

// See stack.h.
const (
	StackGuard = 256
	StackLimit = 128
)

// Test stack split logic by calling functions of every frame size
// from near 0 up to and beyond the default segment size (4k).
// Each of those functions reports its SP + stack limit, and then
// the test (the caller) checks that those make sense.  By not
// doing the actual checking and reporting from the suspect functions,
// we minimize the possibility of crashes during the test itself.
//
// Exhaustive test for http://golang.org/issue/3310.
// The linker used to get a few sizes near the segment size wrong:
//
// --- FAIL: TestStackSplit (0.01 seconds)
// 	stack_test.go:22: after runtime_test.stack3812: sp=0x7f7818d5d078 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3816: sp=0x7f7818d5d078 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3820: sp=0x7f7818d5d070 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3824: sp=0x7f7818d5d070 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3828: sp=0x7f7818d5d068 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3832: sp=0x7f7818d5d068 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3836: sp=0x7f7818d5d060 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3840: sp=0x7f7818d5d060 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3844: sp=0x7f7818d5d058 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3848: sp=0x7f7818d5d058 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3852: sp=0x7f7818d5d050 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3856: sp=0x7f7818d5d050 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3860: sp=0x7f7818d5d048 < limit=0x7f7818d5d080
// 	stack_test.go:22: after runtime_test.stack3864: sp=0x7f7818d5d048 < limit=0x7f7818d5d080
// FAIL
func TestStackSplit(t *testing.T) {
	for _, f := range splitTests {
		sp, guard := f()
		bottom := guard - StackGuard
		if sp < bottom+StackLimit {
			fun := FuncForPC(**(**uintptr)(unsafe.Pointer(&f)))
			t.Errorf("after %s: sp=%#x < limit=%#x (guard=%#x, bottom=%#x)",
				fun.Name(), sp, bottom+StackLimit, guard, bottom)
		}
	}
}

var Used byte

func use(buf []byte) {
	for _, c := range buf {
		Used += c
	}
}

// TestStackMem measures per-thread stack segment cache behavior.
// The test consumed up to 500MB in the past.
func TestStackMem(t *testing.T) {
	const (
		BatchSize      = 32
		BatchCount     = 256
		ArraySize      = 1024
		RecursionDepth = 128
	)
	if testing.Short() {
		return
	}
	defer GOMAXPROCS(GOMAXPROCS(BatchSize))
	s0 := new(MemStats)
	ReadMemStats(s0)
	for b := 0; b < BatchCount; b++ {
		c := make(chan bool, BatchSize)
		for i := 0; i < BatchSize; i++ {
			go func() {
				var f func(k int, a [ArraySize]byte)
				f = func(k int, a [ArraySize]byte) {
					if k == 0 {
						time.Sleep(time.Millisecond)
						return
					}
					f(k-1, a)
				}
				f(RecursionDepth, [ArraySize]byte{})
				c <- true
			}()
		}
		for i := 0; i < BatchSize; i++ {
			<-c
		}

		// The goroutines have signaled via c that they are ready to exit.
		// Give them a chance to exit by sleeping. If we don't wait, we
		// might not reuse them on the next batch.
		time.Sleep(10 * time.Millisecond)
	}
	s1 := new(MemStats)
	ReadMemStats(s1)
	consumed := s1.StackSys - s0.StackSys
	t.Logf("Consumed %vMB for stack mem", consumed>>20)
	estimate := uint64(8 * BatchSize * ArraySize * RecursionDepth) // 8 is to reduce flakiness.
	if consumed > estimate {
		t.Fatalf("Stack mem: want %v, got %v", estimate, consumed)
	}
	// Due to broken stack memory accounting (http://golang.org/issue/7468),
	// StackInuse can decrease during function execution, so we cast the values to int64.
	inuse := int64(s1.StackInuse) - int64(s0.StackInuse)
	t.Logf("Inuse %vMB for stack mem", inuse>>20)
	if inuse > 4<<20 {
		t.Fatalf("Stack inuse: want %v, got %v", 4<<20, inuse)
	}
}

// Test stack growing in different contexts.
func TestStackGrowth(t *testing.T) {
	t.Parallel()
	var wg sync.WaitGroup

	// in a normal goroutine
	wg.Add(1)
	go func() {
		defer wg.Done()
		growStack()
	}()
	wg.Wait()

	// in locked goroutine
	wg.Add(1)
	go func() {
		defer wg.Done()
		LockOSThread()
		growStack()
		UnlockOSThread()
	}()
	wg.Wait()

	// in finalizer
	wg.Add(1)
	go func() {
		defer wg.Done()
		done := make(chan bool)
		go func() {
			s := new(string)
			SetFinalizer(s, func(ss *string) {
				growStack()
				done <- true
			})
			s = nil
			done <- true
		}()
		<-done
		GC()
		select {
		case <-done:
		case <-time.After(20 * time.Second):
			t.Fatal("finalizer did not run")
		}
	}()
	wg.Wait()
}

// ... and in init
//func init() {
//	growStack()
//}

func growStack() {
	n := 1 << 10
	if testing.Short() {
		n = 1 << 8
	}
	for i := 0; i < n; i++ {
		x := 0
		growStackIter(&x, i)
		if x != i+1 {
			panic("stack is corrupted")
		}
	}
	GC()
}

// This function is not an anonimous func, so that the compiler can do escape
// analysis and place x on stack (and subsequently stack growth update the pointer).
func growStackIter(p *int, n int) {
	if n == 0 {
		*p = n + 1
		GC()
		return
	}
	*p = n + 1
	x := 0
	growStackIter(&x, n-1)
	if x != n {
		panic("stack is corrupted")
	}
}

func TestStackGrowthCallback(t *testing.T) {
	t.Parallel()
	var wg sync.WaitGroup

	// test stack growth at chan op
	wg.Add(1)
	go func() {
		defer wg.Done()
		c := make(chan int, 1)
		growStackWithCallback(func() {
			c <- 1
			<-c
		})
	}()

	// test stack growth at map op
	wg.Add(1)
	go func() {
		defer wg.Done()
		m := make(map[int]int)
		growStackWithCallback(func() {
			_, _ = m[1]
			m[1] = 1
		})
	}()

	// test stack growth at goroutine creation
	wg.Add(1)
	go func() {
		defer wg.Done()
		growStackWithCallback(func() {
			done := make(chan bool)
			go func() {
				done <- true
			}()
			<-done
		})
	}()

	wg.Wait()
}

func growStackWithCallback(cb func()) {
	var f func(n int)
	f = func(n int) {
		if n == 0 {
			cb()
			return
		}
		f(n - 1)
	}
	for i := 0; i < 1<<10; i++ {
		f(i)
	}
}

// TestDeferPtrs tests the adjustment of Defer's argument pointers (p aka &y)
// during a stack copy.
func set(p *int, x int) {
	*p = x
}
func TestDeferPtrs(t *testing.T) {
	var y int

	defer func() {
		if y != 42 {
			t.Errorf("defer's stack references were not adjusted appropriately")
		}
	}()
	defer set(&y, 42)
	growStack()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
)

// Lots of the code here are learnt from the http package

type Proxy struct {
	addr string // listen address
}

type clientConn struct {
	keepAlive bool
	buf       *bufio.ReadWriter
	netconn   net.Conn // connection to the proxy client
	// TODO is it possible that one proxy connection is used to server all the client request?
	// Make things simple at this moment and disable http request keep-alive
	// srvconn net.Conn // connection to the server
}

type ProxyError struct {
	msg string
}

func (pe *ProxyError) Error() string { return pe.msg }

func newProxyError(msg string, err error) *ProxyError {
	return &ProxyError{fmt.Sprintln(msg, err)}
}

func NewProxy(addr string) (proxy *Proxy) {
	proxy = &Proxy{addr: addr}
	return
}

func (py *Proxy) Serve() {
	ln, err := net.Listen("tcp", py.addr)
	if err != nil {
		log.Println("Server create failed:", err)
		os.Exit(1)
	}
	info.Println("COW proxy listening", py.addr)

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Println("Client connection:", err)
			continue
		}
		info.Println("New Client:", conn.RemoteAddr())

		c := newClientConn(conn)
		go c.serve()
	}
}

func newClientConn(rwc net.Conn) (c *clientConn) {
	c = &clientConn{netconn: rwc}
	// http pkg uses io.LimitReader with no limit to create a reader, why?
	br := bufio.NewReader(rwc)
	bw := bufio.NewWriter(rwc)
	c.buf = bufio.NewReadWriter(br, bw)
	return
}

func (c *clientConn) serve() {
	defer c.close()
	var r *Request
	var err error
	for {
		if r, err = parseRequest(c.buf.Reader); err != nil {
			// io.EOF means the client connection is closed
			if err != io.EOF {
				log.Println("Reading client request", err)
			}
			return
		}
		if debug {
			debug.Printf("%v", r)
		} else {
			info.Println(r)
		}

		// TODO Need to do the request in a goroutine to support pipelining?
		// If so, how to maintain the order of finishing request?
		// Consider pipelining later as this is just performance improvement.
		if err = c.doRequest(r); err != nil {
			// Server connection error should not close client connection
			// because this connection is between the proxy and the web
			// server, e.g. EOF. If this kind of error lead to client
			// connection close, the proxy will detect this later.
			log.Println("Doing http request", err)
		}

		// How to detect closed client connection?
		// Reading client connection will encounter EOF and detect that the
		// connection has been closed.

		// Firefox will create 6 persistent connections to the proxy server.
		// If opening many connections is not a problem, then nothing need
		// to be done.
		// Otherwise, set a read time out and close connection upon timeout.
		// This should not cause problem as
		// 1. I didn't see any independent message sent by firefox in order to
		//    close a persistent connection
		// 2. Sending Connection: Keep-Alive but actually closing the
		//    connection cause no problem for firefox. (The client should be
		//    able to detect closed connection and open a new one.)
		if !r.KeepAlive {
			break
		}
	}
}

func (c *clientConn) doRequest(r *Request) (err error) {
	host := r.URL.Host
	if !hostHasPort(host) {
		host += ":80"
	}
	srvconn, err := net.Dial("tcp", host)
	if err != nil {
		return newProxyError("Connecting to %s:", err)
	}
	debug.Printf("Connected to %s\n", r.URL.Host)
	// TODO revisit here when implementing keep-alive
	defer srvconn.Close()

	// Send request to the server
	if _, err := srvconn.Write(r.raw.Bytes()); err != nil {
		return err
	}
	// Send request body
	if r.Method == "POST" {
		// TODO we may need to parse content-length
		if _, err = io.Copy(srvconn, c.buf.Reader); err != nil {
			return newProxyError("Sending request body to server", err)
		}
	}

	// Read server reply
	// parse status line
	srvReader := bufio.NewReader(srvconn)
	rp, err := parseResponse(srvReader, r.Method)
	if err != nil {
		return err
	}
	c.buf.WriteString(rp.raw.String())
	// Flush response header to the client earlier
	if err = c.buf.Flush(); err != nil {
		return newProxyError("Flushing response header to client", err)
	}

	// Wrap inside if to avoid function argument evaluation. Would this work?
	if debug {
		debug.Printf("[Response] %s %v\n%v", r.Method, r.URL, rp)
	}

	if err = c.sendResponseBody(srvReader, rp); err != nil {
		return err
	}
	debug.Printf("Finished request %s %s\n", r.Method, r.URL)
	return nil
}

// Send response body if header specifies content length
func (c *clientConn) sendResponseBodyWithContLen(srvReader *bufio.Reader, contLen int64) (err error) {
	debug.Printf("Sending response body to client, content length %d\n", contLen)
	// CopyN will copy n bytes unless there's error of EOF. For EOF, it means
	// the connection is closed, return will propagate till serv function and
	// close client connection.
	if _, err = io.CopyN(c.buf.Writer, srvReader, contLen); err != nil {
		return newProxyError("Sending response body to client", err)
	}
	return nil
}

// Send response body if header specifies chunked encoding
func (c *clientConn) sendResponseBodyChunked(srvReader *bufio.Reader) (err error) {
	debug.Printf("Sending chunked response to client\n")

	b := make([]byte, 2) // buffer for chunked transfer
	for {
		var s string
		// Read chunk size line, ignore chunk extension if any
		if s, err = ReadLine(srvReader); err != nil {
			return newProxyError("Reading chunk size", err)
		}
		// debug.Printf("chunk size line %s", s)
		f := strings.SplitN(s, ";", 2)
		var size int64
		if size, err = strconv.ParseInt(f[0], 16, 64); err != nil {
			return newProxyError("Chunk size not valid", err)
		}
		c.buf.WriteString(s)
		c.buf.WriteString("\r\n")

		if size == 0 { // end of chunked data, ignore any trailers
			goto END
		}

		// Read chunk data and send to client
		if _, err = io.CopyN(c.buf.Writer, srvReader, size); err != nil {
			return newProxyError("Reading chunked data from server", err)
		}
	END:
		// XXX maybe this kind of error handling should be passed to the
		// client? But if the proxy doesn't know when to stop reading from the
		// server, the only way to avoid blocked reading is to set read time
		// out on server connection. Would that be easier?
		if _, err = io.ReadFull(srvReader, b); err != nil {
			return newProxyError("Reading chunked data ending CRLF from server", err)
		}
		if b[0] != '\r' || b[1] != '\n' {
			return newProxyError("Chunked data not ending with CRLF", err)
		}
		c.buf.WriteString("\r\n")
	}
	return nil
}

// Send response body to client.
func (c *clientConn) sendResponseBody(srvReader *bufio.Reader, rp *Response) (err error) {
	if !rp.HasBody {
		return
	}

	if rp.Chunking {
		err = c.sendResponseBodyChunked(srvReader)
	} else if rp.ContLen != 0 {
		err = c.sendResponseBodyWithContLen(srvReader, rp.ContLen)
	} else {
		// Maybe because this is an HTTP/1.0 server. Just read and wait connection close
		debug.Printf("Can't determine response length and not chunked encoding\n")
		if _, err = io.Copy(c.buf.Writer, srvReader); err != nil {
			return err
		}
	}

	if err = c.buf.Flush(); err != nil {
		return newProxyError("Flushing response body to client", err)
	}
	return
}

func (c *clientConn) close() {
	if c.buf != nil {
		c.buf.Flush()
		c.buf = nil
	}
	if c.netconn != nil {
		info.Printf("Client %v connection closed\n", c.netconn.RemoteAddr())
		c.netconn.Close()
		c.netconn = nil
	}
}
<commit_msg>Refactor: make send body function independent.<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
)

// Lots of the code here are learnt from the http package

type Proxy struct {
	addr string // listen address
}

type clientConn struct {
	keepAlive bool
	buf       *bufio.ReadWriter
	netconn   net.Conn // connection to the proxy client
	// TODO is it possible that one proxy connection is used to server all the client request?
	// Make things simple at this moment and disable http request keep-alive
	// srvconn net.Conn // connection to the server
}

type ProxyError struct {
	msg string
}

func (pe *ProxyError) Error() string { return pe.msg }

func newProxyError(msg string, err error) *ProxyError {
	return &ProxyError{fmt.Sprintln(msg, err)}
}

func NewProxy(addr string) (proxy *Proxy) {
	proxy = &Proxy{addr: addr}
	return
}

func (py *Proxy) Serve() {
	ln, err := net.Listen("tcp", py.addr)
	if err != nil {
		log.Println("Server create failed:", err)
		os.Exit(1)
	}
	info.Println("COW proxy listening", py.addr)

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Println("Client connection:", err)
			continue
		}
		info.Println("New Client:", conn.RemoteAddr())

		c := newClientConn(conn)
		go c.serve()
	}
}

func newClientConn(rwc net.Conn) (c *clientConn) {
	c = &clientConn{netconn: rwc}
	// http pkg uses io.LimitReader with no limit to create a reader, why?
	br := bufio.NewReader(rwc)
	bw := bufio.NewWriter(rwc)
	c.buf = bufio.NewReadWriter(br, bw)
	return
}

func (c *clientConn) serve() {
	defer c.close()
	var r *Request
	var err error
	for {
		if r, err = parseRequest(c.buf.Reader); err != nil {
			// io.EOF means the client connection is closed
			if err != io.EOF {
				log.Println("Reading client request", err)
			}
			return
		}
		if debug {
			debug.Printf("%v", r)
		} else {
			info.Println(r)
		}

		// TODO Need to do the request in a goroutine to support pipelining?
		// If so, how to maintain the order of finishing request?
		// Consider pipelining later as this is just performance improvement.
		if err = c.doRequest(r); err != nil {
			// Server connection error should not close client connection
			// because this connection is between the proxy and the web
			// server, e.g. EOF. If this kind of error lead to client
			// connection close, the proxy will detect this later.
			log.Println("Doing http request", err)
		}

		// How to detect closed client connection?
		// Reading client connection will encounter EOF and detect that the
		// connection has been closed.

		// Firefox will create 6 persistent connections to the proxy server.
		// If opening many connections is not a problem, then nothing need
		// to be done.
		// Otherwise, set a read time out and close connection upon timeout.
		// This should not cause problem as
		// 1. I didn't see any independent message sent by firefox in order to
		//    close a persistent connection
		// 2. Sending Connection: Keep-Alive but actually closing the
		//    connection cause no problem for firefox. (The client should be
		//    able to detect closed connection and open a new one.)
		if !r.KeepAlive {
			break
		}
	}
}

func (c *clientConn) doRequest(r *Request) (err error) {
	host := r.URL.Host
	if !hostHasPort(host) {
		host += ":80"
	}
	srvconn, err := net.Dial("tcp", host)
	if err != nil {
		return newProxyError("Connecting to %s:", err)
	}
	// TODO revisit here when implementing keep-alive
	defer srvconn.Close()
	debug.Printf("Connected to %s\n", r.URL.Host)

	// Send request to the server
	if _, err := srvconn.Write(r.raw.Bytes()); err != nil {
		return err
	}
	// Send request body
	if r.Method == "POST" {
		// TODO we may need to parse content-length
		if _, err = io.Copy(srvconn, c.buf.Reader); err != nil {
			return newProxyError("Sending request body to server", err)
		}
	}

	// Read server reply
	// parse status line
	srvReader := bufio.NewReader(srvconn)
	rp, err := parseResponse(srvReader, r.Method)
	if err != nil {
		return err
	}
	c.buf.WriteString(rp.raw.String())
	// Flush response header to the client earlier
	if err = c.buf.Flush(); err != nil {
		return newProxyError("Flushing response header to client", err)
	}

	// Wrap inside if to avoid function argument evaluation. Would this work?
	if debug {
		debug.Printf("[Response] %s %v\n%v", r.Method, r.URL, rp)
	}

	if err = c.sendResponseBody(srvReader, rp); err != nil {
		return err
	}
	debug.Printf("Finished request %s %s\n", r.Method, r.URL)
	return nil
}

// Send response body if header specifies content length
func sendBodyWithContLen(writer *bufio.Writer, reader *bufio.Reader, contLen int64) (err error) {
	debug.Printf("Sending body with content length %d\n", contLen)
	// CopyN will copy n bytes unless there's error of EOF. For EOF, it means
	// the connection is closed, return will propagate till serv function and
	// close client connection.
	if _, err = io.CopyN(writer, reader, contLen); err != nil {
		return newProxyError("Sending response body to client", err)
	}
	return nil
}

// Send response body if header specifies chunked encoding
func sendBodyChunked(writer *bufio.Writer, reader *bufio.Reader) (err error) {
	debug.Printf("Sending chunked body\n")

	b := make([]byte, 2) // buffer for chunked transfer
	for {
		var s string
		// Read chunk size line, ignore chunk extension if any
		if s, err = ReadLine(reader); err != nil {
			return newProxyError("Reading chunk size", err)
		}
		// debug.Printf("chunk size line %s", s)
		f := strings.SplitN(s, ";", 2)
		var size int64
		if size, err = strconv.ParseInt(f[0], 16, 64); err != nil {
			return newProxyError("Chunk size not valid", err)
		}
		writer.WriteString(s)
		writer.WriteString("\r\n")

		if size == 0 { // end of chunked data, ignore any trailers
			goto END
		}

		// Read chunk data and send to client
		if _, err = io.CopyN(writer, reader, size); err != nil {
			return newProxyError("Reading chunked data from server", err)
		}
	END:
		// XXX maybe this kind of error handling should be passed to the
		// client? But if the proxy doesn't know when to stop reading from the
		// server, the only way to avoid blocked reading is to set read time
		// out on server connection. Would that be easier?
		if _, err = io.ReadFull(reader, b); err != nil {
			return newProxyError("Reading chunked data CRLF", err)
		}
		if b[0] != '\r' || b[1] != '\n' {
			return newProxyError("Reading chunked data, not ending with CRLF", err)
		}
		writer.WriteString("\r\n")
	}
	return nil
}

// Send response body to client.
func (c *clientConn) sendResponseBody(srvReader *bufio.Reader, rp *Response) (err error) {
	if !rp.HasBody {
		return
	}

	if rp.Chunking {
		err = sendBodyChunked(c.buf.Writer, srvReader)
	} else if rp.ContLen != 0 {
		err = sendBodyWithContLen(c.buf.Writer, srvReader, rp.ContLen)
	} else {
		// Maybe because this is an HTTP/1.0 server. Just read and wait connection close
		debug.Printf("Can't determine response length and not chunked encoding\n")
		if _, err = io.Copy(c.buf.Writer, srvReader); err != nil {
			return err
		}
	}

	if err = c.buf.Flush(); err != nil {
		return newProxyError("Flushing response body to client", err)
	}
	return
}

func (c *clientConn) close() {
	if c.buf != nil {
		c.buf.Flush()
		c.buf = nil
	}
	if c.netconn != nil {
		info.Printf("Client %v connection closed\n", c.netconn.RemoteAddr())
		c.netconn.Close()
		c.netconn = nil
	}
}
<|endoftext|>


<commit_before>package main

//presharding
//compute hash ,find target redis instance,construct new key with hashtag
import (
	"errors"
	"fmt"
	"hash/crc32"
	"net"
)

var (
	crcTable *crc32.Table
	target   RespReaderWriter
	shards   Shards
)

func Hash(key string) uint32 {
	crc := crc32.New(crcTable)
	crc.Write([]byte(key))
	v := crc.Sum32()
	return v
}

type Shard struct {
	Host          string
	Slot0         uint32
	Slot1         uint32
	ShardRespPool *ObjectPool
}

type Shards struct {
	Slots        uint32
	ShardServers []Shard
}

func NewShards() Shards {
	sh := Shards{}
	s1 := Shard{}
	s1.Host = "172.31.6.58:6479"
	s1.Slot0 = 0
	s1.Slot1 = 512
	s2 := Shard{}
	s2.Host = "172.31.6.58:6579"
	s2.Slot0 = 512
	s2.Slot1 = 1024
	s1.ShardRespPool = NewProxyClientPool(s1.Host)
	s2.ShardRespPool = NewProxyClientPool(s2.Host)
	sh.ShardServers = []Shard{s1, s2}
	return sh

}
func (shards Shards) GetPooledObject(hash uint32) (obj *PooledObject, shard Shard, err error) {
	for _, shard := range shards.ShardServers {
		if hash >= shard.Slot0 && hash < shard.Slot1 {
			fmt.Printf("get shard:%v\n", shard)
			obj, er := shard.ShardRespPool.Borrow()
			return obj, shard, er
		}
	}
	return nil, Shard{}, nil
}

func (shards Shards) GetShard(hash uint32) (Shard, error) {
	for i, shard := range shards.ShardServers {
		if hash >= shard.Slot0 && hash < shard.Slot1 {
			return shards.ShardServers[i], nil
		}
	}
	return Shard{}, errors.New("no shard found")
}

func HandleConn(conn net.Conn) {
	//	timeout := time.Now()
	//	timeout.Add(10 *time.Minute )
	//	conn.SetReadDeadline(timeout)
	//conn.
	client := NewRespReadWriter(conn)
	cch := make(chan int) //close flag chan
	dch := client.LoopRead(cch)
	shardMap := make(map[Shard]*PooledObject)
	var lastHash uint32
	defer func() {
		for shard, o := range shardMap {
			if lastHash >= shard.Slot0 && lastHash < shard.Slot1 {
				serverConn:=o.Value.(*RespReaderWriter)
				s:="*1\r\n$4\r\nQUIT\r\n"
				err:=serverConn.ProxyWrite(s)
				if err!=nil{
					fmt.Printf("send quit error %v\n",err)
				}
				o.Broken = true //set last one to broken
				fmt.Printf("set conn to broken,lastHash :%d\n", lastHash)
			}
			o.Release()

		}
		client.Close()
	}()
	for {
		select {
		case <-cch:
			return
		case params := <-dch:
			if len(params) == 0 {
				continue
			}
			cmd := params[0].(string)
			if cmd == "PING" {
				client.ProxyWrite("+PONG\r\n")
			} else if cmd == "QUIT" {
				client.ProxyWrite("+OK\r\n")
				client.Close()
				return
			} else if cmd == "ECHO" {
				s:=fmt.Sprintf("$%d%s\r\n",len(params[1].(string)),params[1].(string))
				client.ProxyWrite(s)
			} else {
				if len(params) < 2 {
					fmt.Println(params)
				}
				key := params[1].(string)
				hash := Hash(key) % uint32(1024)
				lastHash = hash
				key = fmt.Sprint(hash) + "_" + key
				//fmt.Println(key)
				s := "*" + fmt.Sprint(len(params)) + "\r\n"
				s += "$" + fmt.Sprint(len(cmd)) + "\r\n" + cmd + "\r\n"
				s += "$" + fmt.Sprint(len(key)) + "\r\n" + key + "\r\n"

				//here should check nil and empty string
				for i := 2; i < len(params); i++ {
					s += "$" + fmt.Sprint(len(params[i].(string))) + "\r\n" + params[i].(string) + "\r\n"
				}
				shard, err := shards.GetShard(hash)
				if err != nil {
					//handle error
					client.Close()
					return
				}
				if ob, ok := shardMap[shard]; !ok {
					fmt.Println("get new shard")
					ob, err = shard.ShardRespPool.Borrow()
					if err != nil {
						client.Close()
						return
					} else if ob != nil {
						shardMap[shard] = ob
					}
				}

				server := shardMap[shard]
				ss := server.Value.(*RespReaderWriter)
				err = ss.ProxyWrite(s)
				if err != nil {
					fmt.Printf("backend server write error:%v\n", err)
					server.Broken = true
					server.Release()
					return
				}
				resp, err := ss.ProxyRead()
				if err != nil {
					fmt.Printf("backend server read error:%v\n", err)
					server.Broken = true
					return
				}
				err = client.ProxyWrite(resp)
				if err != nil {
					fmt.Println("clent error")
					server.Broken = true
					server.Release()
				}
			}

		}

	}

}

func init() {
	shards = NewShards()
}
<commit_msg>release all connection on client exit<commit_after>package main

//presharding
//compute hash ,find target redis instance,construct new key with hashtag
import (
	"errors"
	"fmt"
	"hash/crc32"
	"net"
)

var (
	crcTable *crc32.Table
	target   RespReaderWriter
	shards   Shards
)

func Hash(key string) uint32 {
	crc := crc32.New(crcTable)
	crc.Write([]byte(key))
	v := crc.Sum32()
	return v
}

type Shard struct {
	Host          string
	Slot0         uint32
	Slot1         uint32
	ShardRespPool *ObjectPool
}

type Shards struct {
	Slots        uint32
	ShardServers []Shard
}

func NewShards() Shards {
	sh := Shards{}
	s1 := Shard{}
	s1.Host = "172.31.6.58:6479"
	s1.Slot0 = 0
	s1.Slot1 = 512
	s2 := Shard{}
	s2.Host = "172.31.6.58:6579"
	s2.Slot0 = 512
	s2.Slot1 = 1024
	s1.ShardRespPool = NewProxyClientPool(s1.Host)
	s2.ShardRespPool = NewProxyClientPool(s2.Host)
	sh.ShardServers = []Shard{s1, s2}
	return sh

}
func (shards Shards) GetPooledObject(hash uint32) (obj *PooledObject, shard Shard, err error) {
	for _, shard := range shards.ShardServers {
		if hash >= shard.Slot0 && hash < shard.Slot1 {
			fmt.Printf("get shard:%v\n", shard)
			obj, er := shard.ShardRespPool.Borrow()
			return obj, shard, er
		}
	}
	return nil, Shard{}, nil
}

func (shards Shards) GetShard(hash uint32) (Shard, error) {
	for i, shard := range shards.ShardServers {
		if hash >= shard.Slot0 && hash < shard.Slot1 {
			return shards.ShardServers[i], nil
		}
	}
	return Shard{}, errors.New("no shard found")
}

func HandleConn(conn net.Conn) {
	//	timeout := time.Now()
	//	timeout.Add(10 *time.Minute )
	//	conn.SetReadDeadline(timeout)
	//conn.
	client := NewRespReadWriter(conn)
	cch := make(chan int) //close flag chan
	dch := client.LoopRead(cch)
	shardMap := make(map[Shard]*PooledObject)
	var lastHash uint32
	defer func() {
		for _, o := range shardMap {
			//if lastHash >= shard.Slot0 && lastHash < shard.Slot1 {
				serverConn:=o.Value.(*RespReaderWriter)
				s:="*1\r\n$4\r\nQUIT\r\n"
				err:=serverConn.ProxyWrite(s)
				if err!=nil{
					fmt.Printf("send quit error %v\n",err)
				}
				o.Broken = true //set last one to broken
				fmt.Printf("set conn to broken,lastHash :%d\n", lastHash)
		//	}
			o.Release()

		}
		client.Close()
	}()
	for {
		select {
		case <-cch:
			return
		case params := <-dch:
			if len(params) == 0 {
				continue
			}
			cmd := params[0].(string)
			if cmd == "PING" {
				client.ProxyWrite("+PONG\r\n")
			} else if cmd == "QUIT" {
				client.ProxyWrite("+OK\r\n")
				client.Close()
				return
			} else if cmd == "ECHO" {
				s:=fmt.Sprintf("$%d%s\r\n",len(params[1].(string)),params[1].(string))
				client.ProxyWrite(s)
			} else {
				if len(params) < 2 {
					fmt.Println(params)
				}
				key := params[1].(string)
				hash := Hash(key) % uint32(1024)
				lastHash = hash
				key = fmt.Sprint(hash) + "_" + key
				//fmt.Println(key)
				s := "*" + fmt.Sprint(len(params)) + "\r\n"
				s += "$" + fmt.Sprint(len(cmd)) + "\r\n" + cmd + "\r\n"
				s += "$" + fmt.Sprint(len(key)) + "\r\n" + key + "\r\n"

				//here should check nil and empty string
				for i := 2; i < len(params); i++ {
					s += "$" + fmt.Sprint(len(params[i].(string))) + "\r\n" + params[i].(string) + "\r\n"
				}
				shard, err := shards.GetShard(hash)
				if err != nil {
					//handle error
					client.Close()
					return
				}
				if ob, ok := shardMap[shard]; !ok {
					fmt.Println("get new shard")
					ob, err = shard.ShardRespPool.Borrow()
					if err != nil {
						client.Close()
						return
					} else if ob != nil {
						shardMap[shard] = ob
					}
				}

				server := shardMap[shard]
				ss := server.Value.(*RespReaderWriter)
				err = ss.ProxyWrite(s)
				if err != nil {
					fmt.Printf("backend server write error:%v\n", err)
					server.Broken = true
					server.Release()
					return
				}
				resp, err := ss.ProxyRead()
				if err != nil {
					fmt.Printf("backend server read error:%v\n", err)
					server.Broken = true
					return
				}
				err = client.ProxyWrite(resp)
				if err != nil {
					fmt.Println("clent error")
					server.Broken = true
					server.Release()
				}
			}

		}

	}

}

func init() {
	shards = NewShards()
}
<|endoftext|>