<commit_before>package cobra

import (
	"bytes"
	"fmt"
	"io"
	"os"
)

// GenZshCompletionFile generates zsh completion file including descriptions.
func (c *Command) GenZshCompletionFile(filename string) error {
	return c.genZshCompletionFile(filename, true)
}

// GenZshCompletion generates zsh completion file including descriptions
// and writes it to the passed writer.
func (c *Command) GenZshCompletion(w io.Writer) error {
	return c.genZshCompletion(w, true)
}

// GenZshCompletionFileNoDesc generates zsh completion file without descriptions.
func (c *Command) GenZshCompletionFileNoDesc(filename string) error {
	return c.genZshCompletionFile(filename, false)
}

// GenZshCompletionNoDesc generates zsh completion file without descriptions
// and writes it to the passed writer.
func (c *Command) GenZshCompletionNoDesc(w io.Writer) error {
	return c.genZshCompletion(w, false)
}

// MarkZshCompPositionalArgumentFile only worked for zsh and its behavior was
// not consistent with Bash completion. It has therefore been disabled.
// Instead, when no other completion is specified, file completion is done by
// default for every argument. One can disable file completion on a per-argument
// basis by using ValidArgsFunction and ShellCompDirectiveNoFileComp.
// To achieve file extension filtering, one can use ValidArgsFunction and
// ShellCompDirectiveFilterFileExt.
//
// Deprecated
func (c *Command) MarkZshCompPositionalArgumentFile(argPosition int, patterns ...string) error {
	return nil
}

// MarkZshCompPositionalArgumentWords only worked for zsh. It has therefore
// been disabled.
// To achieve the same behavior across all shells, one can use
// ValidArgs (for the first argument only) or ValidArgsFunction for
// any argument (can include the first one also).
//
// Deprecated
func (c *Command) MarkZshCompPositionalArgumentWords(argPosition int, words ...string) error {
	return nil
}

func (c *Command) genZshCompletionFile(filename string, includeDesc bool) error {
	outFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer outFile.Close()

	return c.genZshCompletion(outFile, includeDesc)
}

func (c *Command) genZshCompletion(w io.Writer, includeDesc bool) error {
	buf := new(bytes.Buffer)
	genZshComp(buf, c.Name(), includeDesc)
	_, err := buf.WriteTo(w)
	return err
}

func genZshComp(buf io.StringWriter, name string, includeDesc bool) {
	compCmd := ShellCompRequestCmd
	if !includeDesc {
		compCmd = ShellCompNoDescRequestCmd
	}
	WriteStringAndCheck(buf, fmt.Sprintf(`#compdef _%[1]s %[1]s

# zsh completion for %-36[1]s -*- shell-script -*-

__%[1]s_debug()
{
    local file="$BASH_COMP_DEBUG_FILE"
    if [[ -n ${file} ]]; then
        echo "$*" >> "${file}"
    fi
}

_%[1]s()
{
    local shellCompDirectiveError=%[3]d
    local shellCompDirectiveNoSpace=%[4]d
    local shellCompDirectiveNoFileComp=%[5]d
    local shellCompDirectiveFilterFileExt=%[6]d
    local shellCompDirectiveFilterDirs=%[7]d

    local lastParam lastChar flagPrefix requestComp out directive comp lastComp noSpace
    local -a completions

    __%[1]s_debug "\n========= starting completion logic =========="
    __%[1]s_debug "CURRENT: ${CURRENT}, words[*]: ${words[*]}"

    # The user could have moved the cursor backwards on the command-line.
    # We need to trigger completion from the $CURRENT location, so we need
    # to truncate the command-line ($words) up to the $CURRENT location.
    # (We cannot use $CURSOR as its value does not work when a command is an alias.)
    words=("${=words[1,CURRENT]}")
    __%[1]s_debug "Truncated words[*]: ${words[*]},"

    lastParam=${words[-1]}
    lastChar=${lastParam[-1]}
    __%[1]s_debug "lastParam: ${lastParam}, lastChar: ${lastChar}"

    # For zsh, when completing a flag with an = (e.g., %[1]s -n=<TAB>)
    # completions must be prefixed with the flag
    setopt local_options BASH_REMATCH
    if [[ "${lastParam}" =~ '-.*=' ]]; then
        # We are dealing with a flag with an =
        flagPrefix="-P ${BASH_REMATCH}"
    fi

    # Prepare the command to obtain completions
    requestComp="${words[1]} %[2]s ${words[2,-1]}"
    if [ "${lastChar}" = "" ]; then
        # If the last parameter is complete (there is a space following it)
        # We add an extra empty parameter so we can indicate this to the go completion code.
        __%[1]s_debug "Adding extra empty parameter"
        requestComp="${requestComp} \"\""
    fi

    __%[1]s_debug "About to call: eval ${requestComp}"

    # Use eval to handle any environment variables and such
    out=$(eval ${requestComp} 2>/dev/null)
    __%[1]s_debug "completion output: ${out}"

    # Extract the directive integer following a : from the last line
    local lastLine
    while IFS='\n' read -r line; do
        lastLine=${line}
    done < <(printf "%%s\n" "${out[@]}")
    __%[1]s_debug "last line: ${lastLine}"

    if [ "${lastLine[1]}" = : ]; then
        directive=${lastLine[2,-1]}
        # Remove the directive including the : and the newline
        local suffix
        (( suffix=${#lastLine}+2))
        out=${out[1,-$suffix]}
    else
        # There is no directive specified.  Leave $out as is.
        __%[1]s_debug "No directive found.  Setting do default"
        directive=0
    fi

    __%[1]s_debug "directive: ${directive}"
    __%[1]s_debug "completions: ${out}"
    __%[1]s_debug "flagPrefix: ${flagPrefix}"

    if [ $((directive & shellCompDirectiveError)) -ne 0 ]; then
        __%[1]s_debug "Completion received error. Ignoring completions."
        return
    fi

    while IFS='\n' read -r comp; do
        if [ -n "$comp" ]; then
            # If requested, completions are returned with a description.
            # The description is preceded by a TAB character.
            # For zsh's _describe, we need to use a : instead of a TAB.
            # We first need to escape any : as part of the completion itself.
            comp=${comp//:/\\:}

            local tab=$(printf '\t')
            comp=${comp//$tab/:}

            __%[1]s_debug "Adding completion: ${comp}"
            completions+=${comp}
            lastComp=$comp
        fi
    done < <(printf "%%s\n" "${out[@]}")

    if [ $((directive & shellCompDirectiveNoSpace)) -ne 0 ]; then
        __%[1]s_debug "Activating nospace."
        noSpace="-S ''"
    fi

    if [ $((directive & shellCompDirectiveFilterFileExt)) -ne 0 ]; then
        # File extension filtering
        local filteringCmd
        filteringCmd='_files'
        for filter in ${completions[@]}; do
            if [ ${filter[1]} != '*' ]; then
                # zsh requires a glob pattern to do file filtering
                filter="\*.$filter"
            fi
            filteringCmd+=" -g $filter"
        done
        filteringCmd+=" ${flagPrefix}"

        __%[1]s_debug "File filtering command: $filteringCmd"
        _arguments '*:filename:'"$filteringCmd"
    elif [ $((directive & shellCompDirectiveFilterDirs)) -ne 0 ]; then
        # File completion for directories only
        local subdir
        subdir="${completions[1]}"
        if [ -n "$subdir" ]; then
            __%[1]s_debug "Listing directories in $subdir"
            pushd "${subdir}" >/dev/null 2>&1
        else
            __%[1]s_debug "Listing directories in ."
        fi

        local result
        _arguments '*:dirname:_files -/'" ${flagPrefix}"
        result=$?
        if [ -n "$subdir" ]; then
            popd >/dev/null 2>&1
        fi
        return $result
    else
        __%[1]s_debug "Calling _describe"
        if eval _describe "completions" completions $flagPrefix $noSpace; then
            __%[1]s_debug "_describe found some completions"

            # Return the success of having called _describe
            return 0
        else
            __%[1]s_debug "_describe did not find completions."
            __%[1]s_debug "Checking if we should do file completion."
            if [ $((directive & shellCompDirectiveNoFileComp)) -ne 0 ]; then
                __%[1]s_debug "deactivating file completion"

                # We must return an error code here to let zsh know that there were no
                # completions found by _describe; this is what will trigger other
                # matching algorithms to attempt to find completions.
                # For example zsh can match letters in the middle of words.
                return 1
            else
                # Perform file completion
                __%[1]s_debug "Activating file completion"

                # We must return the result of this command, so it must be the
                # last command, or else we must store its result to return it.
                _arguments '*:filename:_files'" ${flagPrefix}"
            fi
        fi
    fi
}

# don't run the completion function when being source-ed or eval-ed
if [ "$funcstack[1]" = "_%[1]s" ]; then
    _%[1]s
fi
`, name, compCmd,
		ShellCompDirectiveError, ShellCompDirectiveNoSpace, ShellCompDirectiveNoFileComp,
		ShellCompDirectiveFilterFileExt, ShellCompDirectiveFilterDirs))
}
<commit_msg>Wrap printf tab with quotes (#1665)<commit_after>package cobra

import (
	"bytes"
	"fmt"
	"io"
	"os"
)

// GenZshCompletionFile generates zsh completion file including descriptions.
func (c *Command) GenZshCompletionFile(filename string) error {
	return c.genZshCompletionFile(filename, true)
}

// GenZshCompletion generates zsh completion file including descriptions
// and writes it to the passed writer.
func (c *Command) GenZshCompletion(w io.Writer) error {
	return c.genZshCompletion(w, true)
}

// GenZshCompletionFileNoDesc generates zsh completion file without descriptions.
func (c *Command) GenZshCompletionFileNoDesc(filename string) error {
	return c.genZshCompletionFile(filename, false)
}

// GenZshCompletionNoDesc generates zsh completion file without descriptions
// and writes it to the passed writer.
func (c *Command) GenZshCompletionNoDesc(w io.Writer) error {
	return c.genZshCompletion(w, false)
}

// MarkZshCompPositionalArgumentFile only worked for zsh and its behavior was
// not consistent with Bash completion. It has therefore been disabled.
// Instead, when no other completion is specified, file completion is done by
// default for every argument. One can disable file completion on a per-argument
// basis by using ValidArgsFunction and ShellCompDirectiveNoFileComp.
// To achieve file extension filtering, one can use ValidArgsFunction and
// ShellCompDirectiveFilterFileExt.
//
// Deprecated
func (c *Command) MarkZshCompPositionalArgumentFile(argPosition int, patterns ...string) error {
	return nil
}

// MarkZshCompPositionalArgumentWords only worked for zsh. It has therefore
// been disabled.
// To achieve the same behavior across all shells, one can use
// ValidArgs (for the first argument only) or ValidArgsFunction for
// any argument (can include the first one also).
//
// Deprecated
func (c *Command) MarkZshCompPositionalArgumentWords(argPosition int, words ...string) error {
	return nil
}

func (c *Command) genZshCompletionFile(filename string, includeDesc bool) error {
	outFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer outFile.Close()

	return c.genZshCompletion(outFile, includeDesc)
}

func (c *Command) genZshCompletion(w io.Writer, includeDesc bool) error {
	buf := new(bytes.Buffer)
	genZshComp(buf, c.Name(), includeDesc)
	_, err := buf.WriteTo(w)
	return err
}

func genZshComp(buf io.StringWriter, name string, includeDesc bool) {
	compCmd := ShellCompRequestCmd
	if !includeDesc {
		compCmd = ShellCompNoDescRequestCmd
	}
	WriteStringAndCheck(buf, fmt.Sprintf(`#compdef _%[1]s %[1]s

# zsh completion for %-36[1]s -*- shell-script -*-

__%[1]s_debug()
{
    local file="$BASH_COMP_DEBUG_FILE"
    if [[ -n ${file} ]]; then
        echo "$*" >> "${file}"
    fi
}

_%[1]s()
{
    local shellCompDirectiveError=%[3]d
    local shellCompDirectiveNoSpace=%[4]d
    local shellCompDirectiveNoFileComp=%[5]d
    local shellCompDirectiveFilterFileExt=%[6]d
    local shellCompDirectiveFilterDirs=%[7]d

    local lastParam lastChar flagPrefix requestComp out directive comp lastComp noSpace
    local -a completions

    __%[1]s_debug "\n========= starting completion logic =========="
    __%[1]s_debug "CURRENT: ${CURRENT}, words[*]: ${words[*]}"

    # The user could have moved the cursor backwards on the command-line.
    # We need to trigger completion from the $CURRENT location, so we need
    # to truncate the command-line ($words) up to the $CURRENT location.
    # (We cannot use $CURSOR as its value does not work when a command is an alias.)
    words=("${=words[1,CURRENT]}")
    __%[1]s_debug "Truncated words[*]: ${words[*]},"

    lastParam=${words[-1]}
    lastChar=${lastParam[-1]}
    __%[1]s_debug "lastParam: ${lastParam}, lastChar: ${lastChar}"

    # For zsh, when completing a flag with an = (e.g., %[1]s -n=<TAB>)
    # completions must be prefixed with the flag
    setopt local_options BASH_REMATCH
    if [[ "${lastParam}" =~ '-.*=' ]]; then
        # We are dealing with a flag with an =
        flagPrefix="-P ${BASH_REMATCH}"
    fi

    # Prepare the command to obtain completions
    requestComp="${words[1]} %[2]s ${words[2,-1]}"
    if [ "${lastChar}" = "" ]; then
        # If the last parameter is complete (there is a space following it)
        # We add an extra empty parameter so we can indicate this to the go completion code.
        __%[1]s_debug "Adding extra empty parameter"
        requestComp="${requestComp} \"\""
    fi

    __%[1]s_debug "About to call: eval ${requestComp}"

    # Use eval to handle any environment variables and such
    out=$(eval ${requestComp} 2>/dev/null)
    __%[1]s_debug "completion output: ${out}"

    # Extract the directive integer following a : from the last line
    local lastLine
    while IFS='\n' read -r line; do
        lastLine=${line}
    done < <(printf "%%s\n" "${out[@]}")
    __%[1]s_debug "last line: ${lastLine}"

    if [ "${lastLine[1]}" = : ]; then
        directive=${lastLine[2,-1]}
        # Remove the directive including the : and the newline
        local suffix
        (( suffix=${#lastLine}+2))
        out=${out[1,-$suffix]}
    else
        # There is no directive specified.  Leave $out as is.
        __%[1]s_debug "No directive found.  Setting do default"
        directive=0
    fi

    __%[1]s_debug "directive: ${directive}"
    __%[1]s_debug "completions: ${out}"
    __%[1]s_debug "flagPrefix: ${flagPrefix}"

    if [ $((directive & shellCompDirectiveError)) -ne 0 ]; then
        __%[1]s_debug "Completion received error. Ignoring completions."
        return
    fi

    while IFS='\n' read -r comp; do
        if [ -n "$comp" ]; then
            # If requested, completions are returned with a description.
            # The description is preceded by a TAB character.
            # For zsh's _describe, we need to use a : instead of a TAB.
            # We first need to escape any : as part of the completion itself.
            comp=${comp//:/\\:}

            local tab="$(printf '\t')"
            comp=${comp//$tab/:}

            __%[1]s_debug "Adding completion: ${comp}"
            completions+=${comp}
            lastComp=$comp
        fi
    done < <(printf "%%s\n" "${out[@]}")

    if [ $((directive & shellCompDirectiveNoSpace)) -ne 0 ]; then
        __%[1]s_debug "Activating nospace."
        noSpace="-S ''"
    fi

    if [ $((directive & shellCompDirectiveFilterFileExt)) -ne 0 ]; then
        # File extension filtering
        local filteringCmd
        filteringCmd='_files'
        for filter in ${completions[@]}; do
            if [ ${filter[1]} != '*' ]; then
                # zsh requires a glob pattern to do file filtering
                filter="\*.$filter"
            fi
            filteringCmd+=" -g $filter"
        done
        filteringCmd+=" ${flagPrefix}"

        __%[1]s_debug "File filtering command: $filteringCmd"
        _arguments '*:filename:'"$filteringCmd"
    elif [ $((directive & shellCompDirectiveFilterDirs)) -ne 0 ]; then
        # File completion for directories only
        local subdir
        subdir="${completions[1]}"
        if [ -n "$subdir" ]; then
            __%[1]s_debug "Listing directories in $subdir"
            pushd "${subdir}" >/dev/null 2>&1
        else
            __%[1]s_debug "Listing directories in ."
        fi

        local result
        _arguments '*:dirname:_files -/'" ${flagPrefix}"
        result=$?
        if [ -n "$subdir" ]; then
            popd >/dev/null 2>&1
        fi
        return $result
    else
        __%[1]s_debug "Calling _describe"
        if eval _describe "completions" completions $flagPrefix $noSpace; then
            __%[1]s_debug "_describe found some completions"

            # Return the success of having called _describe
            return 0
        else
            __%[1]s_debug "_describe did not find completions."
            __%[1]s_debug "Checking if we should do file completion."
            if [ $((directive & shellCompDirectiveNoFileComp)) -ne 0 ]; then
                __%[1]s_debug "deactivating file completion"

                # We must return an error code here to let zsh know that there were no
                # completions found by _describe; this is what will trigger other
                # matching algorithms to attempt to find completions.
                # For example zsh can match letters in the middle of words.
                return 1
            else
                # Perform file completion
                __%[1]s_debug "Activating file completion"

                # We must return the result of this command, so it must be the
                # last command, or else we must store its result to return it.
                _arguments '*:filename:_files'" ${flagPrefix}"
            fi
        fi
    fi
}

# don't run the completion function when being source-ed or eval-ed
if [ "$funcstack[1]" = "_%[1]s" ]; then
    _%[1]s
fi
`, name, compCmd,
		ShellCompDirectiveError, ShellCompDirectiveNoSpace, ShellCompDirectiveNoFileComp,
		ShellCompDirectiveFilterFileExt, ShellCompDirectiveFilterDirs))
}
<|endoftext|>


<commit_before>package encoding_test

import (
	"testing"

	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/serial"
	"v2ray.com/core/common/uuid"
	"v2ray.com/core/proxy/vmess"
	. "v2ray.com/core/proxy/vmess/encoding"
	. "v2ray.com/ext/assert"
)

func TestRequestSerialization(t *testing.T) {
	assert := With(t)

	user := &protocol.User{
		Level: 0,
		Email: "test@v2ray.com",
	}
	id := uuid.New()
	account := &vmess.Account{
		Id:      id.String(),
		AlterId: 0,
	}
	user.Account = serial.ToTypedMessage(account)

	expectedRequest := &protocol.RequestHeader{
		Version:  1,
		User:     user,
		Command:  protocol.RequestCommandTCP,
		Address:  net.DomainAddress("www.v2ray.com"),
		Port:     net.Port(443),
		Security: protocol.Security(protocol.SecurityType_AES128_GCM),
	}

	buffer := buf.New()
	client := NewClientSession(protocol.DefaultIDHash)
	common.Must(client.EncodeRequestHeader(expectedRequest, buffer))

	buffer2 := buf.New()
	buffer2.Append(buffer.Bytes())

	sessionHistory := NewSessionHistory()
	defer common.Close(sessionHistory)

	userValidator := vmess.NewTimedUserValidator(protocol.DefaultIDHash)
	userValidator.Add(user)
	defer common.Close(userValidator)

	server := NewServerSession(userValidator, sessionHistory)
	actualRequest, err := server.DecodeRequestHeader(buffer)
	assert(err, IsNil)

	assert(expectedRequest.Version, Equals, actualRequest.Version)
	assert(byte(expectedRequest.Command), Equals, byte(actualRequest.Command))
	assert(byte(expectedRequest.Option), Equals, byte(actualRequest.Option))
	assert(expectedRequest.Address, Equals, actualRequest.Address)
	assert(expectedRequest.Port, Equals, actualRequest.Port)
	assert(byte(expectedRequest.Security), Equals, byte(actualRequest.Security))

	_, err = server.DecodeRequestHeader(buffer2)
	// anti replay attack
	assert(err, IsNotNil)
}

func TestInvalidRequest(t *testing.T) {
	assert := With(t)

	user := &protocol.User{
		Level: 0,
		Email: "test@v2ray.com",
	}
	id := uuid.New()
	account := &vmess.Account{
		Id:      id.String(),
		AlterId: 0,
	}
	user.Account = serial.ToTypedMessage(account)

	expectedRequest := &protocol.RequestHeader{
		Version:  1,
		User:     user,
		Command:  protocol.RequestCommand(100),
		Address:  net.DomainAddress("www.v2ray.com"),
		Port:     net.Port(443),
		Security: protocol.Security(protocol.SecurityType_AES128_GCM),
	}

	buffer := buf.New()
	client := NewClientSession(protocol.DefaultIDHash)
	common.Must(client.EncodeRequestHeader(expectedRequest, buffer))

	buffer2 := buf.New()
	buffer2.Append(buffer.Bytes())

	sessionHistory := NewSessionHistory()
	defer common.Close(sessionHistory)

	userValidator := vmess.NewTimedUserValidator(protocol.DefaultIDHash)
	userValidator.Add(user)
	defer common.Close(userValidator)

	server := NewServerSession(userValidator, sessionHistory)
	_, err := server.DecodeRequestHeader(buffer)
	assert(err, IsNotNil)
}

func TestMuxRequest(t *testing.T) {
	assert := With(t)

	user := &protocol.User{
		Level: 0,
		Email: "test@v2ray.com",
	}
	id := uuid.New()
	account := &vmess.Account{
		Id:      id.String(),
		AlterId: 0,
	}
	user.Account = serial.ToTypedMessage(account)

	expectedRequest := &protocol.RequestHeader{
		Version:  1,
		User:     user,
		Command:  protocol.RequestCommandMux,
		Security: protocol.Security(protocol.SecurityType_AES128_GCM),
	}

	buffer := buf.New()
	client := NewClientSession(protocol.DefaultIDHash)
	common.Must(client.EncodeRequestHeader(expectedRequest, buffer))

	buffer2 := buf.New()
	buffer2.Append(buffer.Bytes())

	sessionHistory := NewSessionHistory()
	defer common.Close(sessionHistory)

	userValidator := vmess.NewTimedUserValidator(protocol.DefaultIDHash)
	userValidator.Add(user)
	defer common.Close(userValidator)

	server := NewServerSession(userValidator, sessionHistory)
	actualRequest, err := server.DecodeRequestHeader(buffer)
	assert(err, IsNil)

	assert(expectedRequest.Version, Equals, actualRequest.Version)
	assert(byte(expectedRequest.Command), Equals, byte(actualRequest.Command))
	assert(byte(expectedRequest.Option), Equals, byte(actualRequest.Option))
	assert(byte(expectedRequest.Security), Equals, byte(actualRequest.Security))
}
<commit_msg>fix test break<commit_after>package encoding_test

import (
	"testing"

	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/serial"
	"v2ray.com/core/common/uuid"
	"v2ray.com/core/proxy/vmess"
	. "v2ray.com/core/proxy/vmess/encoding"
	. "v2ray.com/ext/assert"
)

func TestRequestSerialization(t *testing.T) {
	assert := With(t)

	user := &protocol.User{
		Level: 0,
		Email: "test@v2ray.com",
	}
	id := uuid.New()
	account := &vmess.Account{
		Id:      id.String(),
		AlterId: 0,
	}
	user.Account = serial.ToTypedMessage(account)

	expectedRequest := &protocol.RequestHeader{
		Version:  1,
		User:     user,
		Command:  protocol.RequestCommandTCP,
		Address:  net.DomainAddress("www.v2ray.com"),
		Port:     net.Port(443),
		Security: protocol.SecurityType_AES128_GCM,
	}

	buffer := buf.New()
	client := NewClientSession(protocol.DefaultIDHash)
	common.Must(client.EncodeRequestHeader(expectedRequest, buffer))

	buffer2 := buf.New()
	buffer2.Append(buffer.Bytes())

	sessionHistory := NewSessionHistory()
	defer common.Close(sessionHistory)

	userValidator := vmess.NewTimedUserValidator(protocol.DefaultIDHash)
	userValidator.Add(user)
	defer common.Close(userValidator)

	server := NewServerSession(userValidator, sessionHistory)
	actualRequest, err := server.DecodeRequestHeader(buffer)
	assert(err, IsNil)

	assert(expectedRequest.Version, Equals, actualRequest.Version)
	assert(byte(expectedRequest.Command), Equals, byte(actualRequest.Command))
	assert(byte(expectedRequest.Option), Equals, byte(actualRequest.Option))
	assert(expectedRequest.Address, Equals, actualRequest.Address)
	assert(expectedRequest.Port, Equals, actualRequest.Port)
	assert(byte(expectedRequest.Security), Equals, byte(actualRequest.Security))

	_, err = server.DecodeRequestHeader(buffer2)
	// anti replay attack
	assert(err, IsNotNil)
}

func TestInvalidRequest(t *testing.T) {
	assert := With(t)

	user := &protocol.User{
		Level: 0,
		Email: "test@v2ray.com",
	}
	id := uuid.New()
	account := &vmess.Account{
		Id:      id.String(),
		AlterId: 0,
	}
	user.Account = serial.ToTypedMessage(account)

	expectedRequest := &protocol.RequestHeader{
		Version:  1,
		User:     user,
		Command:  protocol.RequestCommand(100),
		Address:  net.DomainAddress("www.v2ray.com"),
		Port:     net.Port(443),
		Security: protocol.SecurityType_AES128_GCM,
	}

	buffer := buf.New()
	client := NewClientSession(protocol.DefaultIDHash)
	common.Must(client.EncodeRequestHeader(expectedRequest, buffer))

	buffer2 := buf.New()
	buffer2.Append(buffer.Bytes())

	sessionHistory := NewSessionHistory()
	defer common.Close(sessionHistory)

	userValidator := vmess.NewTimedUserValidator(protocol.DefaultIDHash)
	userValidator.Add(user)
	defer common.Close(userValidator)

	server := NewServerSession(userValidator, sessionHistory)
	_, err := server.DecodeRequestHeader(buffer)
	assert(err, IsNotNil)
}

func TestMuxRequest(t *testing.T) {
	assert := With(t)

	user := &protocol.User{
		Level: 0,
		Email: "test@v2ray.com",
	}
	id := uuid.New()
	account := &vmess.Account{
		Id:      id.String(),
		AlterId: 0,
	}
	user.Account = serial.ToTypedMessage(account)

	expectedRequest := &protocol.RequestHeader{
		Version:  1,
		User:     user,
		Command:  protocol.RequestCommandMux,
		Security: protocol.SecurityType_AES128_GCM,
	}

	buffer := buf.New()
	client := NewClientSession(protocol.DefaultIDHash)
	common.Must(client.EncodeRequestHeader(expectedRequest, buffer))

	buffer2 := buf.New()
	buffer2.Append(buffer.Bytes())

	sessionHistory := NewSessionHistory()
	defer common.Close(sessionHistory)

	userValidator := vmess.NewTimedUserValidator(protocol.DefaultIDHash)
	userValidator.Add(user)
	defer common.Close(userValidator)

	server := NewServerSession(userValidator, sessionHistory)
	actualRequest, err := server.DecodeRequestHeader(buffer)
	assert(err, IsNil)

	assert(expectedRequest.Version, Equals, actualRequest.Version)
	assert(byte(expectedRequest.Command), Equals, byte(actualRequest.Command))
	assert(byte(expectedRequest.Option), Equals, byte(actualRequest.Option))
	assert(byte(expectedRequest.Security), Equals, byte(actualRequest.Security))
}
<|endoftext|>


<commit_before>package adifparser

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"strconv"
	"strings"
)

// Public interface for ADIFRecords
type ADIFRecord interface {
	// Print as ADIF String
	ToString() string
	// Fingerprint for duplication detection
	Fingerprint() string
	// Setters and getters
	GetValue(string) (string, error)
	SetValue(string, string)
	// Get all of the present field names
	GetFields() []string
}

// Internal implementation for ADIFRecord
type baseADIFRecord struct {
	values map[string]string
}

type fieldData struct {
	name     string
	value    string
	typecode byte
	hasType  bool
}

// Errors
var NoData = errors.New("No data to parse.")
var NoSuchField = errors.New("No such field.")
var InvalidField = errors.New("Invalid field definition.")

// Create a new ADIFRecord from scratch
func NewADIFRecord() *baseADIFRecord {
	record := &baseADIFRecord{}
	record.values = make(map[string]string)
	return record
}

// Parse an ADIFRecord
func ParseADIFRecord(buf []byte) (*baseADIFRecord, error) {
	record := NewADIFRecord()

	if len(buf) == 0 {
		return nil, NoData
	}

	for len(buf) > 0 {
		var data *fieldData
		var err error
		if buf[0] == '/' && buf[1] == '/' { // Recognize comments and skip them.
			end_of_line := bytes.IndexByte(buf, 13)
			if end_of_line > -1 {
				buf = buf[end_of_line+1:]
			} else { // The comment ends the record so there's no end-of-line.
				buf = buf[len(buf):]
			}
		} else {
			data, buf, err = getNextField(buf)
			if err != nil {
				return nil, err
			}
			// TODO: accomodate types
			record.values[data.name] = data.value
		}
	}

	return record, nil
}

// Get the next field, return field data, leftover data, and optional error
func getNextField(buf []byte) (*fieldData, []byte, error) {
	data := &fieldData{}

	// Extract name
	start_of_name := tagStartPos(buf) + 1
	end_of_name := bytes.IndexByte(buf, ':')
	end_of_tag := bytes.IndexByte(buf, '>')
	if end_of_name == -1 || end_of_tag < end_of_name || end_of_name < start_of_name {
		return nil, buf, InvalidField
	}
	data.name = strings.ToLower(string(buf[start_of_name:end_of_name]))
	buf = buf[end_of_name+1:]
	// Adjust to new buffer
	end_of_tag -= end_of_name + 1

	// Length
	var length int
	var err error
	start_type := bytes.IndexByte(buf, ':')
	if start_type == -1 || start_type > end_of_tag {
		end_of_length := bytes.IndexByte(buf, '>')
		length, err = strconv.Atoi(string(buf[:end_of_length]))
		buf = buf[end_of_length+1:]
		data.hasType = false
	} else {
		length, err = strconv.Atoi(string(buf[:start_type]))
		data.typecode = buf[start_type+1]
		buf = buf[start_type+3:]
		data.hasType = true
	}
	if err != nil {
		// TODO: log the error
		return nil, buf, err
	}

	// Value
	data.value = string(buf[:length])
	buf = bytes.TrimSpace(buf[length:])

	return data, buf, nil
}

func serializeField(name string, value string) string {
	return fmt.Sprintf("<%s:%d>%s", name, len(value), value)
}

// Print an ADIFRecord as a string
func (r *baseADIFRecord) ToString() string {
	var record bytes.Buffer
	for _, n := range ADIFfieldOrder {
		if v, ok := r.values[n]; ok {
			record.WriteString(serializeField(n, v))
		}
	}
	// Handle custom fields
	for n, v := range r.values {
		if !isStandardADIFField(n) {
			record.WriteString(serializeField(n, v))
		}
	}
	return record.String()
}

// Get fingerprint of ADIFRecord
func (r *baseADIFRecord) Fingerprint() string {
	fpfields := []string{
		"call", "station_callsign", "band",
		"freq", "mode", "qso_date", "time_on",
		"time_off"}
	fpvals := make([]string, 0, len(fpfields))
	for _, f := range fpfields {
		if n, ok := r.values[f]; ok {
			fpvals = append(fpvals, n)
		}
	}
	fptext := strings.Join(fpvals, "|")
	h := sha256.New()
	h.Write([]byte(fptext))
	return hex.EncodeToString(h.Sum(nil))
}

// Get a value
func (r *baseADIFRecord) GetValue(name string) (string, error) {
	if v, ok := r.values[name]; ok {
		return v, nil
	}
	return "", NoSuchField
}

// Set a value
func (r *baseADIFRecord) SetValue(name string, value string) {
	r.values[strings.ToLower(name)] = value
}

// Get all of the present field names
func (r *baseADIFRecord) GetFields() []string {
	keys := make([]string, len(r.values))
	i := 0
	for k := range r.values {
		keys[i] = k
		i++
	}
	return keys
}
<commit_msg>Add ADIFRecord.DeleteField()<commit_after>package adifparser

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"strconv"
	"strings"
)

// Public interface for ADIFRecords
type ADIFRecord interface {
	// Print as ADIF String
	ToString() string
	// Fingerprint for duplication detection
	Fingerprint() string
	// Setters and getters
	GetValue(string) (string, error)
	SetValue(string, string)
	// Get all of the present field names
	GetFields() []string
}

// Internal implementation for ADIFRecord
type baseADIFRecord struct {
	values map[string]string
}

type fieldData struct {
	name     string
	value    string
	typecode byte
	hasType  bool
}

// Errors
var NoData = errors.New("No data to parse.")
var NoSuchField = errors.New("No such field.")
var InvalidField = errors.New("Invalid field definition.")

// Create a new ADIFRecord from scratch
func NewADIFRecord() *baseADIFRecord {
	record := &baseADIFRecord{}
	record.values = make(map[string]string)
	return record
}

// Parse an ADIFRecord
func ParseADIFRecord(buf []byte) (*baseADIFRecord, error) {
	record := NewADIFRecord()

	if len(buf) == 0 {
		return nil, NoData
	}

	for len(buf) > 0 {
		var data *fieldData
		var err error
		if buf[0] == '/' && buf[1] == '/' { // Recognize comments and skip them.
			end_of_line := bytes.IndexByte(buf, 13)
			if end_of_line > -1 {
				buf = buf[end_of_line+1:]
			} else { // The comment ends the record so there's no end-of-line.
				buf = buf[len(buf):]
			}
		} else {
			data, buf, err = getNextField(buf)
			if err != nil {
				return nil, err
			}
			// TODO: accomodate types
			record.values[data.name] = data.value
		}
	}

	return record, nil
}

// Get the next field, return field data, leftover data, and optional error
func getNextField(buf []byte) (*fieldData, []byte, error) {
	data := &fieldData{}

	// Extract name
	start_of_name := tagStartPos(buf) + 1
	end_of_name := bytes.IndexByte(buf, ':')
	end_of_tag := bytes.IndexByte(buf, '>')
	if end_of_name == -1 || end_of_tag < end_of_name || end_of_name < start_of_name {
		return nil, buf, InvalidField
	}
	data.name = strings.ToLower(string(buf[start_of_name:end_of_name]))
	buf = buf[end_of_name+1:]
	// Adjust to new buffer
	end_of_tag -= end_of_name + 1

	// Length
	var length int
	var err error
	start_type := bytes.IndexByte(buf, ':')
	if start_type == -1 || start_type > end_of_tag {
		end_of_length := bytes.IndexByte(buf, '>')
		length, err = strconv.Atoi(string(buf[:end_of_length]))
		buf = buf[end_of_length+1:]
		data.hasType = false
	} else {
		length, err = strconv.Atoi(string(buf[:start_type]))
		data.typecode = buf[start_type+1]
		buf = buf[start_type+3:]
		data.hasType = true
	}
	if err != nil {
		// TODO: log the error
		return nil, buf, err
	}

	// Value
	data.value = string(buf[:length])
	buf = bytes.TrimSpace(buf[length:])

	return data, buf, nil
}

func serializeField(name string, value string) string {
	return fmt.Sprintf("<%s:%d>%s", name, len(value), value)
}

// Print an ADIFRecord as a string
func (r *baseADIFRecord) ToString() string {
	var record bytes.Buffer
	for _, n := range ADIFfieldOrder {
		if v, ok := r.values[n]; ok {
			record.WriteString(serializeField(n, v))
		}
	}
	// Handle custom fields
	for n, v := range r.values {
		if !isStandardADIFField(n) {
			record.WriteString(serializeField(n, v))
		}
	}
	return record.String()
}

// Get fingerprint of ADIFRecord
func (r *baseADIFRecord) Fingerprint() string {
	fpfields := []string{
		"call", "station_callsign", "band",
		"freq", "mode", "qso_date", "time_on",
		"time_off"}
	fpvals := make([]string, 0, len(fpfields))
	for _, f := range fpfields {
		if n, ok := r.values[f]; ok {
			fpvals = append(fpvals, n)
		}
	}
	fptext := strings.Join(fpvals, "|")
	h := sha256.New()
	h.Write([]byte(fptext))
	return hex.EncodeToString(h.Sum(nil))
}

// Get a value
func (r *baseADIFRecord) GetValue(name string) (string, error) {
	if v, ok := r.values[name]; ok {
		return v, nil
	}
	return "", NoSuchField
}

// Set a value
func (r *baseADIFRecord) SetValue(name string, value string) {
	r.values[strings.ToLower(name)] = value
}

// Get all of the present field names
func (r *baseADIFRecord) GetFields() []string {
	keys := make([]string, len(r.values))
	i := 0
	for k := range r.values {
		keys[i] = k
		i++
	}
	return keys
}

// Delete a field (from the internal map)
func (r *baseADIFRecord) DeleteField(name string) (bool, error) {
	if _, ok := r.values[name]; ok {
		delete(r.values, name)
		return true, nil
	}
	return false, ErrNoSuchField
}
<|endoftext|>


<commit_before>package extdirect
import (
	"net/http"
	"fmt"
	"strings"
	"io"
	"io/ioutil"
	"encoding/json"
	"reflect"
	"github.com/mitchellh/mapstructure"
	"time"
	"golang.org/x/net/context"
	"net/url"
	"strconv"
	"github.com/nbgo/fail"
)

// ErrDecodeFromPostRequest has information about decoding error.
type ErrDecodeFromPostRequest struct {
	Details string
	Reason  error
}
var _ fail.CompositeError = ErrDecodeFromPostRequest{}
func (err ErrDecodeFromPostRequest) Error() string {
	return fmt.Sprintf("failed to decode form post: %v: %v", err.Details, err.Reason)
}
// InnerError implements CompositeError.InnerError().
func (err ErrDecodeFromPostRequest) InnerError() error {
	return err.Reason
}

// ErrInvalidContentType occurs when client request contains invalid content type.
type ErrInvalidContentType string
func (err ErrInvalidContentType) Error() string {
	return fmt.Sprintf("invalid content type: %s", string(err))
}

// ErrTypeConversion contains information about type conversion error.
type ErrTypeConversion struct {
	SourceType reflect.Type
	TargetType reflect.Type
}
func (err ErrTypeConversion) Error() string {
	return fmt.Sprintf("cannot convert type %v to type %v", err.SourceType, err.TargetType)
}

// ErrDirectActionMethod contains information about error occurred during direct method execution,
type ErrDirectActionMethod struct {
	Action  string
	Method  string
	Err     interface{}
	isPanic bool
}
func (err ErrDirectActionMethod) Error() string {
	return fmt.Sprintf("error executing %v.%v(): %v", err.Action, err.Method, err.Err)
}

type request struct {
	Type   string `json:"type"`
	Tid    int `json:"tid"`
	Action string `json:"action"`
	Method string `json:"method"`
	Upload bool `json:"upload"`
	Data   interface{} `json:"data"`
}

type response struct {
	Type    string `json:"type"`
	Tid     int `json:"tid"`
	Action  string `json:"action"`
	Method  string `json:"method"`
	Message *string `json:"message,omitempty"`
	Result  interface{} `json:"result,omitempty"`
}

// API is routes for getting Ext.Direct API script.
func API(provider *directServiceProvider) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/javascript; charset=utf-8")
		if js, err := provider.JavaScript(); err != nil {
			panic(err)
		} else {
			if _, err := w.Write([]byte(js)); err != nil {
				panic(err)
			}
		}
	}
}

// ActionsHandler is route for handling Ext.Direct requests.
func ActionsHandler(provider *directServiceProvider) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		actionHandler(provider, nil, w, r)
	}
}

// ActionsHandlerCtx is route with context suppor for handling Ext.Direct requests.
func ActionsHandlerCtx(provider *directServiceProvider) func(c context.Context, w http.ResponseWriter, r *http.Request) {
	return func(c context.Context, w http.ResponseWriter, r *http.Request) {
		actionHandler(provider, c, w, r)
	}
}

func actionHandler(provider *directServiceProvider, c context.Context, w http.ResponseWriter, r *http.Request) {
	var reqs []*request
	var err error
	contentType := r.Header.Get("Content-Type")
	isFormHandler := false

	switch {
	case strings.HasPrefix(contentType, "application/json"):
		reqs = mustDecodeTransaction(r.Body)
	case strings.HasPrefix(contentType, "application/x-www-form-urlencoded"):
		r.ParseForm()
		reqs = mustDecodeFormPost(r.Form)
		isFormHandler = true
	default:
		panic(ErrInvalidContentType(contentType))
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	if !isFormHandler {
		err = json.NewEncoder(w).Encode(provider.processRequests(c, r, reqs))
	} else {
		resps := provider.processRequests(c, r, reqs)
		err = json.NewEncoder(w).Encode(resps[0])

	}
	if err != nil {
		panic(err)
	}
}

func (provider *directServiceProvider) processRequests(c context.Context, r *http.Request, reqs []*request) []*response {
	resps := make([]*response, len(reqs))
	respsChannel := make(chan *response, len(reqs))
	for _, req := range reqs {
		go func(req *request) {
			resp := &response{
				Tid: req.Tid,
				Action: req.Action,
				Method: req.Method,
				Type: req.Type,
			}
			var tStart time.Time
			profilingStarted := false
			defer func() {
				if profilingStarted {
					duration := time.Now().Sub(tStart)
					log.Print(logLevelInfo, fmt.Sprintf("%s.%s() %v ", req.Action, req.Method, duration), map[string]interface{}{"duration":duration, "action": req.Action, "method": req.Method})
					profilingStarted = false
				}
				if err := recover(); err != nil {
					log.Print(fail.New(ErrDirectActionMethod{req.Action, req.Method, err, true}))
					resp.Type = "exception"
					respMessage := fmt.Sprintf("%v", err)
					resp.Message = &respMessage
				}
				respsChannel <- resp
			}()

			// Create instance of action type
			actionInfo := provider.actionsInfo[req.Action]
			if provider.debug {
				log.Print(fmt.Sprintf("Create instance of action %s (type %v)", req.Action, actionInfo.Type))
			}
			actionVal := reflect.New(actionInfo.Type).Elem()

			// Set context and request
			if c != nil || r != nil {
				if provider.debug {
					log.Print("Set action context/request.")
				}
				contextType := reflect.TypeOf((*context.Context)(nil)).Elem()
				requestType := reflect.TypeOf(&http.Request{})
				fieldsLen := actionInfo.Type.NumField()
				for i := 0; i < fieldsLen; i++ {
					t := actionInfo.Type.Field(i).Type

					if t.Implements(contextType) {
						if c != nil {
							if provider.debug {
								log.Print("Set action context.")
							}
							actionVal.Field(i).Set(reflect.ValueOf(c))
						} else {
							if provider.debug {
								log.Print(logLevelWarn, "Context cannot be set to action instance because context is nil.")
							}
						}
					}

					if t == requestType {
						if r != nil {
							if provider.debug {
								log.Print("Set action request.")
							}
							actionVal.Field(i).Set(reflect.ValueOf(r))
						}
					}
				}
			}

			if provider.debug {
				log.Print(fmt.Sprintf("Prepare arguments for method %s.%s", req.Action, req.Method))
			}
			methodInfo := actionInfo.Methods[req.Method]
			directMethod := actionInfo.DirectMethods[req.Method]
			if provider.debug {
				log.Print(fmt.Sprintf("Direct method to use: %s, formhandler=%v", directMethod.Name, directMethod.FormHandler))
			}
			methodArgsLen := methodInfo.Type.NumIn() - 1
			var args []reflect.Value
			if req.Data != nil {
				if provider.debug {
					log.Print(fmt.Sprintf("Type of request data is %T", req.Data))
				}
				if directMethod.FormHandler != nil {
					if provider.debug {
						log.Print("Prepare arguments for form handler call.")
					}
					args = make([]reflect.Value, 1)
					args[0] = reflect.ValueOf(req.Data.(map[string]string))
					// TODO: Support structure type argument for form handler.
				} else {
					args = make([]reflect.Value, methodArgsLen)
					for i, arg := range req.Data.([]interface{}) {
						if provider.debug {
							log.Print(fmt.Sprintf("Initial arg #%v type is %T", i, arg))
						}
						convertedArg := convertArg(methodInfo.Type.In(i + 1), arg)
						if provider.debug {
							log.Print(fmt.Sprintf("Converted arg #%v type is %T", i, convertedArg))
						}
						args[i] = reflect.ValueOf(convertedArg)
					}
				}
			}

			if provider.profile {
				profilingStarted = true
				tStart = time.Now()
			}

			if provider.debug {
				log.Print(fmt.Sprintf("Call method %s.%s", req.Action, req.Method))
			}
			// Call action method.
			resultsValues := actionVal.MethodByName(methodInfo.Name).Call(args)

			if profilingStarted {
				duration := time.Now().Sub(tStart)
				log.Print(logLevelInfo, fmt.Sprintf("%s.%s() %v ", req.Action, req.Method, duration), map[string]interface{}{"duration":duration, "action": req.Action, "method": req.Method})
				profilingStarted = false
			}
			for i, resultValue := range resultsValues {
				if methodInfo.Type.Out(i).Name() == "error" {
					if err, isErr := resultValue.Interface().(error); isErr {
						log.Print(&ErrDirectActionMethod{req.Action, req.Method, err, false})
						resp.Type = "exception"
						respMessage := fmt.Sprintf("%v", err)
						resp.Message = &respMessage
						resp.Result = nil
						break;
					}
				} else {
					result := resultValue.Interface()
					resp.Result = result
				}
			}
		}(req)
	}

	for i := 0; i < len(reqs); i++ {
		var resp = <-respsChannel
		resps[i] = resp
	}

	return resps
}

func mustDecodeFormPost(f url.Values) []*request {
	req := &request{
		Type:f["extType"][0],
		Action: f["extAction"][0],
		Method: f["extMethod"][0],
	}
	tid, tidErr := strconv.Atoi(f["extTID"][0]);
	if tidErr != nil {
		panic(fail.New(ErrDecodeFromPostRequest{"could not parse TID", tidErr}))
	}
	req.Tid = tid
	upload, hasUpload := f["extUpload"]
	req.Upload = hasUpload && strings.ToLower(upload[0]) == "true"

	data := make(map[string]string, 0)
	for k, v := range f {
		if k == "extType" || k == "extTID" || k == "extAction" || k == "extMethod" || k == "extUpload" {
			continue
		}
		data[k] = v[0]
	}
	req.Data = data

	return []*request{req}
}

func mustDecodeTransaction(r io.Reader) []*request {
	if jsonData, err := ioutil.ReadAll(r); err != nil {
		panic(err)
	} else {
		var reqs []*request
		if err := json.Unmarshal(jsonData, &reqs); err != nil {
			// Attempt to unmarshal as a single request.
			var req request
			if err := json.Unmarshal(jsonData, &req); err != nil {
				panic(err)
			} else {
				reqs = make([]*request, 1)
				reqs[0] = &req
			}
		}
		return reqs
	}
}

func convertArg(argType reflect.Type, argValue interface{}) interface{} {
	sourceType := reflect.TypeOf(argValue)
	if sourceType != argType {
		switch v := argValue.(type) {
		case float64:
			switch argType.Kind() {
			case reflect.Int: return int(v)
			case reflect.Int8: return int8(v)
			case reflect.Int16: return int16(v)
			case reflect.Int32: return int32(v)
			case reflect.Float32: return float32(v)
			default: panic(fail.New(ErrTypeConversion{sourceType, argType}))
			}
		case nil:
			switch argType.Kind() {
			case reflect.Int: return int(0)
			case reflect.Int8: return int8(0)
			case reflect.Int16: return int16(0)
			case reflect.Int32: return int32(0)
			case reflect.String: return ""
			default: panic(fail.New(ErrTypeConversion{sourceType, argType}))
			}
		case map[string]interface{}:
			switch argType.Kind() {
			case reflect.Ptr: fallthrough
			case reflect.Struct:
				structInstanceValue := reflect.New(argType).Elem()
				structInstanceRef := structInstanceValue.Addr().Interface()
				mapstructure.Decode(v, structInstanceRef)
				return structInstanceValue.Interface()
			default: panic(fail.New(ErrTypeConversion{sourceType, argType}))
			}
		}
	}

	return argValue
}<commit_msg>IsNil info added for converted argument.<commit_after>package extdirect

import (
	"net/http"
	"fmt"
	"strings"
	"io"
	"io/ioutil"
	"encoding/json"
	"reflect"
	"github.com/mitchellh/mapstructure"
	"time"
	"golang.org/x/net/context"
	"net/url"
	"strconv"
	"github.com/nbgo/fail"
)

// ErrDecodeFromPostRequest has information about decoding error.
type ErrDecodeFromPostRequest struct {
	Details string
	Reason  error
}

var _ fail.CompositeError = ErrDecodeFromPostRequest{}

func (err ErrDecodeFromPostRequest) Error() string {
	return fmt.Sprintf("failed to decode form post: %v: %v", err.Details, err.Reason)
}
// InnerError implements CompositeError.InnerError().
func (err ErrDecodeFromPostRequest) InnerError() error {
	return err.Reason
}

// ErrInvalidContentType occurs when client request contains invalid content type.
type ErrInvalidContentType string

func (err ErrInvalidContentType) Error() string {
	return fmt.Sprintf("invalid content type: %s", string(err))
}

// ErrTypeConversion contains information about type conversion error.
type ErrTypeConversion struct {
	SourceType reflect.Type
	TargetType reflect.Type
}

func (err ErrTypeConversion) Error() string {
	return fmt.Sprintf("cannot convert type %v to type %v", err.SourceType, err.TargetType)
}

// ErrDirectActionMethod contains information about error occurred during direct method execution,
type ErrDirectActionMethod struct {
	Action  string
	Method  string
	Err     interface{}
	isPanic bool
}

func (err ErrDirectActionMethod) Error() string {
	return fmt.Sprintf("error executing %v.%v(): %v", err.Action, err.Method, err.Err)
}

type request struct {
	Type   string `json:"type"`
	Tid    int `json:"tid"`
	Action string `json:"action"`
	Method string `json:"method"`
	Upload bool `json:"upload"`
	Data   interface{} `json:"data"`
}

type response struct {
	Type    string `json:"type"`
	Tid     int `json:"tid"`
	Action  string `json:"action"`
	Method  string `json:"method"`
	Message *string `json:"message,omitempty"`
	Result  interface{} `json:"result,omitempty"`
}

// API is routes for getting Ext.Direct API script.
func API(provider *directServiceProvider) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/javascript; charset=utf-8")
		if js, err := provider.JavaScript(); err != nil {
			panic(err)
		} else {
			if _, err := w.Write([]byte(js)); err != nil {
				panic(err)
			}
		}
	}
}

// ActionsHandler is route for handling Ext.Direct requests.
func ActionsHandler(provider *directServiceProvider) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		actionHandler(provider, nil, w, r)
	}
}

// ActionsHandlerCtx is route with context suppor for handling Ext.Direct requests.
func ActionsHandlerCtx(provider *directServiceProvider) func(c context.Context, w http.ResponseWriter, r *http.Request) {
	return func(c context.Context, w http.ResponseWriter, r *http.Request) {
		actionHandler(provider, c, w, r)
	}
}

func actionHandler(provider *directServiceProvider, c context.Context, w http.ResponseWriter, r *http.Request) {
	var reqs []*request
	var err error
	contentType := r.Header.Get("Content-Type")
	isFormHandler := false

	switch {
	case strings.HasPrefix(contentType, "application/json"):
		reqs = mustDecodeTransaction(r.Body)
	case strings.HasPrefix(contentType, "application/x-www-form-urlencoded"):
		r.ParseForm()
		reqs = mustDecodeFormPost(r.Form)
		isFormHandler = true
	default:
		panic(ErrInvalidContentType(contentType))
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	if !isFormHandler {
		err = json.NewEncoder(w).Encode(provider.processRequests(c, r, reqs))
	} else {
		resps := provider.processRequests(c, r, reqs)
		err = json.NewEncoder(w).Encode(resps[0])

	}
	if err != nil {
		panic(err)
	}
}

func (provider *directServiceProvider) processRequests(c context.Context, r *http.Request, reqs []*request) []*response {
	resps := make([]*response, len(reqs))
	respsChannel := make(chan *response, len(reqs))
	for _, req := range reqs {
		go func(req *request) {
			resp := &response{
				Tid: req.Tid,
				Action: req.Action,
				Method: req.Method,
				Type: req.Type,
			}
			var tStart time.Time
			profilingStarted := false
			defer func() {
				if profilingStarted {
					duration := time.Now().Sub(tStart)
					log.Print(logLevelInfo, fmt.Sprintf("%s.%s() %v ", req.Action, req.Method, duration), map[string]interface{}{"duration":duration, "action": req.Action, "method": req.Method})
					profilingStarted = false
				}
				if err := recover(); err != nil {
					log.Print(fail.New(ErrDirectActionMethod{req.Action, req.Method, err, true}))
					resp.Type = "exception"
					respMessage := fmt.Sprintf("%v", err)
					resp.Message = &respMessage
				}
				respsChannel <- resp
			}()

			// Create instance of action type
			actionInfo := provider.actionsInfo[req.Action]
			if provider.debug {
				log.Print(fmt.Sprintf("Create instance of action %s (type %v)", req.Action, actionInfo.Type))
			}
			actionVal := reflect.New(actionInfo.Type).Elem()

			// Set context and request
			if c != nil || r != nil {
				if provider.debug {
					log.Print("Set action context/request.")
				}
				contextType := reflect.TypeOf((*context.Context)(nil)).Elem()
				requestType := reflect.TypeOf(&http.Request{})
				fieldsLen := actionInfo.Type.NumField()
				for i := 0; i < fieldsLen; i++ {
					t := actionInfo.Type.Field(i).Type

					if t.Implements(contextType) {
						if c != nil {
							if provider.debug {
								log.Print("Set action context.")
							}
							actionVal.Field(i).Set(reflect.ValueOf(c))
						} else {
							if provider.debug {
								log.Print(logLevelWarn, "Context cannot be set to action instance because context is nil.")
							}
						}
					}

					if t == requestType {
						if r != nil {
							if provider.debug {
								log.Print("Set action request.")
							}
							actionVal.Field(i).Set(reflect.ValueOf(r))
						}
					}
				}
			}

			if provider.debug {
				log.Print(fmt.Sprintf("Prepare arguments for method %s.%s", req.Action, req.Method))
			}
			methodInfo := actionInfo.Methods[req.Method]
			directMethod := actionInfo.DirectMethods[req.Method]
			if provider.debug {
				log.Print(fmt.Sprintf("Direct method to use: %s, formhandler=%v", directMethod.Name, directMethod.FormHandler))
			}
			methodArgsLen := methodInfo.Type.NumIn() - 1
			var args []reflect.Value
			if req.Data != nil {
				if provider.debug {
					log.Print(fmt.Sprintf("Type of request data is %T", req.Data))
				}
				if directMethod.FormHandler != nil {
					if provider.debug {
						log.Print("Prepare arguments for form handler call.")
					}
					args = make([]reflect.Value, 1)
					args[0] = reflect.ValueOf(req.Data.(map[string]string))
					// TODO: Support structure type argument for form handler.
				} else {
					args = make([]reflect.Value, methodArgsLen)
					for i, arg := range req.Data.([]interface{}) {
						if provider.debug {
							log.Print(fmt.Sprintf("Initial arg #%v type is %T", i, arg))
						}
						convertedArg := convertArg(methodInfo.Type.In(i + 1), arg)
						if provider.debug {
							log.Print(fmt.Sprintf("Converted arg #%v type is %T, IsNil=%v", i, convertedArg, convertedArg == nil))
						}
						args[i] = reflect.ValueOf(convertedArg)
					}
				}
			}

			if provider.profile {
				profilingStarted = true
				tStart = time.Now()
			}

			if provider.debug {
				log.Print(fmt.Sprintf("Call method %s.%s", req.Action, req.Method))
			}
			// Call action method.
			resultsValues := actionVal.MethodByName(methodInfo.Name).Call(args)

			if profilingStarted {
				duration := time.Now().Sub(tStart)
				log.Print(logLevelInfo, fmt.Sprintf("%s.%s() %v ", req.Action, req.Method, duration), map[string]interface{}{"duration":duration, "action": req.Action, "method": req.Method})
				profilingStarted = false
			}
			for i, resultValue := range resultsValues {
				if methodInfo.Type.Out(i).Name() == "error" {
					if err, isErr := resultValue.Interface().(error); isErr {
						log.Print(&ErrDirectActionMethod{req.Action, req.Method, err, false})
						resp.Type = "exception"
						respMessage := fmt.Sprintf("%v", err)
						resp.Message = &respMessage
						resp.Result = nil
						break;
					}
				} else {
					result := resultValue.Interface()
					resp.Result = result
				}
			}
		}(req)
	}

	for i := 0; i < len(reqs); i++ {
		var resp = <-respsChannel
		resps[i] = resp
	}

	return resps
}

func mustDecodeFormPost(f url.Values) []*request {
	req := &request{
		Type:f["extType"][0],
		Action: f["extAction"][0],
		Method: f["extMethod"][0],
	}
	tid, tidErr := strconv.Atoi(f["extTID"][0]);
	if tidErr != nil {
		panic(fail.New(ErrDecodeFromPostRequest{"could not parse TID", tidErr}))
	}
	req.Tid = tid
	upload, hasUpload := f["extUpload"]
	req.Upload = hasUpload && strings.ToLower(upload[0]) == "true"

	data := make(map[string]string, 0)
	for k, v := range f {
		if k == "extType" || k == "extTID" || k == "extAction" || k == "extMethod" || k == "extUpload" {
			continue
		}
		data[k] = v[0]
	}
	req.Data = data

	return []*request{req}
}

func mustDecodeTransaction(r io.Reader) []*request {
	if jsonData, err := ioutil.ReadAll(r); err != nil {
		panic(err)
	} else {
		var reqs []*request
		if err := json.Unmarshal(jsonData, &reqs); err != nil {
			// Attempt to unmarshal as a single request.
			var req request
			if err := json.Unmarshal(jsonData, &req); err != nil {
				panic(err)
			} else {
				reqs = make([]*request, 1)
				reqs[0] = &req
			}
		}
		return reqs
	}
}

func convertArg(argType reflect.Type, argValue interface{}) interface{} {
	sourceType := reflect.TypeOf(argValue)
	if sourceType != argType {
		switch v := argValue.(type) {
		case float64:
			switch argType.Kind() {
			case reflect.Int: return int(v)
			case reflect.Int8: return int8(v)
			case reflect.Int16: return int16(v)
			case reflect.Int32: return int32(v)
			case reflect.Float32: return float32(v)
			default: panic(fail.New(ErrTypeConversion{sourceType, argType}))
			}
		case nil:
			switch argType.Kind() {
			case reflect.Int: return int(0)
			case reflect.Int8: return int8(0)
			case reflect.Int16: return int16(0)
			case reflect.Int32: return int32(0)
			case reflect.String: return ""
			default: panic(fail.New(ErrTypeConversion{sourceType, argType}))
			}
		case map[string]interface{}:
			switch argType.Kind() {
			case reflect.Ptr: fallthrough
			case reflect.Struct:
				structInstanceValue := reflect.New(argType).Elem()
				structInstanceRef := structInstanceValue.Addr().Interface()
				mapstructure.Decode(v, structInstanceRef)
				return structInstanceValue.Interface()
			default: panic(fail.New(ErrTypeConversion{sourceType, argType}))
			}
		}
	}

	return argValue
}<|endoftext|>


<commit_before>package ai

import (
	"bytes"
	"log"

	"nelhage.com/tak/ptn"
	"nelhage.com/tak/tak"
)

const (
	maxEval int64 = 1 << 30
	minEval       = -maxEval
)

type MinimaxAI struct {
	depth int

	Debug bool
}

func formatpv(ms []tak.Move) string {
	var out bytes.Buffer
	out.WriteString("[")
	for i, m := range ms {
		if i != 0 {
			out.WriteString(" ")
		}
		out.WriteString(ptn.FormatMove(&m))
	}
	out.WriteString("]")
	return out.String()
}

func (m *MinimaxAI) GetMove(p *tak.Position) tak.Move {
	ms, _ := m.Analyze(p)
	return ms[0]
}

func (m *MinimaxAI) Analyze(p *tak.Position) ([]tak.Move, int64) {
	var ms []tak.Move
	var v int64
	for i := 1; i <= m.depth; i++ {
		ms, v = m.minimax(p, i, ms, minEval-1, maxEval+1)
		if m.Debug {
			log.Printf("[minimax] depth=%d val=%d pv=%s",
				i, v, formatpv(ms))
		}
	}
	return ms, v
}

func (ai *MinimaxAI) minimax(
	p *tak.Position,
	depth int,
	pv []tak.Move,
	α, β int64) ([]tak.Move, int64) {
	over, _ := p.GameOver()
	if depth == 0 || over {
		return nil, ai.evaluate(p)
	}

	if p.MoveNumber() < 2 {
		for _, c := range [][]int{{0, 0}, {p.Size() - 1, 0}, {0, p.Size() - 1}, {p.Size() - 1, p.Size() - 1}} {
			x, y := c[0], c[1]
			if len(p.At(x, y)) == 0 {
				return []tak.Move{{X: x, Y: y, Type: tak.PlaceFlat}}, 0
			}
		}
	}
	moves := p.AllMoves()
	if len(pv) > 0 {
		for i, m := range moves {
			if m.Equal(&pv[0]) {
				moves[0], moves[i] = moves[i], moves[0]
				break
			}
		}
	}

	best := make([]tak.Move, 0, depth)
	max := minEval - 1
	for _, m := range moves {
		child, e := p.Move(&m)
		if e != nil {
			continue
		}
		var ms []tak.Move
		var v int64
		if len(best) == 0 {
			ms, v = ai.minimax(child, depth-1, nil, -β, -α)
		} else {
			ms, v = ai.minimax(child, depth-1, best[1:], -β, -α)
		}
		v = -v
		if v > max {
			max = v
			best = append(best[:0], m)
			best = append(best, ms...)
		}
		if v > α {
			α = v
			if α >= β {
				break
			}
		}
	}
	return best, max
}

func imin(a, b int) int {
	if a < b {
		return a
	}
	return b
}

const (
	weightFlat       = 1
	weightCaptured   = 1
	weightControlled = 5
)

func (m *MinimaxAI) evaluate(p *tak.Position) int64 {
	if over, winner := p.GameOver(); over {
		switch winner {
		case tak.NoColor:
			return 0
		case p.ToMove():
			return maxEval - int64(p.MoveNumber())
		default:
			return minEval + int64(p.MoveNumber())
		}
	}
	mine, theirs := 0, 0
	me := p.ToMove()
	addw := func(c tak.Color, w int) {
		if c == me {
			mine += w
		} else {
			theirs += w
		}
	}
	for x := 0; x < p.Size(); x++ {
		for y := 0; y < p.Size(); y++ {
			sq := p.At(x, y)
			if len(sq) == 0 {
				continue
			}
			addw(sq[0].Color(), weightControlled)
			for i, stone := range sq {
				if i > 0 && i < p.Size() {
					addw(sq[0].Color(), weightCaptured)
				}
				if stone.Kind() == tak.Flat {
					addw(stone.Color(), weightFlat)
				}
			}
		}
	}

	return int64(mine - theirs)
}

func NewMinimax(depth int) *MinimaxAI {
	return &MinimaxAI{depth: depth}
}
<commit_msg>minor value function tweaking<commit_after>package ai

import (
	"bytes"
	"log"

	"nelhage.com/tak/ptn"
	"nelhage.com/tak/tak"
)

const (
	maxEval int64 = 1 << 30
	minEval       = -maxEval
)

type MinimaxAI struct {
	depth int

	Debug bool
}

func formatpv(ms []tak.Move) string {
	var out bytes.Buffer
	out.WriteString("[")
	for i, m := range ms {
		if i != 0 {
			out.WriteString(" ")
		}
		out.WriteString(ptn.FormatMove(&m))
	}
	out.WriteString("]")
	return out.String()
}

func (m *MinimaxAI) GetMove(p *tak.Position) tak.Move {
	ms, _ := m.Analyze(p)
	return ms[0]
}

func (m *MinimaxAI) Analyze(p *tak.Position) ([]tak.Move, int64) {
	var ms []tak.Move
	var v int64
	for i := 1; i <= m.depth; i++ {
		ms, v = m.minimax(p, i, ms, minEval-1, maxEval+1)
		if m.Debug {
			log.Printf("[minimax] depth=%d val=%d pv=%s",
				i, v, formatpv(ms))
		}
	}
	return ms, v
}

func (ai *MinimaxAI) minimax(
	p *tak.Position,
	depth int,
	pv []tak.Move,
	α, β int64) ([]tak.Move, int64) {
	over, _ := p.GameOver()
	if depth == 0 || over {
		return nil, ai.evaluate(p)
	}

	if p.MoveNumber() < 2 {
		for _, c := range [][]int{{0, 0}, {p.Size() - 1, 0}, {0, p.Size() - 1}, {p.Size() - 1, p.Size() - 1}} {
			x, y := c[0], c[1]
			if len(p.At(x, y)) == 0 {
				return []tak.Move{{X: x, Y: y, Type: tak.PlaceFlat}}, 0
			}
		}
	}
	moves := p.AllMoves()
	if len(pv) > 0 {
		for i, m := range moves {
			if m.Equal(&pv[0]) {
				moves[0], moves[i] = moves[i], moves[0]
				break
			}
		}
	}

	best := make([]tak.Move, 0, depth)
	max := minEval - 1
	for _, m := range moves {
		child, e := p.Move(&m)
		if e != nil {
			continue
		}
		var ms []tak.Move
		var v int64
		if len(best) == 0 {
			ms, v = ai.minimax(child, depth-1, nil, -β, -α)
		} else {
			ms, v = ai.minimax(child, depth-1, best[1:], -β, -α)
		}
		v = -v
		if v > max {
			max = v
			best = append(best[:0], m)
			best = append(best, ms...)
		}
		if v > α {
			α = v
			if α >= β {
				break
			}
		}
	}
	return best, max
}

func imin(a, b int) int {
	if a < b {
		return a
	}
	return b
}

const (
	weightFlat       = 2
	weightCaptured   = 1
	weightControlled = 5
	weightCapstone   = -1
)

func (m *MinimaxAI) evaluate(p *tak.Position) int64 {
	if over, winner := p.GameOver(); over {
		switch winner {
		case tak.NoColor:
			return 0
		case p.ToMove():
			return maxEval - int64(p.MoveNumber())
		default:
			return minEval + int64(p.MoveNumber())
		}
	}
	mine, theirs := 0, 0
	me := p.ToMove()
	addw := func(c tak.Color, w int) {
		if c == me {
			mine += w
		} else {
			theirs += w
		}
	}
	for x := 0; x < p.Size(); x++ {
		for y := 0; y < p.Size(); y++ {
			sq := p.At(x, y)
			if len(sq) == 0 {
				continue
			}
			addw(sq[0].Color(), weightControlled)
			if sq[0].Kind() == tak.Capstone {
				addw(sq[0].Color(), weightCapstone)
			}
			for i, stone := range sq {
				if i > 0 && i < p.Size() {
					addw(sq[0].Color(), weightCaptured)
				}
				if stone.Kind() == tak.Flat {
					addw(stone.Color(), weightFlat)
				}
			}
		}
	}

	return int64(mine - theirs)
}

func NewMinimax(depth int) *MinimaxAI {
	return &MinimaxAI{depth: depth}
}
<|endoftext|>


<commit_before>package jsonrpc

import (
	"errors"

	"github.com/neptulon/cmap"
	"github.com/neptulon/shortid"
)

// Router is a JSON-RPC message routing middleware.
type Router struct {
	server    *Server
	reqRoutes map[string]func(ctx *ReqCtx) error // method name -> handler func(ctx *ReqCtx) error
	notRoutes map[string]func(ctx *NotCtx) error // method name -> handler func(ctx *NotCtx) error
	resRoutes *cmap.CMap                         // message ID (string) -> handler func(ctx *ResCtx) error : requests sent from the router that are pending responses from clients
}

// NewRouter creates a JSON-RPC router instance and registers it with the Neptulon JSON-RPC server.
func NewRouter(s *Server) (*Router, error) {
	if s == nil {
		return nil, errors.New("Given Neptulon server instance is nil.")
	}

	r := Router{
		server:    s,
		reqRoutes: make(map[string]func(ctx *ReqCtx) error),
		notRoutes: make(map[string]func(ctx *NotCtx) error),
		resRoutes: cmap.New(),
	}

	s.ReqMiddleware(r.reqMiddleware)
	s.NotMiddleware(r.notMiddleware)
	s.ResMiddleware(r.resMiddleware)
	return &r, nil
}

// Request adds a new incoming request route registry.
func (r *Router) Request(route string, handler func(ctx *ReqCtx) error) {
	r.reqRoutes[route] = handler
}

// Notification adds a new incoming notification route registry.
func (r *Router) Notification(route string, handler func(ctx *NotCtx) error) {
	r.notRoutes[route] = handler
}

// SendRequest sends a JSON-RPC request throught the connection denoted by the connection ID.
// resHandler is called when a response is returned.
func (r *Router) SendRequest(connID string, method string, params interface{}, resHandler func(ctx *ResCtx)) error {
	id, err := shortid.UUID()
	if err != nil {
		return err
	}

	req := Request{ID: id, Method: method, Params: params}
	if err = r.server.send(connID, req); err != nil {
		return err
	}

	r.resRoutes.Set(req.ID, resHandler)
	return nil
}

// SendNotification sends a JSON-RPC notification through the connection denoted by the connection ID.
func (r *Router) SendNotification(connID string, method string, params interface{}) error {
	return r.server.send(connID, Notification{Method: method, Params: params})
}

func (r *Router) reqMiddleware(ctx *ReqCtx) error {
	if handler, ok := r.reqRoutes[ctx.method]; ok {
		return handler(ctx)
	}

	return nil
}

func (r *Router) notMiddleware(ctx *NotCtx) error {
	if handler, ok := r.notRoutes[ctx.method]; ok {
		handler(ctx)
	}

	return nil
}

func (r *Router) resMiddleware(ctx *ResCtx) error {
	if handler, ok := r.resRoutes.GetOk(ctx.id); ok {
		handler.(func(ctx *ResCtx) error)(ctx)
		r.resRoutes.Delete(ctx.id)
	}

	return nil
}
<commit_msg>fix router return erros<commit_after>package jsonrpc

import (
	"errors"

	"github.com/neptulon/cmap"
	"github.com/neptulon/shortid"
)

// Router is a JSON-RPC message routing middleware.
type Router struct {
	server    *Server
	reqRoutes map[string]func(ctx *ReqCtx) error // method name -> handler func(ctx *ReqCtx) error
	notRoutes map[string]func(ctx *NotCtx) error // method name -> handler func(ctx *NotCtx) error
	resRoutes *cmap.CMap                         // message ID (string) -> handler func(ctx *ResCtx) error : requests sent from the router that are pending responses from clients
}

// NewRouter creates a JSON-RPC router instance and registers it with the Neptulon JSON-RPC server.
func NewRouter(s *Server) (*Router, error) {
	if s == nil {
		return nil, errors.New("Given Neptulon server instance is nil.")
	}

	r := Router{
		server:    s,
		reqRoutes: make(map[string]func(ctx *ReqCtx) error),
		notRoutes: make(map[string]func(ctx *NotCtx) error),
		resRoutes: cmap.New(),
	}

	s.ReqMiddleware(r.reqMiddleware)
	s.NotMiddleware(r.notMiddleware)
	s.ResMiddleware(r.resMiddleware)
	return &r, nil
}

// Request adds a new incoming request route registry.
func (r *Router) Request(route string, handler func(ctx *ReqCtx) error) {
	r.reqRoutes[route] = handler
}

// Notification adds a new incoming notification route registry.
func (r *Router) Notification(route string, handler func(ctx *NotCtx) error) {
	r.notRoutes[route] = handler
}

// SendRequest sends a JSON-RPC request throught the connection denoted by the connection ID.
// resHandler is called when a response is returned.
func (r *Router) SendRequest(connID string, method string, params interface{}, resHandler func(ctx *ResCtx)) error {
	id, err := shortid.UUID()
	if err != nil {
		return err
	}

	req := Request{ID: id, Method: method, Params: params}
	if err = r.server.send(connID, req); err != nil {
		return err
	}

	r.resRoutes.Set(req.ID, resHandler)
	return nil
}

// SendNotification sends a JSON-RPC notification through the connection denoted by the connection ID.
func (r *Router) SendNotification(connID string, method string, params interface{}) error {
	return r.server.send(connID, Notification{Method: method, Params: params})
}

func (r *Router) reqMiddleware(ctx *ReqCtx) error {
	if handler, ok := r.reqRoutes[ctx.method]; ok {
		return handler(ctx)
	}

	return nil
}

func (r *Router) notMiddleware(ctx *NotCtx) error {
	if handler, ok := r.notRoutes[ctx.method]; ok {
		return handler(ctx)
	}

	return nil
}

func (r *Router) resMiddleware(ctx *ResCtx) error {
	if handler, ok := r.resRoutes.GetOk(ctx.id); ok {
		err := handler.(func(ctx *ResCtx) error)(ctx)
		r.resRoutes.Delete(ctx.id)
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package baa

import (
	"fmt"
	"sync"
)

const (
	// method key in routeMap
	GET int = iota
	POST
	PUT
	DELETE
	PATCH
	OPTIONS
	HEAD
	// RouteLength route table length
	RouteLength
)

// methods declare support HTTP method
var methods = map[string]bool{
	"GET":     true,
	"POST":    true,
	"PUT":     true,
	"DELETE":  true,
	"PATCH":   true,
	"OPTIONS": true,
	"HEAD":    true,
}

// methodKeys declare method key in routeMap
var methodKeys = map[string]int{
	"GET":     GET,
	"POST":    POST,
	"PUT":     PUT,
	"DELETE":  DELETE,
	"PATCH":   PATCH,
	"OPTIONS": OPTIONS,
	"HEAD":    HEAD,
}

// Router provlider router for baa
type Router struct {
	autoHead        bool
	mu              sync.RWMutex
	notFoundHandler HandlerFunc
	routeMap        [RouteLength]*Route
	routeNamedMap   map[string]string
	groups          []*group
}

// Route is a tree node
// route use radix tree
type Route struct {
	pattern  string
	hasParam bool
	parent   *Route
	router   *Router
	children []*Route
	handlers []HandlerFunc
}

// group route
type group struct {
	pattern  string
	handlers []HandlerFunc
}

// newRouter create a router instance
func newRouter() *Router {
	r := new(Router)
	for i := 0; i < len(r.routeMap); i++ {
		r.routeMap[i] = newRoute("/", nil, nil)
	}
	r.routeNamedMap = make(map[string]string)
	r.groups = make([]*group, 0)
	return r
}

// newRoute create a route item
func newRoute(pattern string, handles []HandlerFunc, router *Router) *Route {
	r := new(Route)
	r.pattern = pattern
	r.handlers = handles
	r.router = router
	r.children = make([]*Route, 0)
	return r
}

// newGroup create a group router
func newGroup() *group {
	g := new(group)
	g.handlers = make([]HandlerFunc, 0)
	return g
}

// Match match the uri for handler
func (r *Router) match(method, uri string, c *Context) *Route {
	return r.lookup(uri, r.routeMap[methodKeys[method]], c)
}

// urlFor use named route return format url
func (r *Router) urlFor(name string, args ...interface{}) string {
	if name == "" {
		return ""
	}
	url := r.routeNamedMap[name]
	if url == "" {
		return ""
	}
	return fmt.Sprintf(url, args...)
}

// groupAdd add a group route has same prefix and handle chain
func (r *Router) groupAdd(pattern string, f func(), handlers []HandlerFunc) {
	g := newGroup()
	g.pattern = pattern
	g.handlers = handlers
	r.groups = append(r.groups, g)

	f()

	r.groups = r.groups[:len(r.groups)-1]
}

// Handle registers a new request handle with the given pattern, method and handlers.
func (r *Router) add(method string, pattern string, handlers []HandlerFunc) *Route {
	if _, ok := methods[method]; !ok {
		panic("unsupport http method [" + method + "]")
	}
	if pattern == "" {
		panic("route pattern can not be emtpy!")
	}
	if pattern[0] != '/' {
		panic("route pattern must begin /")
	}
	r.mu.Lock()
	defer r.mu.Unlock()

	// check group set
	if len(r.groups) > 0 {
		var gpattern string
		var ghandlers []HandlerFunc
		for i := range r.groups {
			gpattern += r.groups[i].pattern
			if len(r.groups[i].handlers) > 0 {
				ghandlers = append(ghandlers, r.groups[i].handlers...)
			}
		}
		pattern = gpattern + pattern
		ghandlers = append(ghandlers, handlers...)
		handlers = ghandlers
	}

	for i := 0; i < len(handlers); i++ {
		handlers[i] = wrapHandlerFunc(handlers[i])
	}

	root := r.routeMap[methodKeys[method]]
	var radix []byte
	var param []byte
	var i, k int
	var tru *Route
	for i = 0; i < len(pattern); i++ {
		//param route
		if pattern[i] == ':' {
			// clear static route
			if len(radix) > 0 {
				root = r.insert(root, newRoute(string(radix), nil, nil))
				radix = radix[:0]
			}
			// set param route
			param = param[:0]
			k = 0
			for i = i + 1; i < len(pattern); i++ {
				if !isParamChar(pattern[i]) {
					i--
					break
				}
				param = append(param, pattern[i])
				k++
			}
			if k == 0 {
				panic("route pattern param is empty")
			}
			// check last character
			p := ":" + string(param[:k])
			if i == len(pattern) {
				tru = newRoute(p, handlers, r)
			} else {
				tru = newRoute(p, nil, nil)
			}
			tru.hasParam = true
			root = r.insert(root, tru)
			continue
		}
		radix = append(radix, pattern[i])
	}

	// static route
	if len(radix) > 0 {
		tru = newRoute(string(radix), handlers, r)
		r.insert(root, tru)
	}

	return newRoute(pattern, handlers, r)
}

// insert build the route tree
func (r *Router) insert(root *Route, node *Route) *Route {
	// same route
	if root.pattern == node.pattern {
		if node.handlers != nil {
			root.handlers = node.handlers
		}
		return root
	}

	// param route
	if root.hasParam && node.hasParam {
		if root.parent == nil {
			panic("Router.insert error route has no parent")
		}
		return root.parent.insertChild(node)
	}

	if !root.hasParam && node.hasParam {
		return root.insertChild(node)
	}

	var i int
	if root.hasParam && !node.hasParam {
		for i = range root.children {
			if root.children[i].pattern == node.pattern {
				if node.handlers != nil {
					root.children[i].handlers = node.handlers
				}
				return root.children[i]
			}
			if root.children[i].hasPrefixString(node.pattern) > 0 {
				return r.insert(root.children[i], node)
			}
		}

		root.insertChild(node)
		return node
	}

	// find radix
	pos := root.hasPrefixString(node.pattern)
	if pos == 0 {
		panic("Router.insert error root[" + root.pattern + "] and node[" + node.pattern + "] not have both prefix")
	}
	if pos == len(node.pattern) {
		root.parent.deleteChild(root)
		root.parent.insertChild(node)
		root.pattern = root.pattern[pos:]
		node.insertChild(root)
		return node
	}

	node.pattern = node.pattern[pos:]
	if pos == len(root.pattern) {
		for i = range root.children {
			if root.children[i].pattern == node.pattern {
				if node.handlers != nil {
					root.children[i].handlers = node.handlers
				}
				return root.children[i]
			}
			if root.children[i].hasPrefixString(node.pattern) > 0 {
				return r.insert(root.children[i], node)
			}
		}

		root.insertChild(node)
		return node
	}

	// _parent root and ru has new parent
	_parent := newRoute(root.pattern[:pos], nil, r)
	root.parent.deleteChild(root)
	root.parent.insertChild(_parent)
	root.pattern = root.pattern[pos:]
	_parent.insertChild(root)
	_parent.insertChild(node)

	return node
}

func (r *Router) lookup(pattern string, root *Route, c *Context) *Route {
	var i int
	// static route
	if !root.hasParam {
		if pattern == root.pattern {
			return root
		}
		if len(root.children) == 0 {
			return nil
		}
		if i = root.hasPrefixString(pattern); i == len(root.pattern) {
			pattern = pattern[len(root.pattern):]
		} else {
			return nil
		}
	} else {
		if len(root.children) == 0 {
			i = len(pattern)
		} else {
			for i = 0; i < len(pattern) && isParamChar(pattern[i]); i++ {
			}
		}
		c.SetParam(root.pattern[1:], pattern[:i])
		if i == len(pattern) {
			return root
		}
		pattern = pattern[i:]
	}

	// first, static route
	for i = 0; i < len(root.children); i++ {
		if !root.children[i].hasParam && len(pattern) < len(root.children[i].pattern) {
			continue
		}
		if node := r.lookup(pattern, root.children[i], c); node != nil && node.handlers != nil {
			return node
		}
	}

	return nil
}

// print the route map
func (r *Router) print(prefix string, root *Route) {
	if root == nil {
		for m := range r.routeMap {
			fmt.Println(m)
			r.print("", r.routeMap[m])
		}
		return
	}
	fmt.Println(prefix + " -> " + root.pattern)
	for i := range root.children {
		r.print(prefix+" -> "+root.pattern, root.children[i])
	}
}

// Name set name of route
func (r *Route) Name(name string) {
	if name == "" {
		return
	}
	p := make([]byte, 0, len(r.pattern))
	for i := 0; i < len(r.pattern); i++ {
		if r.pattern[i] != ':' {
			p = append(p, r.pattern[i])
			continue
		}
		p = append(p, '%')
		p = append(p, 'v')
		for i = i + 1; i < len(r.pattern); i++ {
			if !isParamChar(r.pattern[i]) {
				i--
				break
			}
		}
	}
	r.router.routeNamedMap[name] = string(p)
}

// deleteChild find child and delete from root route
func (r *Route) deleteChild(child *Route) {
	for i := 0; i < len(r.children); i++ {
		if r.children[i].pattern != child.pattern {
			continue
		}
		if len(r.children) == 1 {
			r.children = r.children[:0]
			return
		}
		if i == 0 {
			r.children = r.children[1:]
			return
		}
		if i+1 == len(r.children) {
			r.children = r.children[:i]
			return
		}
		r.children = append(r.children[:i], r.children[i+1:]...)
		return
	}
}

// insertChild insert child into root route, and returns the child route
func (r *Route) insertChild(child *Route) *Route {
	for i := 0; i < len(r.children); i++ {
		if r.children[i].pattern == child.pattern {
			if child.handlers != nil {
				r.children[i].handlers = child.handlers
			}
			return r.children[i]
		}
	}
	child.parent = r
	r.children = append(r.children, child)
	return child
}

// hasChild check root has child, if yes return child route, or reutrn nil
func (r *Route) hasChild(child *Route) *Route {
	for i := 0; i < len(r.children); i++ {
		if r.children[i].pattern == child.pattern {
			return r.children[i]
		}
	}
	return nil
}

// hasPrefix returns the same prefix position, if none return 0
func (r *Route) hasPrefix(child *Route) int {
	return r.hasPrefixString(child.pattern)
}

// hasPrefixString returns the same prefix position, if none return 0
func (r *Route) hasPrefixString(s string) int {
	var i, l int
	l = len(r.pattern)
	for i = 0; i < len(s) && i < l && s[i] == r.pattern[i]; i++ {
	}
	return i
}

// wrapHandlerFunc wrap for context handler chain
func wrapHandlerFunc(h HandlerFunc) HandlerFunc {
	return func(c *Context) {
		h(c)
		c.Next()
	}
}

// isParamChar check the char can used for route params
// a-z->65:90, A-Z->97:122, 0-9->48->57, _->95, :->58
func isParamChar(c byte) bool {
	if (c >= 65 && c <= 90) || (c >= 97 && c <= 122) || (c >= 48 && c <= 58) || c == 95 {
		return true
	}
	return false
}
<commit_msg>optimize router<commit_after>package baa

import (
	"fmt"
	"sync"
)

const (
	// method key in routeMap
	GET int = iota
	POST
	PUT
	DELETE
	PATCH
	OPTIONS
	HEAD
	// RouteLength route table length
	RouteLength
)

// methods declare support HTTP method
var methods = map[string]bool{
	"GET":     true,
	"POST":    true,
	"PUT":     true,
	"DELETE":  true,
	"PATCH":   true,
	"OPTIONS": true,
	"HEAD":    true,
}

// methodKeys declare method key in routeMap
var methodKeys = map[string]int{
	"GET":     GET,
	"POST":    POST,
	"PUT":     PUT,
	"DELETE":  DELETE,
	"PATCH":   PATCH,
	"OPTIONS": OPTIONS,
	"HEAD":    HEAD,
}

// Router provlider router for baa
type Router struct {
	autoHead        bool
	mu              sync.RWMutex
	notFoundHandler HandlerFunc
	groups          []*group
	routeMap        [RouteLength]*Route
	routeNamedMap   map[string]string
}

// Route is a tree node
// route use radix tree
type Route struct {
	hasParam  bool
	parent    *Route
	router    *Router
	children  []*Route
	handlers  []HandlerFunc
	pattern   string
	paramName string
}

// group route
type group struct {
	pattern  string
	handlers []HandlerFunc
}

// newRouter create a router instance
func newRouter() *Router {
	r := new(Router)
	for i := 0; i < len(r.routeMap); i++ {
		r.routeMap[i] = newRoute("/", nil, nil)
	}
	r.routeNamedMap = make(map[string]string)
	r.groups = make([]*group, 0)
	return r
}

// newRoute create a route item
func newRoute(pattern string, handles []HandlerFunc, router *Router) *Route {
	r := new(Route)
	r.pattern = pattern
	r.handlers = handles
	r.router = router
	r.children = make([]*Route, 0)
	return r
}

// newGroup create a group router
func newGroup() *group {
	g := new(group)
	g.handlers = make([]HandlerFunc, 0)
	return g
}

// match match the uri for handler
func (r *Router) match(method, uri string, c *Context) *Route {
	return r.lookup(uri, r.routeMap[methodKeys[method]], c)
}

// urlFor use named route return format url
func (r *Router) urlFor(name string, args ...interface{}) string {
	if name == "" {
		return ""
	}
	url := r.routeNamedMap[name]
	if url == "" {
		return ""
	}
	return fmt.Sprintf(url, args...)
}

// groupAdd add a group route has same prefix and handle chain
func (r *Router) groupAdd(pattern string, f func(), handlers []HandlerFunc) {
	g := newGroup()
	g.pattern = pattern
	g.handlers = handlers
	r.groups = append(r.groups, g)

	f()

	r.groups = r.groups[:len(r.groups)-1]
}

// Handle registers a new request handle with the given pattern, method and handlers.
func (r *Router) add(method string, pattern string, handlers []HandlerFunc) *Route {
	if _, ok := methods[method]; !ok {
		panic("unsupport http method [" + method + "]")
	}
	if pattern == "" {
		panic("route pattern can not be emtpy!")
	}
	if pattern[0] != '/' {
		panic("route pattern must begin /")
	}
	r.mu.Lock()
	defer r.mu.Unlock()

	// check group set
	if len(r.groups) > 0 {
		var gpattern string
		var ghandlers []HandlerFunc
		for i := range r.groups {
			gpattern += r.groups[i].pattern
			if len(r.groups[i].handlers) > 0 {
				ghandlers = append(ghandlers, r.groups[i].handlers...)
			}
		}
		pattern = gpattern + pattern
		ghandlers = append(ghandlers, handlers...)
		handlers = ghandlers
	}

	for i := 0; i < len(handlers); i++ {
		handlers[i] = wrapHandlerFunc(handlers[i])
	}

	root := r.routeMap[methodKeys[method]]
	var radix []byte
	var param []byte
	var i, k int
	var tru *Route
	for i = 0; i < len(pattern); i++ {
		//param route
		if pattern[i] == ':' {
			// clear static route
			if len(radix) > 0 {
				root = r.insert(root, newRoute(string(radix), nil, nil))
				radix = radix[:0]
			}
			// set param route
			param = param[:0]
			k = 0
			for i = i + 1; i < len(pattern); i++ {
				if !isParamChar(pattern[i]) {
					i--
					break
				}
				param = append(param, pattern[i])
				k++
			}
			if k == 0 {
				panic("route pattern param is empty")
			}
			// check last character
			if i == len(pattern) {
				tru = newRoute(":", handlers, r)
			} else {
				tru = newRoute(":", nil, nil)
			}
			tru.paramName = string(param[:k])
			tru.hasParam = true
			root = r.insert(root, tru)
			continue
		}
		radix = append(radix, pattern[i])
	}

	// static route
	if len(radix) > 0 {
		tru = newRoute(string(radix), handlers, r)
		r.insert(root, tru)
	}

	return newRoute(pattern, handlers, r)
}

// insert build the route tree
func (r *Router) insert(root *Route, node *Route) *Route {
	// same route
	if root.pattern == node.pattern {
		if node.handlers != nil {
			root.handlers = node.handlers
		}
		return root
	}

	// param route
	if root.hasParam && node.hasParam {
		if root.parent == nil {
			panic("Router.insert error route has no parent")
		}
		return root.parent.insertChild(node)
	}

	if !root.hasParam && node.hasParam {
		return root.insertChild(node)
	}

	var i int
	if root.hasParam && !node.hasParam {
		for i = range root.children {
			if root.children[i].pattern == node.pattern {
				if node.handlers != nil {
					root.children[i].handlers = node.handlers
				}
				return root.children[i]
			}
			if root.children[i].hasPrefixString(node.pattern) > 0 {
				return r.insert(root.children[i], node)
			}
		}

		root.insertChild(node)
		return node
	}

	// find radix
	pos := root.hasPrefixString(node.pattern)
	if pos == 0 {
		panic("Router.insert error root[" + root.pattern + "] and node[" + node.pattern + "] not have both prefix")
	}
	if pos == len(node.pattern) {
		root.parent.deleteChild(root)
		root.parent.insertChild(node)
		root.pattern = root.pattern[pos:]
		node.insertChild(root)
		return node
	}

	node.pattern = node.pattern[pos:]
	if pos == len(root.pattern) {
		for i = range root.children {
			if root.children[i].pattern == node.pattern {
				if node.handlers != nil {
					root.children[i].handlers = node.handlers
				}
				return root.children[i]
			}
			if root.children[i].hasPrefixString(node.pattern) > 0 {
				return r.insert(root.children[i], node)
			}
		}

		root.insertChild(node)
		return node
	}

	// _parent root and ru has new parent
	_parent := newRoute(root.pattern[:pos], nil, r)
	root.parent.deleteChild(root)
	root.parent.insertChild(_parent)
	root.pattern = root.pattern[pos:]
	_parent.insertChild(root)
	_parent.insertChild(node)

	return node
}

// lookup find a route
func (r *Router) lookup(pattern string, root *Route, c *Context) *Route {
	var i, l int
	// static route
	if !root.hasParam {
		if pattern == root.pattern {
			return root
		}
		if len(root.children) == 0 {
			return nil
		}
		if root.hasPrefixString(pattern) == len(root.pattern) {
			pattern = pattern[len(root.pattern):]
		} else {
			return nil
		}
	} else {
		l = len(pattern)
		if len(root.children) == 0 {
			i = l
		} else {
			for i = 0; i < l && isParamChar(pattern[i]); i++ {
			}
		}
		c.SetParam(root.paramName, pattern[:i])
		if i == l {
			return root
		}
		pattern = pattern[i:]
	}

	// static route
	l = len(pattern)
	for i = 0; i < len(root.children); i++ {
		if l < len(root.children[i].pattern) {
			continue
		}
		if node := r.lookup(pattern, root.children[i], c); node != nil && node.handlers != nil {
			return node
		}
	}

	return nil
}

// print the route map
func (r *Router) print(prefix string, root *Route) {
	if root == nil {
		for m := range r.routeMap {
			fmt.Println(m)
			r.print("", r.routeMap[m])
		}
		return
	}
	fmt.Println(prefix + " -> " + root.pattern)
	for i := range root.children {
		r.print(prefix+" -> "+root.pattern, root.children[i])
	}
}

// Name set name of route
func (r *Route) Name(name string) {
	if name == "" {
		return
	}
	p := make([]byte, 0, len(r.pattern))
	for i := 0; i < len(r.pattern); i++ {
		if r.pattern[i] != ':' {
			p = append(p, r.pattern[i])
			continue
		}
		p = append(p, '%')
		p = append(p, 'v')
		for i = i + 1; i < len(r.pattern); i++ {
			if !isParamChar(r.pattern[i]) {
				i--
				break
			}
		}
	}
	r.router.routeNamedMap[name] = string(p)
}

// deleteChild find child and delete from root route
func (r *Route) deleteChild(node *Route) {
	for i := 0; i < len(r.children); i++ {
		if r.children[i].pattern != node.pattern {
			continue
		}
		if len(r.children) == 1 {
			r.children = r.children[:0]
			return
		}
		if i == 0 {
			r.children = r.children[1:]
			return
		}
		if i+1 == len(r.children) {
			r.children = r.children[:i]
			return
		}
		r.children = append(r.children[:i], r.children[i+1:]...)
		return
	}
}

// insertChild insert child into root route, and returns the child route
func (r *Route) insertChild(node *Route) *Route {
	for i := 0; i < len(r.children); i++ {
		if r.children[i].pattern == node.pattern {
			if r.children[i].hasParam && node.hasParam && r.children[i].paramName != node.paramName {
				panic("Router.insert error cannot use two param [:" + r.children[i].paramName + ", " + node.paramName + "]with same prefix!")
			}
			if node.handlers != nil {
				r.children[i].handlers = node.handlers
			}
			return r.children[i]
		}
	}
	node.parent = r
	r.children = append(r.children, node)
	return node
}

// hasChild check root has child, if yes return child route, or reutrn nil
func (r *Route) hasChild(node *Route) *Route {
	for i := 0; i < len(r.children); i++ {
		if r.children[i].pattern == node.pattern {
			return r.children[i]
		}
	}
	return nil
}

// hasPrefix returns the same prefix position, if none return 0
func (r *Route) hasPrefix(node *Route) int {
	return r.hasPrefixString(node.pattern)
}

// hasPrefixString returns the same prefix position, if none return 0
func (r *Route) hasPrefixString(s string) int {
	var i, l int
	l = len(r.pattern)
	if len(s) < l {
		l = len(s)
	}
	for i = 0; i < l && s[i] == r.pattern[i]; i++ {
	}
	return i
}

// wrapHandlerFunc wrap for context handler chain
func wrapHandlerFunc(h HandlerFunc) HandlerFunc {
	return func(c *Context) {
		h(c)
		c.Next()
	}
}

// isParamChar check the char can used for route params
// a-z->65:90, A-Z->97:122, 0-9->48->57, _->95, :->58
func isParamChar(c byte) bool {
	if (c >= 65 && c <= 90) || (c >= 97 && c <= 122) || (c >= 48 && c <= 58) || c == 95 {
		return true
	}
	return false
}
<|endoftext|>