<commit_before>package handler

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/astaxie/beego/logs"
	"gopkg.in/macaron.v1"

	"github.com/containerops/dockyard/models"
	"github.com/containerops/dockyard/module"
	"github.com/containerops/wrench/setting"
	"github.com/containerops/wrench/utils"
)

var ManifestCtx []byte

func PutManifestsV2Handler(ctx *macaron.Context, log *logs.BeeLogger) (int, []byte) {
	namespace := ctx.Params(":namespace")
	repository := ctx.Params(":repository")

	agent := ctx.Req.Header.Get("User-Agent")

	repo := new(models.Repository)
	if err := repo.Put(namespace, repository, "", agent, setting.APIVERSION_V2); err != nil {
		log.Error("[REGISTRY API V2] Save repository failed: %v", err.Error())

		result, _ := json.Marshal(map[string]string{"message": err.Error()})
		return http.StatusBadRequest, result
	}

	if len(ManifestCtx) == 0 {
		ManifestCtx, _ = ctx.Req.Body().Bytes()
	}

	if err := module.ParseManifest(ManifestCtx); err != nil {
		log.Error("[REGISTRY API V2] Decode Manifest Error: %v", err.Error())

		result, _ := json.Marshal(map[string]string{"message": "Manifest converted failed"})
		return http.StatusBadRequest, result
	}

	digest, err := utils.DigestManifest(ManifestCtx)
	if err != nil {
		log.Error("[REGISTRY API V2] Get manifest digest failed: %v", err.Error())

		result, _ := json.Marshal(map[string]string{"message": "Get manifest digest failed"})
		return http.StatusBadRequest, result
	}

	random := fmt.Sprintf("%s://%s/v2/%s/%s/manifests/%s",
		setting.ListenMode,
		setting.Domains,
		namespace,
		repository,
		digest)

	ctx.Resp.Header().Set("Docker-Content-Digest", digest)
	ctx.Resp.Header().Set("Location", random)

	result, _ := json.Marshal(map[string]string{})
	return http.StatusAccepted, result
}

func GetTagsListV2Handler(ctx *macaron.Context, log *logs.BeeLogger) (int, []byte) {
	namespace := ctx.Params(":namespace")
	repository := ctx.Params(":repository")

	r := new(models.Repository)
	if has, _, err := r.Has(namespace, repository); err != nil || has == false {
		log.Error("[REGISTRY API V2] Repository not found: %v", repository)

		result, _ := json.Marshal(map[string]string{"message": "Repository not found"})
		return http.StatusNotFound, result
	}

	data := map[string]interface{}{}
	tags := []string{}

	data["name"] = fmt.Sprintf("%s/%s", namespace, repository)

	for _, value := range r.Tags {
		t := new(models.Tag)
		if err := t.GetByKey(value); err != nil {
			log.Error("[REGISTRY API V2] Tag not found: %v", err.Error())

			result, _ := json.Marshal(map[string]string{"message": "Tag not found"})
			return http.StatusNotFound, result
		}

		tags = append(tags, t.Name)
	}

	data["tags"] = tags

	result, _ := json.Marshal(data)
	return http.StatusOK, result
}

func GetManifestsV2Handler(ctx *macaron.Context, log *logs.BeeLogger) (int, []byte) {
	t := new(models.Tag)

	namespace := ctx.Params(":namespace")
	repository := ctx.Params(":repository")
	if err := t.Get(namespace, repository, ctx.Params(":tag")); err != nil {
		log.Error("[REGISTRY API V2] Manifest not found: %v", err.Error())

		result, _ := json.Marshal(map[string]string{"message": "Manifest not found"})
		return http.StatusNotFound, result
	}

	digest, err := utils.DigestManifest([]byte(t.Manifest))
	if err != nil {
		log.Error("[REGISTRY API V2] Get manifest digest failed: %v", err.Error())

		result, _ := json.Marshal(map[string]string{"message": "Get manifest digest failed"})
		return http.StatusBadRequest, result
	}

	ctx.Resp.Header().Set("Content-Type", "application/json; charset=utf-8")
	ctx.Resp.Header().Set("Docker-Content-Digest", digest)
	ctx.Resp.Header().Set("Content-Length", fmt.Sprint(len(t.Manifest)))

	return http.StatusOK, []byte(t.Manifest)
}
<commit_msg>Bugfix about Docker V2 PutManifestsV2Handler<commit_after>package handler

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/astaxie/beego/logs"
	"gopkg.in/macaron.v1"

	"github.com/containerops/dockyard/models"
	"github.com/containerops/dockyard/module"
	"github.com/containerops/wrench/setting"
	"github.com/containerops/wrench/utils"
)

var ManifestCtx []byte

func PutManifestsV2Handler(ctx *macaron.Context, log *logs.BeeLogger) (int, []byte) {
	namespace := ctx.Params(":namespace")
	repository := ctx.Params(":repository")

	agent := ctx.Req.Header.Get("User-Agent")

	repo := new(models.Repository)
	if err := repo.Put(namespace, repository, "", agent, setting.APIVERSION_V2); err != nil {
		log.Error("[REGISTRY API V2] Save repository failed: %v", err.Error())

		result, _ := json.Marshal(map[string]string{"message": err.Error()})
		return http.StatusBadRequest, result
	}

	if len(ManifestCtx) == 0 {
		ManifestCtx, _ = ctx.Req.Body().Bytes()
	}

	if err := module.ParseManifest(ManifestCtx); err != nil {
		log.Error("[REGISTRY API V2] Decode Manifest Error: %v", err.Error())

		ManifestCtx = []byte{}
		result, _ := json.Marshal(map[string]string{"message": "Manifest converted failed"})
		return http.StatusBadRequest, result
	}

	digest, err := utils.DigestManifest(ManifestCtx)
	if err != nil {
		log.Error("[REGISTRY API V2] Get manifest digest failed: %v", err.Error())

		ManifestCtx = []byte{}
		result, _ := json.Marshal(map[string]string{"message": "Get manifest digest failed"})
		return http.StatusBadRequest, result
	}

	random := fmt.Sprintf("%s://%s/v2/%s/%s/manifests/%s",
		setting.ListenMode,
		setting.Domains,
		namespace,
		repository,
		digest)

	ctx.Resp.Header().Set("Docker-Content-Digest", digest)
	ctx.Resp.Header().Set("Location", random)

	ManifestCtx = []byte{}
	result, _ := json.Marshal(map[string]string{})
	return http.StatusAccepted, result
}

func GetTagsListV2Handler(ctx *macaron.Context, log *logs.BeeLogger) (int, []byte) {
	namespace := ctx.Params(":namespace")
	repository := ctx.Params(":repository")

	r := new(models.Repository)
	if has, _, err := r.Has(namespace, repository); err != nil || has == false {
		log.Error("[REGISTRY API V2] Repository not found: %v", repository)

		result, _ := json.Marshal(map[string]string{"message": "Repository not found"})
		return http.StatusNotFound, result
	}

	data := map[string]interface{}{}
	tags := []string{}

	data["name"] = fmt.Sprintf("%s/%s", namespace, repository)

	for _, value := range r.Tags {
		t := new(models.Tag)
		if err := t.GetByKey(value); err != nil {
			log.Error("[REGISTRY API V2] Tag not found: %v", err.Error())

			result, _ := json.Marshal(map[string]string{"message": "Tag not found"})
			return http.StatusNotFound, result
		}

		tags = append(tags, t.Name)
	}

	data["tags"] = tags

	result, _ := json.Marshal(data)
	return http.StatusOK, result
}

func GetManifestsV2Handler(ctx *macaron.Context, log *logs.BeeLogger) (int, []byte) {
	t := new(models.Tag)

	namespace := ctx.Params(":namespace")
	repository := ctx.Params(":repository")
	if err := t.Get(namespace, repository, ctx.Params(":tag")); err != nil {
		log.Error("[REGISTRY API V2] Manifest not found: %v", err.Error())

		result, _ := json.Marshal(map[string]string{"message": "Manifest not found"})
		return http.StatusNotFound, result
	}

	digest, err := utils.DigestManifest([]byte(t.Manifest))
	if err != nil {
		log.Error("[REGISTRY API V2] Get manifest digest failed: %v", err.Error())

		result, _ := json.Marshal(map[string]string{"message": "Get manifest digest failed"})
		return http.StatusBadRequest, result
	}

	ctx.Resp.Header().Set("Content-Type", "application/json; charset=utf-8")
	ctx.Resp.Header().Set("Docker-Content-Digest", digest)
	ctx.Resp.Header().Set("Content-Length", fmt.Sprint(len(t.Manifest)))

	return http.StatusOK, []byte(t.Manifest)
}
<|endoftext|>


<commit_before>package dex

import (
	"fmt"
	"html"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"regexp"
	"strconv"

	"github.com/argoproj/argo-cd/errors"
)

var messageRe = regexp.MustCompile(`<p>(.*)([\s\S]*?)<\/p>`)

// NewDexHTTPReverseProxy returns a reverse proxy to the Dex server. Dex is assumed to be configured
// with the external issuer URL muxed to the same path configured in server.go. In other words, if
// Argo CD API server wants to proxy requests at /api/dex, then the dex config yaml issuer URL should
// also be /api/dex (e.g. issuer: https://argocd.example.com/api/dex)
func NewDexHTTPReverseProxy(serverAddr string) func(writer http.ResponseWriter, request *http.Request) {
	target, err := url.Parse(serverAddr)
	errors.CheckError(err)
	proxy := httputil.NewSingleHostReverseProxy(target)
	proxy.ModifyResponse = func(resp *http.Response) error {
		if resp.StatusCode == 500 {
			b, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return err
			}
			err = resp.Body.Close()
			if err != nil {
				return err
			}
			var message string
			matches := messageRe.FindSubmatch(b)
			if len(matches) > 1 {
				message = html.UnescapeString(string(matches[1]))
			} else {
				message = "Unknown error"
			}
			resp.ContentLength = 0
			resp.Header.Set("Content-Length", strconv.Itoa(0))
			resp.Header.Set("Location", fmt.Sprintf("/login?sso_error=%s", url.QueryEscape(message)))
			resp.StatusCode = http.StatusSeeOther
			return nil
		}
		return nil
	}
	return func(w http.ResponseWriter, r *http.Request) {
		proxy.ServeHTTP(w, r)
	}
}
<commit_msg>Issue #1039 - Correct redirect to login page if dex authentication is not successful (#1040)<commit_after>package dex

import (
	"bytes"
	"fmt"
	"html"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"regexp"
	"strconv"

	"github.com/argoproj/argo-cd/errors"
)

var messageRe = regexp.MustCompile(`<p>(.*)([\s\S]*?)<\/p>`)

// NewDexHTTPReverseProxy returns a reverse proxy to the Dex server. Dex is assumed to be configured
// with the external issuer URL muxed to the same path configured in server.go. In other words, if
// Argo CD API server wants to proxy requests at /api/dex, then the dex config yaml issuer URL should
// also be /api/dex (e.g. issuer: https://argocd.example.com/api/dex)
func NewDexHTTPReverseProxy(serverAddr string) func(writer http.ResponseWriter, request *http.Request) {
	target, err := url.Parse(serverAddr)
	errors.CheckError(err)
	proxy := httputil.NewSingleHostReverseProxy(target)
	proxy.ModifyResponse = func(resp *http.Response) error {
		if resp.StatusCode == 500 {
			b, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return err
			}
			err = resp.Body.Close()
			if err != nil {
				return err
			}
			var message string
			matches := messageRe.FindSubmatch(b)
			if len(matches) > 1 {
				message = html.UnescapeString(string(matches[1]))
			} else {
				message = "Unknown error"
			}
			resp.ContentLength = 0
			resp.Header.Set("Content-Length", strconv.Itoa(0))
			resp.Header.Set("Location", fmt.Sprintf("/login?sso_error=%s", url.QueryEscape(message)))
			resp.StatusCode = http.StatusSeeOther
			resp.Body = ioutil.NopCloser(bytes.NewReader(make([]byte, 0)))
			return nil
		}
		return nil
	}
	return func(w http.ResponseWriter, r *http.Request) {
		proxy.ServeHTTP(w, r)
	}
}
<|endoftext|>


<commit_before>// Package dns ...
package dns

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/nanobox-io/nanobox/commands/server"
)

type DomainName struct {
	IP     string
	Domain string
}

type DomainRPC struct{}

type Request struct {
	Entry string
}

type Response struct {
	Message string
	Success bool
}

var (
	hostsFile = detectHostsFile()
	newline   = detectNewlineChar()
)

func init() {
	server.Register(&DomainRPC{})
}

// Entry generate the DNS entry to be added
func Entry(ip, name, env string) string {
	return fmt.Sprintf("%s     %s # dns added for '%s' by nanobox", ip, name, env)
}

// Exists ...
func Exists(entry string) bool {

	// open the hosts file for scanning...
	f, err := os.Open(hostsFile)
	if err != nil {
		return false
	}
	defer f.Close()

	// scan each line of the hosts file to see if there is a match for this
	// entry
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		if scanner.Text() == entry {
			return true
		}
	}

	return false
}

func List(filter string) []DomainName {

	// open the hosts file
	f, err := os.Open(hostsFile)
	if err != nil {
		return nil
	}
	defer f.Close()

	entries := []DomainName{}

	// scan each line of the hosts file to see if there is a match for this
	// entry
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), filter) {
			fields := strings.Fields(scanner.Text())
			if len(fields) >= 2 {
				entries = append(entries, DomainName{IP: fields[0], Domain: fields[1]})
			}
		}
	}

	return entries
}

// Add ...
func Add(entry string) error {
	// break early if there is no entry
	// or we have already added this entry
	if entry == "" || Exists(entry) {
		return nil
	}

	req := Request{entry}
	resp := &Response{}

	err := server.ClientRun("DomainRPC.Add", req, resp)
	if !resp.Success {
		err = fmt.Errorf("failed to add domain: %v %v", err, resp.Message)
	}

	return err
}

// the rpc function run from the server
func (drpc *DomainRPC) Add(req Request, resp *Response) error {

	// open hosts file
	f, err := os.OpenFile(hostsFile, os.O_RDWR|os.O_APPEND, 0644)
	if err != nil {
		return err
	}
	defer f.Close()

	// write the DNS entry to the file
	if _, err := f.WriteString(fmt.Sprintf("%s%s", req.Entry, newline)); err != nil {
		return err
	}

	resp.Success = true
	return nil
}

// Remove ...
func Remove(entry string) error {
	if entry == "" {
		return nil
	}
	req := Request{entry}
	resp := &Response{}

	err := server.ClientRun("DomainRPC.Remove", req, resp)
	if !resp.Success {
		err = fmt.Errorf("failed to remove domain: %v %v", err, resp.Message)
	}
	return nil
}

// the rpc function run from the server
func (drpc *DomainRPC) Remove(req Request, resp *Response) error {

	// "contents" will end up storing the entire contents of the file excluding the
	// entry that is trying to be removed
	var contents string

	// open hosts file
	f, err := os.OpenFile(hostsFile, os.O_RDWR, 0644)
	if err != nil {
		return err
	}
	defer f.Close()

	// remove entry from /etc/hosts
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {

		// if the line contain the entry skip it
		// make it do a loose string check
		// if its exactly the entry then remove it.
		// if it contains the same environment
		// also remove it
		if strings.Contains(scanner.Text(), req.Entry) {
			continue
		}

		// add each line back into the file
		contents += fmt.Sprintf("%s%s", scanner.Text(), newline)
	}

	// trim the contents to avoid any extra newlines
	contents = strings.TrimSpace(contents)

	// add a single newline for completeness
	contents += newline

	// write back the contents of the hosts file minus the removed entry
	if err := ioutil.WriteFile(hostsFile, []byte(contents), 0644); err != nil {
		return err
	}

	resp.Success = true
	return nil

}

// RemoveAll removes all dns entries added by nanobox
func RemoveAll() error {

	// short-circuit if no entries were added by nanobox
	if len(List("by nanobox")) == 0 {
		return nil
	}

	return Remove("by nanobox")
}
<commit_msg>allow dns's to work when something doesnt have a newline fixes #485<commit_after>// Package dns ...
package dns

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/nanobox-io/nanobox/commands/server"
)

type DomainName struct {
	IP     string
	Domain string
}

type DomainRPC struct{}

type Request struct {
	Entry string
}

type Response struct {
	Message string
	Success bool
}

var (
	hostsFile = detectHostsFile()
	newline   = detectNewlineChar()
)

func init() {
	server.Register(&DomainRPC{})
}

// Entry generate the DNS entry to be added
func Entry(ip, name, env string) string {
	return fmt.Sprintf("%s     %s # dns added for '%s' by nanobox", ip, name, env)
}

// Exists ...
func Exists(entry string) bool {

	// open the hosts file for scanning...
	f, err := os.Open(hostsFile)
	if err != nil {
		return false
	}
	defer f.Close()

	// scan each line of the hosts file to see if there is a match for this
	// entry
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		if scanner.Text() == entry {
			return true
		}
	}

	return false
}

func List(filter string) []DomainName {

	// open the hosts file
	f, err := os.Open(hostsFile)
	if err != nil {
		return nil
	}
	defer f.Close()

	entries := []DomainName{}

	// scan each line of the hosts file to see if there is a match for this
	// entry
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), filter) {
			fields := strings.Fields(scanner.Text())
			if len(fields) >= 2 {
				entries = append(entries, DomainName{IP: fields[0], Domain: fields[1]})
			}
		}
	}

	return entries
}

// Add ...
func Add(entry string) error {
	// break early if there is no entry
	// or we have already added this entry
	if entry == "" || Exists(entry) {
		return nil
	}

	req := Request{entry}
	resp := &Response{}

	err := server.ClientRun("DomainRPC.Add", req, resp)
	if !resp.Success {
		err = fmt.Errorf("failed to add domain: %v %v", err, resp.Message)
	}

	return err
}

// the rpc function run from the server
func (drpc *DomainRPC) Add(req Request, resp *Response) error {

	// open hosts file
	f, err := os.OpenFile(hostsFile, os.O_RDWR|os.O_APPEND, 0644)
	if err != nil {
		return err
	}
	defer f.Close()

	// write the DNS entry to the file
	// we place a newline before and after because
	// the extra new lines wont hurt and it ensures success
	if _, err := f.WriteString(fmt.Sprintf("%s%s%s", newline, req.Entry, newline)); err != nil {
		return err
	}

	resp.Success = true
	return nil
}

// Remove ...
func Remove(entry string) error {
	if entry == "" {
		return nil
	}
	req := Request{entry}
	resp := &Response{}

	err := server.ClientRun("DomainRPC.Remove", req, resp)
	if !resp.Success {
		err = fmt.Errorf("failed to remove domain: %v %v", err, resp.Message)
	}
	return nil
}

// the rpc function run from the server
func (drpc *DomainRPC) Remove(req Request, resp *Response) error {

	// "contents" will end up storing the entire contents of the file excluding the
	// entry that is trying to be removed
	var contents string

	// open hosts file
	f, err := os.OpenFile(hostsFile, os.O_RDWR, 0644)
	if err != nil {
		return err
	}
	defer f.Close()

	// remove entry from /etc/hosts
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {

		// if the line contain the entry skip it
		// make it do a loose string check
		// if its exactly the entry then remove it.
		// if it contains the same environment
		// also remove it
		if strings.Contains(scanner.Text(), req.Entry) {
			continue
		}

		// add each line back into the file
		contents += fmt.Sprintf("%s%s", scanner.Text(), newline)
	}

	// trim the contents to avoid any extra newlines
	contents = strings.TrimSpace(contents)

	// add a single newline for completeness
	contents += newline

	// write back the contents of the hosts file minus the removed entry
	if err := ioutil.WriteFile(hostsFile, []byte(contents), 0644); err != nil {
		return err
	}

	resp.Success = true
	return nil

}

// RemoveAll removes all dns entries added by nanobox
func RemoveAll() error {

	// short-circuit if no entries were added by nanobox
	if len(List("by nanobox")) == 0 {
		return nil
	}

	return Remove("by nanobox")
}
<|endoftext|>


<commit_before>// Copyright © 2017 Alexander Sosna <alexander@xxor.de>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
package util

import (
	"errors"
	"fmt"
	"io"
	"os"
	"strings"

	log "github.com/Sirupsen/logrus"
)

func AnswerConfirmation(msg string) (confirmed bool, err error) {
	var input string
	log.Warn(msg)

	_, err = fmt.Scanln(&input)
	if err != nil {
		return false, err
	}
	positive := []string{"j", "ja", "y", "yes", "do it", "let's rock"}
	negative := []string{"nein", "n", "no", "hell no", "fuck off"}

	input = strings.ToLower(input)

	for _, element := range positive {
		if element == input {
			return true, nil
		}
	}

	for _, element := range negative {
		if element == input {
			return false, nil
		}
	}
	doesNotParse := errors.New("Answer can not be parsed: " + input)
	return false, doesNotParse
}

func MustAnswerConfirmation(msg string) (confirmed bool) {
	if confirmed, err := AnswerConfirmation(msg); err == nil {
		return confirmed
	}
	return MustAnswerConfirmation(msg)
}

func WatchOutput(input io.Reader, outputFunc func(args ...interface{})) {
	log.Debug("nop")
}

// Exists returns whether the given file or directory exists or not
func Exists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return true, err
}

// IsEmpty returns true if the given dir is empty
func IsEmpty(path string) (bool, error) {
	f, err := os.Open(path)
	if err != nil {
		return false, err
	}
	defer f.Close()

	// Readdirnames returns at most n names, or io.EOF if not enough are available
	_, err = f.Readdirnames(1)
	if err == io.EOF {
		return true, nil
	}
	return false, err
}
<commit_msg>Revert test<commit_after>// Copyright © 2017 Alexander Sosna <alexander@xxor.de>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
package util

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"

	log "github.com/Sirupsen/logrus"
)

func AnswerConfirmation(msg string) (confirmed bool, err error) {
	var input string
	log.Warn(msg)

	_, err = fmt.Scanln(&input)
	if err != nil {
		return false, err
	}
	positive := []string{"j", "ja", "y", "yes", "do it", "let's rock"}
	negative := []string{"nein", "n", "no", "hell no", "fuck off"}

	input = strings.ToLower(input)

	for _, element := range positive {
		if element == input {
			return true, nil
		}
	}

	for _, element := range negative {
		if element == input {
			return false, nil
		}
	}
	doesNotParse := errors.New("Answer can not be parsed: " + input)
	return false, doesNotParse
}

func MustAnswerConfirmation(msg string) (confirmed bool) {
	if confirmed, err := AnswerConfirmation(msg); err == nil {
		return confirmed
	}
	return MustAnswerConfirmation(msg)
}

func WatchOutput(input io.Reader, outputFunc func(args ...interface{})) {
	log.Debug("watchOutput started")
	scanner := bufio.NewScanner(input)
	for scanner.Scan() {
		outputFunc(scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		outputFunc("error reading input:", err)
	}
	log.Debug("watchOutput end")
}

// Exists returns whether the given file or directory exists or not
func Exists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return true, err
}

// IsEmpty returns true if the given dir is empty
func IsEmpty(path string) (bool, error) {
	f, err := os.Open(path)
	if err != nil {
		return false, err
	}
	defer f.Close()

	// Readdirnames returns at most n names, or io.EOF if not enough are available
	_, err = f.Readdirnames(1)
	if err == io.EOF {
		return true, nil
	}
	return false, err
}
<|endoftext|>


<commit_before>/*
The MIT License (MIT)

Copyright (c) 2016 winlin

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

/*
 This the main entrance of https-proxy, proxy to api or other http server.
*/
package main

import (
	"context"
	"crypto/tls"
	"flag"
	"fmt"
	oe "github.com/ossrs/go-oryx-lib/errors"
	"github.com/ossrs/go-oryx-lib/https"
	ol "github.com/ossrs/go-oryx-lib/logger"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"path"
	"strings"
	"sync"
)

const server = "Oryx/0.0.3"

type Proxies []string

func (v *Proxies) String() string {
	return "proxy to backend services"
}

func (v *Proxies) Set(value string) error {
	*v = append(*v, value)
	return nil
}

func run(ctx context.Context) error {
	fmt.Println(server, "HTTP/HTTPS static server with API proxy.")

	var httpPort, httpsPort int
	var httpsDomains, html, cacheFile string
	var useLetsEncrypt bool
	var ssCert, ssKey string
	var oproxies Proxies
	flag.IntVar(&httpPort, "http", 0, "http listen at. 0 to disable http.")
	flag.IntVar(&httpsPort, "https", 0, "https listen at. 0 to disable https. 443 to serve. ")
	flag.StringVar(&httpsDomains, "domains", "", "the allow domains, empty to allow all. for example: ossrs.net,www.ossrs.net")
	flag.StringVar(&html, "root", "./html", "the www web root. support relative dir to argv[0].")
	flag.StringVar(&cacheFile, "cache", "./letsencrypt.cache", "the cache for https. support relative dir to argv[0].")
	flag.BoolVar(&useLetsEncrypt, "lets", false, "whether use letsencrypt CA. self sign if not.")
	flag.StringVar(&ssKey, "ssk", "server.key", "https self-sign key by(before server.cert): openssl genrsa -out server.key 2048")
	flag.StringVar(&ssCert, "ssc", "server.crt", `https self-sign cert by: openssl req -new -x509 -key server.key -out server.crt -days 365 -subj "/C=CN/ST=Beijing/L=Beijing/O=Me/OU=Me/CN=me.org"`)
	flag.Var(&oproxies, "proxy", "one or more proxy the matched path to backend, for example, -proxy http://127.0.0.1:8888/api/webrtc")
	flag.Parse()

	if useLetsEncrypt && (httpsPort != 0 && httpsPort != 443) {
		return oe.Errorf("for letsencrypt, https=%v must be 0(disabled) or 443(enabled)", httpsPort)
	}
	if httpPort == 0 && httpsPort == 0 {
		flag.PrintDefaults()
		os.Exit(-1)
	}

	var proxyUrls []*url.URL
	proxies := map[string]*httputil.ReverseProxy{}
	for _, oproxy := range []string(oproxies) {
		if oproxy == "" {
			return oe.Errorf("empty proxy in %v", oproxies)
		}

		proxyUrl, err := url.Parse(oproxy)
		if err != nil {
			return oe.Wrapf(err, "parse proxy %v", oproxy)
		}

		proxy := &httputil.ReverseProxy{
			Director: func(r *http.Request) {
				r.URL.Scheme = proxyUrl.Scheme
				r.URL.Host = proxyUrl.Host

				// about x-real-ip and x-forwarded-for
				// https://segmentfault.com/q/1010000002409659
				// https://distinctplace.com/2014/04/23/story-behind-x-forwarded-for-and-x-real-ip-headers/
				// @remark http proxy will set the X-Forwarded-For.
				if rip := r.Header.Get("X-Real-IP"); rip != "" {
					r.Header.Set("X-Real-IP", rip)
				} else if ip, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
					r.Header.Set("X-Real-IP", ip)
				}
				//ol.Tf(ctx, "proxy http %v to %v", r.RemoteAddr, r.URL.String())
			},
		}

		if _, ok := proxies[proxyUrl.Path]; ok {
			return oe.Errorf("proxy %v duplicated", proxyUrl.Path)
		}

		proxyUrls = append(proxyUrls, proxyUrl)
		proxies[proxyUrl.Path] = proxy
		ol.Tf(ctx, "Proxy %v to %v", proxyUrl.Path, oproxy)
	}

	if !path.IsAbs(cacheFile) && path.IsAbs(os.Args[0]) {
		cacheFile = path.Join(path.Dir(os.Args[0]), cacheFile)
	}
	if !path.IsAbs(html) && path.IsAbs(os.Args[0]) {
		html = path.Join(path.Dir(os.Args[0]), html)
	}

	fs := http.FileServer(http.Dir(html))
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Server", server)

		if o := r.Header.Get("Origin"); len(o) > 0 {
			w.Header().Set("Access-Control-Allow-Origin", "*")
			w.Header().Set("Access-Control-Allow-Methods", "GET, POST, HEAD, PUT, DELETE, OPTIONS")
			w.Header().Set("Access-Control-Expose-Headers", "Server,range,Content-Length,Content-Range")
			w.Header().Set("Access-Control-Allow-Headers", "origin,range,accept-encoding,referer,Cache-Control,X-Proxy-Authorization,X-Requested-With,Content-Type")
		}

		if proxyUrls == nil {
			fs.ServeHTTP(w, r)
			return
		}

		for _, proxyUrl := range proxyUrls {
			srcPath, proxyPath := r.URL.Path, proxyUrl.Path
			if !strings.HasSuffix(srcPath, "/") {
				// /api to /api/
				// /api.js to /api.js/
				// /api/100 to /api/100/
				srcPath += "/"
			}
			if !strings.HasSuffix(proxyPath, "/") {
				// /api/ to /api/
				// to match /api/ or /api/100
				// and not match /api.js/
				proxyPath += "/"
			}
			if !strings.HasPrefix(srcPath, proxyPath) {
				continue
			}

			// For matched OPTIONS, directly return without response.
			if r.Method == "OPTIONS" {
				return
			}

			if proxy, ok := proxies[proxyUrl.Path]; ok {
				proxy.ServeHTTP(w, r)
				return
			}
		}

		fs.ServeHTTP(w, r)
	})

	var protos []string
	if httpPort != 0 {
		protos = append(protos, fmt.Sprintf("http(:%v)", httpPort))
	}
	if httpsPort != 0 {
		s := httpsDomains
		if httpsDomains == "" {
			s = "all domains"
		}

		if useLetsEncrypt {
			protos = append(protos, fmt.Sprintf("https(:%v, %v, %v)", httpsPort, s, cacheFile))
		} else {
			protos = append(protos, fmt.Sprintf("https(:%v)", httpsPort))
		}

		if useLetsEncrypt {
			protos = append(protos, "letsencrypt")
		} else {
			protos = append(protos, fmt.Sprintf("self-sign(%v, %v)", ssKey, ssCert))
		}
	}
	ol.Tf(ctx, "%v html root at %v", strings.Join(protos, ", "), string(html))

	var hs, hss *http.Server

	wg := sync.WaitGroup{}
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	var err error
	go func() {
		defer wg.Done()
		if httpPort == 0 {
			ol.W(ctx, "http server disabled")
			return
		}

		defer cancel()
		hs = &http.Server{Addr: fmt.Sprintf(":%v", httpPort), Handler: nil}
		ol.Tf(ctx, "http serve at %v", httpPort)

		if err = hs.ListenAndServe(); err != nil {
			err = oe.Wrapf(err, "serve http")
			return
		}
		ol.T("http server ok")
	}()
	wg.Add(1)

	go func() {
		defer wg.Done()
		if httpsPort == 0 {
			ol.W(ctx, "https server disabled")
			return
		}

		defer cancel()

		var m https.Manager
		if useLetsEncrypt {
			var domains []string
			if httpsDomains != "" {
				domains = strings.Split(httpsDomains, ",")
			}

			if m, err = https.NewLetsencryptManager("", domains, cacheFile); err != nil {
				err = oe.Wrapf(err, "create letsencrypt manager")
				return
			}
		} else {
			if m, err = https.NewSelfSignManager(ssCert, ssKey); err != nil {
				err = oe.Wrapf(err, "create self-sign manager")
				return
			}
		}

		hss = &http.Server{
			Addr: fmt.Sprintf(":%v", httpsPort),
			TLSConfig: &tls.Config{
				GetCertificate: m.GetCertificate,
			},
		}
		ol.Tf(ctx, "http serve at %v", httpsPort)

		if err = hss.ListenAndServeTLS("", ""); err != nil {
			err = oe.Wrapf(err, "listen and serve https")
			return
		}
		ol.T("https serve ok")
	}()
	wg.Add(1)

	select {
	case <-ctx.Done():
		if hs != nil {
			hs.Close()
		}
		if hss != nil {
			hss.Close()
		}
	}
	wg.Wait()

	return err
}

func main() {
	ctx := context.Background()
	if err := run(ctx); err != nil {
		ol.Ef(ctx, "run err %+v", err)
		os.Exit(-1)
	}
}
<commit_msg>Refine proxy, add schema to header<commit_after>/*
The MIT License (MIT)

Copyright (c) 2016 winlin

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

/*
 This the main entrance of https-proxy, proxy to api or other http server.
*/
package main

import (
	"context"
	"crypto/tls"
	"flag"
	"fmt"
	oe "github.com/ossrs/go-oryx-lib/errors"
	"github.com/ossrs/go-oryx-lib/https"
	ol "github.com/ossrs/go-oryx-lib/logger"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"path"
	"strings"
	"sync"
)

const server = "Oryx/0.0.3"

type Proxies []string

func (v *Proxies) String() string {
	return "proxy to backend services"
}

func (v *Proxies) Set(value string) error {
	*v = append(*v, value)
	return nil
}

func run(ctx context.Context) error {
	fmt.Println(server, "HTTP/HTTPS static server with API proxy.")

	var httpPort, httpsPort int
	var httpsDomains, html, cacheFile string
	var useLetsEncrypt bool
	var ssCert, ssKey string
	var oproxies Proxies
	flag.IntVar(&httpPort, "http", 0, "http listen at. 0 to disable http.")
	flag.IntVar(&httpsPort, "https", 0, "https listen at. 0 to disable https. 443 to serve. ")
	flag.StringVar(&httpsDomains, "domains", "", "the allow domains, empty to allow all. for example: ossrs.net,www.ossrs.net")
	flag.StringVar(&html, "root", "./html", "the www web root. support relative dir to argv[0].")
	flag.StringVar(&cacheFile, "cache", "./letsencrypt.cache", "the cache for https. support relative dir to argv[0].")
	flag.BoolVar(&useLetsEncrypt, "lets", false, "whether use letsencrypt CA. self sign if not.")
	flag.StringVar(&ssKey, "ssk", "server.key", "https self-sign key by(before server.cert): openssl genrsa -out server.key 2048")
	flag.StringVar(&ssCert, "ssc", "server.crt", `https self-sign cert by: openssl req -new -x509 -key server.key -out server.crt -days 365 -subj "/C=CN/ST=Beijing/L=Beijing/O=Me/OU=Me/CN=me.org"`)
	flag.Var(&oproxies, "proxy", "one or more proxy the matched path to backend, for example, -proxy http://127.0.0.1:8888/api/webrtc")
	flag.Parse()

	if useLetsEncrypt && (httpsPort != 0 && httpsPort != 443) {
		return oe.Errorf("for letsencrypt, https=%v must be 0(disabled) or 443(enabled)", httpsPort)
	}
	if httpPort == 0 && httpsPort == 0 {
		flag.PrintDefaults()
		os.Exit(-1)
	}

	var proxyUrls []*url.URL
	proxies := map[string]*httputil.ReverseProxy{}
	for _, oproxy := range []string(oproxies) {
		if oproxy == "" {
			return oe.Errorf("empty proxy in %v", oproxies)
		}

		proxyUrl, err := url.Parse(oproxy)
		if err != nil {
			return oe.Wrapf(err, "parse proxy %v", oproxy)
		}

		proxy := &httputil.ReverseProxy{
			Director: func(r *http.Request) {
				// about the x-real-schema, we proxy to backend to identify the client schema.
				if rschema := r.Header.Get("X-Real-Schema"); rschema == "" {
					if r.TLS == nil {
						r.Header.Set("X-Real-Schema", "http")
					} else {
						r.Header.Set("X-Real-Schema", "https")
					}
				}

				// about x-real-ip and x-forwarded-for or
				// about X-Real-IP and X-Forwarded-For or
				// https://segmentfault.com/q/1010000002409659
				// https://distinctplace.com/2014/04/23/story-behind-x-forwarded-for-and-x-real-ip-headers/
				// @remark http proxy will set the X-Forwarded-For.
				if rip := r.Header.Get("X-Real-IP"); rip == "" {
					if rip, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
						r.Header.Set("X-Real-IP", rip)
					}
				}

				r.URL.Scheme = proxyUrl.Scheme
				r.URL.Host = proxyUrl.Host

				//ol.Tf(ctx, "proxy http %v to %v", r.RemoteAddr, r.URL.String())
			},
		}

		if _, ok := proxies[proxyUrl.Path]; ok {
			return oe.Errorf("proxy %v duplicated", proxyUrl.Path)
		}

		proxyUrls = append(proxyUrls, proxyUrl)
		proxies[proxyUrl.Path] = proxy
		ol.Tf(ctx, "Proxy %v to %v", proxyUrl.Path, oproxy)
	}

	if !path.IsAbs(cacheFile) && path.IsAbs(os.Args[0]) {
		cacheFile = path.Join(path.Dir(os.Args[0]), cacheFile)
	}
	if !path.IsAbs(html) && path.IsAbs(os.Args[0]) {
		html = path.Join(path.Dir(os.Args[0]), html)
	}

	fs := http.FileServer(http.Dir(html))
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Server", server)

		if o := r.Header.Get("Origin"); len(o) > 0 {
			w.Header().Set("Access-Control-Allow-Origin", "*")
			w.Header().Set("Access-Control-Allow-Methods", "GET, POST, HEAD, PUT, DELETE, OPTIONS")
			w.Header().Set("Access-Control-Expose-Headers", "Server,range,Content-Length,Content-Range")
			w.Header().Set("Access-Control-Allow-Headers", "origin,range,accept-encoding,referer,Cache-Control,X-Proxy-Authorization,X-Requested-With,Content-Type")
		}

		if proxyUrls == nil {
			fs.ServeHTTP(w, r)
			return
		}

		for _, proxyUrl := range proxyUrls {
			srcPath, proxyPath := r.URL.Path, proxyUrl.Path
			if !strings.HasSuffix(srcPath, "/") {
				// /api to /api/
				// /api.js to /api.js/
				// /api/100 to /api/100/
				srcPath += "/"
			}
			if !strings.HasSuffix(proxyPath, "/") {
				// /api/ to /api/
				// to match /api/ or /api/100
				// and not match /api.js/
				proxyPath += "/"
			}
			if !strings.HasPrefix(srcPath, proxyPath) {
				continue
			}

			// For matched OPTIONS, directly return without response.
			if r.Method == "OPTIONS" {
				return
			}

			if proxy, ok := proxies[proxyUrl.Path]; ok {
				proxy.ServeHTTP(w, r)
				return
			}
		}

		fs.ServeHTTP(w, r)
	})

	var protos []string
	if httpPort != 0 {
		protos = append(protos, fmt.Sprintf("http(:%v)", httpPort))
	}
	if httpsPort != 0 {
		s := httpsDomains
		if httpsDomains == "" {
			s = "all domains"
		}

		if useLetsEncrypt {
			protos = append(protos, fmt.Sprintf("https(:%v, %v, %v)", httpsPort, s, cacheFile))
		} else {
			protos = append(protos, fmt.Sprintf("https(:%v)", httpsPort))
		}

		if useLetsEncrypt {
			protos = append(protos, "letsencrypt")
		} else {
			protos = append(protos, fmt.Sprintf("self-sign(%v, %v)", ssKey, ssCert))
		}
	}
	ol.Tf(ctx, "%v html root at %v", strings.Join(protos, ", "), string(html))

	var hs, hss *http.Server

	wg := sync.WaitGroup{}
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	var err error
	go func() {
		defer wg.Done()
		if httpPort == 0 {
			ol.W(ctx, "http server disabled")
			return
		}

		defer cancel()
		hs = &http.Server{Addr: fmt.Sprintf(":%v", httpPort), Handler: nil}
		ol.Tf(ctx, "http serve at %v", httpPort)

		if err = hs.ListenAndServe(); err != nil {
			err = oe.Wrapf(err, "serve http")
			return
		}
		ol.T("http server ok")
	}()
	wg.Add(1)

	go func() {
		defer wg.Done()
		if httpsPort == 0 {
			ol.W(ctx, "https server disabled")
			return
		}

		defer cancel()

		var m https.Manager
		if useLetsEncrypt {
			var domains []string
			if httpsDomains != "" {
				domains = strings.Split(httpsDomains, ",")
			}

			if m, err = https.NewLetsencryptManager("", domains, cacheFile); err != nil {
				err = oe.Wrapf(err, "create letsencrypt manager")
				return
			}
		} else {
			if m, err = https.NewSelfSignManager(ssCert, ssKey); err != nil {
				err = oe.Wrapf(err, "create self-sign manager")
				return
			}
		}

		hss = &http.Server{
			Addr: fmt.Sprintf(":%v", httpsPort),
			TLSConfig: &tls.Config{
				GetCertificate: m.GetCertificate,
			},
		}
		ol.Tf(ctx, "http serve at %v", httpsPort)

		if err = hss.ListenAndServeTLS("", ""); err != nil {
			err = oe.Wrapf(err, "listen and serve https")
			return
		}
		ol.T("https serve ok")
	}()
	wg.Add(1)

	select {
	case <-ctx.Done():
		if hs != nil {
			hs.Close()
		}
		if hss != nil {
			hss.Close()
		}
	}
	wg.Wait()

	return err
}

func main() {
	ctx := context.Background()
	if err := run(ctx); err != nil {
		ol.Ef(ctx, "run err %+v", err)
		os.Exit(-1)
	}
}
<|endoftext|>


<commit_before>package util

import (
	"os"
)

// InTempDir is like WithTempDir, but also cd into the directory before running
// the function, and cd backs after running the function if possible.
//
// It panics if it could not get the working directory or change directory.
//
// It is useful in tests.
func InTempDir(f func(string)) {
	tmpdir, cleanup := TestDir()
	defer cleanup()

	oldpwd, err := os.Getwd()
	if err != nil {
		panic(err)
	}

	mustChdir(tmpdir)
	defer mustChdir(oldpwd)

	f(tmpdir)
}
<commit_msg>util: Remove unused function InTempDir.<commit_after><|endoftext|>


<commit_before>// Copyright (c) 2015 Mattermost, Inc. All Rights Reserved.
// See License.txt for license information.

package utils

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"

	l4g "code.google.com/p/log4go"

	"github.com/mattermost/platform/model"
)

const (
	MODE_DEV        = "dev"
	MODE_BETA       = "beta"
	MODE_PROD       = "prod"
	LOG_ROTATE_SIZE = 10000
)

var Cfg *model.Config = &model.Config{}
var CfgLastModified int64 = 0
var CfgFileName string = ""
var ClientCfg map[string]string = map[string]string{}
var SanitizeOptions map[string]bool = map[string]bool{}

func FindConfigFile(fileName string) string {
	if _, err := os.Stat("/tmp/" + fileName); err == nil {
		fileName, _ = filepath.Abs("/tmp/" + fileName)
	} else if _, err := os.Stat("./config/" + fileName); err == nil {
		fileName, _ = filepath.Abs("./config/" + fileName)
	} else if _, err := os.Stat("../config/" + fileName); err == nil {
		fileName, _ = filepath.Abs("../config/" + fileName)
	} else if _, err := os.Stat(fileName); err == nil {
		fileName, _ = filepath.Abs(fileName)
	}

	return fileName
}

func FindDir(dir string) string {
	fileName := "."
	if _, err := os.Stat("./" + dir + "/"); err == nil {
		fileName, _ = filepath.Abs("./" + dir + "/")
	} else if _, err := os.Stat("../" + dir + "/"); err == nil {
		fileName, _ = filepath.Abs("../" + dir + "/")
	} else if _, err := os.Stat("/tmp/" + dir); err == nil {
		fileName, _ = filepath.Abs("/tmp/" + dir)
	}

	return fileName + "/"
}

func ConfigureCmdLineLog() {
	ls := model.LogSettings{}
	ls.EnableConsole = true
	ls.ConsoleLevel = "ERROR"
	ls.EnableFile = false
	configureLog(&ls)
}

func configureLog(s *model.LogSettings) {

	l4g.Close()

	if s.EnableConsole {
		level := l4g.DEBUG
		if s.ConsoleLevel == "INFO" {
			level = l4g.INFO
		} else if s.ConsoleLevel == "ERROR" {
			level = l4g.ERROR
		}

		l4g.AddFilter("stdout", level, l4g.NewConsoleLogWriter())
	}

	if s.EnableFile {

		var fileFormat = s.FileFormat

		if fileFormat == "" {
			fileFormat = "[%D %T] [%L] %M"
		}

		level := l4g.DEBUG
		if s.FileLevel == "INFO" {
			level = l4g.INFO
		} else if s.FileLevel == "ERROR" {
			level = l4g.ERROR
		}

		flw := l4g.NewFileLogWriter(GetLogFileLocation(s.FileLocation), false)
		flw.SetFormat(fileFormat)
		flw.SetRotate(true)
		flw.SetRotateLines(LOG_ROTATE_SIZE)
		l4g.AddFilter("file", level, flw)
	}
}

func GetLogFileLocation(fileLocation string) string {
	if fileLocation == "" {
		return FindDir("logs") + "mattermost.log"
	} else {
		return fileLocation
	}
}

func SaveConfig(fileName string, config *model.Config) *model.AppError {
	b, err := json.MarshalIndent(config, "", "    ")
	if err != nil {
		return model.NewAppError("SaveConfig", "An error occurred while saving the file to "+fileName, err.Error())
	}

	err = ioutil.WriteFile(fileName, b, 0644)
	if err != nil {
		return model.NewAppError("SaveConfig", "An error occurred while saving the file to "+fileName, err.Error())
	}

	return nil
}

// LoadConfig will try to search around for the corresponding config file.
// It will search /tmp/fileName then attempt ./config/fileName,
// then ../config/fileName and last it will look at fileName
func LoadConfig(fileName string) {

	fileName = FindConfigFile(fileName)

	file, err := os.Open(fileName)
	if err != nil {
		panic("Error opening config file=" + fileName + ", err=" + err.Error())
	}

	decoder := json.NewDecoder(file)
	config := model.Config{}
	err = decoder.Decode(&config)
	if err != nil {
		panic("Error decoding config file=" + fileName + ", err=" + err.Error())
	}

	if info, err := file.Stat(); err != nil {
		panic("Error getting config info file=" + fileName + ", err=" + err.Error())
	} else {
		CfgLastModified = info.ModTime().Unix()
		CfgFileName = fileName
	}

	config.SetDefaults()

	if err := config.IsValid(); err != nil {
		panic("Error validating config file=" + fileName + ", err=" + err.Message)
	}

	configureLog(&config.LogSettings)
	TestConnection(&config)

	Cfg = &config
	SanitizeOptions = getSanitizeOptions(Cfg)
	ClientCfg = getClientConfig(Cfg)
}

func getSanitizeOptions(c *model.Config) map[string]bool {
	options := map[string]bool{}
	options["fullname"] = c.PrivacySettings.ShowFullName
	options["email"] = c.PrivacySettings.ShowEmailAddress

	return options
}

func getClientConfig(c *model.Config) map[string]string {
	props := make(map[string]string)

	props["Version"] = model.CurrentVersion
	props["BuildNumber"] = model.BuildNumber
	props["BuildDate"] = model.BuildDate
	props["BuildHash"] = model.BuildHash

	props["SiteName"] = c.TeamSettings.SiteName
	props["EnableTeamCreation"] = strconv.FormatBool(c.TeamSettings.EnableTeamCreation)
	props["RestrictTeamNames"] = strconv.FormatBool(*c.TeamSettings.RestrictTeamNames)

	props["EnableOAuthServiceProvider"] = strconv.FormatBool(c.ServiceSettings.EnableOAuthServiceProvider)

	props["SegmentDeveloperKey"] = c.ServiceSettings.SegmentDeveloperKey
	props["GoogleDeveloperKey"] = c.ServiceSettings.GoogleDeveloperKey
	props["EnableIncomingWebhooks"] = strconv.FormatBool(c.ServiceSettings.EnableIncomingWebhooks)
	props["EnableOutgoingWebhooks"] = strconv.FormatBool(c.ServiceSettings.EnableOutgoingWebhooks)
	props["EnablePostUsernameOverride"] = strconv.FormatBool(c.ServiceSettings.EnablePostUsernameOverride)
	props["EnablePostIconOverride"] = strconv.FormatBool(c.ServiceSettings.EnablePostIconOverride)

	props["SendEmailNotifications"] = strconv.FormatBool(c.EmailSettings.SendEmailNotifications)
	props["EnableSignUpWithEmail"] = strconv.FormatBool(c.EmailSettings.EnableSignUpWithEmail)
	props["RequireEmailVerification"] = strconv.FormatBool(c.EmailSettings.RequireEmailVerification)
	props["FeedbackEmail"] = c.EmailSettings.FeedbackEmail

	props["EnableSignUpWithGitLab"] = strconv.FormatBool(c.GitLabSettings.Enable)

	props["ShowEmailAddress"] = strconv.FormatBool(c.PrivacySettings.ShowEmailAddress)

	props["EnablePublicLink"] = strconv.FormatBool(c.FileSettings.EnablePublicLink)
	props["ProfileHeight"] = fmt.Sprintf("%v", c.FileSettings.ProfileHeight)
	props["ProfileWidth"] = fmt.Sprintf("%v", c.FileSettings.ProfileWidth)

	return props
}
<commit_msg>adds trailing slash if missing from Directory in config.json<commit_after>// Copyright (c) 2015 Mattermost, Inc. All Rights Reserved.
// See License.txt for license information.

package utils

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"

	l4g "code.google.com/p/log4go"

	"github.com/mattermost/platform/model"
)

const (
	MODE_DEV        = "dev"
	MODE_BETA       = "beta"
	MODE_PROD       = "prod"
	LOG_ROTATE_SIZE = 10000
)

var Cfg *model.Config = &model.Config{}
var CfgLastModified int64 = 0
var CfgFileName string = ""
var ClientCfg map[string]string = map[string]string{}
var SanitizeOptions map[string]bool = map[string]bool{}

func FindConfigFile(fileName string) string {
	if _, err := os.Stat("/tmp/" + fileName); err == nil {
		fileName, _ = filepath.Abs("/tmp/" + fileName)
	} else if _, err := os.Stat("./config/" + fileName); err == nil {
		fileName, _ = filepath.Abs("./config/" + fileName)
	} else if _, err := os.Stat("../config/" + fileName); err == nil {
		fileName, _ = filepath.Abs("../config/" + fileName)
	} else if _, err := os.Stat(fileName); err == nil {
		fileName, _ = filepath.Abs(fileName)
	}

	return fileName
}

func FindDir(dir string) string {
	fileName := "."
	if _, err := os.Stat("./" + dir + "/"); err == nil {
		fileName, _ = filepath.Abs("./" + dir + "/")
	} else if _, err := os.Stat("../" + dir + "/"); err == nil {
		fileName, _ = filepath.Abs("../" + dir + "/")
	} else if _, err := os.Stat("/tmp/" + dir); err == nil {
		fileName, _ = filepath.Abs("/tmp/" + dir)
	}

	return fileName + "/"
}

func ConfigureCmdLineLog() {
	ls := model.LogSettings{}
	ls.EnableConsole = true
	ls.ConsoleLevel = "ERROR"
	ls.EnableFile = false
	configureLog(&ls)
}

func configureLog(s *model.LogSettings) {

	l4g.Close()

	if s.EnableConsole {
		level := l4g.DEBUG
		if s.ConsoleLevel == "INFO" {
			level = l4g.INFO
		} else if s.ConsoleLevel == "ERROR" {
			level = l4g.ERROR
		}

		l4g.AddFilter("stdout", level, l4g.NewConsoleLogWriter())
	}

	if s.EnableFile {

		var fileFormat = s.FileFormat

		if fileFormat == "" {
			fileFormat = "[%D %T] [%L] %M"
		}

		level := l4g.DEBUG
		if s.FileLevel == "INFO" {
			level = l4g.INFO
		} else if s.FileLevel == "ERROR" {
			level = l4g.ERROR
		}

		flw := l4g.NewFileLogWriter(GetLogFileLocation(s.FileLocation), false)
		flw.SetFormat(fileFormat)
		flw.SetRotate(true)
		flw.SetRotateLines(LOG_ROTATE_SIZE)
		l4g.AddFilter("file", level, flw)
	}
}

func GetLogFileLocation(fileLocation string) string {
	if fileLocation == "" {
		return FindDir("logs") + "mattermost.log"
	} else {
		return fileLocation
	}
}

func SaveConfig(fileName string, config *model.Config) *model.AppError {
	b, err := json.MarshalIndent(config, "", "    ")
	if err != nil {
		return model.NewAppError("SaveConfig", "An error occurred while saving the file to "+fileName, err.Error())
	}

	err = ioutil.WriteFile(fileName, b, 0644)
	if err != nil {
		return model.NewAppError("SaveConfig", "An error occurred while saving the file to "+fileName, err.Error())
	}

	return nil
}

// LoadConfig will try to search around for the corresponding config file.
// It will search /tmp/fileName then attempt ./config/fileName,
// then ../config/fileName and last it will look at fileName
func LoadConfig(fileName string) {

	fileName = FindConfigFile(fileName)

	file, err := os.Open(fileName)
	if err != nil {
		panic("Error opening config file=" + fileName + ", err=" + err.Error())
	}

	decoder := json.NewDecoder(file)
	config := model.Config{}
	err = decoder.Decode(&config)
	if err != nil {
		panic("Error decoding config file=" + fileName + ", err=" + err.Error())
	}

	if info, err := file.Stat(); err != nil {
		panic("Error getting config info file=" + fileName + ", err=" + err.Error())
	} else {
		CfgLastModified = info.ModTime().Unix()
		CfgFileName = fileName
	}

	config.SetDefaults()

	if err := config.IsValid(); err != nil {
		panic("Error validating config file=" + fileName + ", err=" + err.Message)
	}

	configureLog(&config.LogSettings)
	TestConnection(&config)

	if config.FileSettings.DriverName == model.IMAGE_DRIVER_LOCAL {
		dir := config.FileSettings.Directory
		if len(dir) > 0 && dir[len(dir)-1:] != "/" {
			config.FileSettings.Directory += "/"
		}
	}

	Cfg = &config
	SanitizeOptions = getSanitizeOptions(Cfg)
	ClientCfg = getClientConfig(Cfg)
}

func getSanitizeOptions(c *model.Config) map[string]bool {
	options := map[string]bool{}
	options["fullname"] = c.PrivacySettings.ShowFullName
	options["email"] = c.PrivacySettings.ShowEmailAddress

	return options
}

func getClientConfig(c *model.Config) map[string]string {
	props := make(map[string]string)

	props["Version"] = model.CurrentVersion
	props["BuildNumber"] = model.BuildNumber
	props["BuildDate"] = model.BuildDate
	props["BuildHash"] = model.BuildHash

	props["SiteName"] = c.TeamSettings.SiteName
	props["EnableTeamCreation"] = strconv.FormatBool(c.TeamSettings.EnableTeamCreation)
	props["RestrictTeamNames"] = strconv.FormatBool(*c.TeamSettings.RestrictTeamNames)

	props["EnableOAuthServiceProvider"] = strconv.FormatBool(c.ServiceSettings.EnableOAuthServiceProvider)

	props["SegmentDeveloperKey"] = c.ServiceSettings.SegmentDeveloperKey
	props["GoogleDeveloperKey"] = c.ServiceSettings.GoogleDeveloperKey
	props["EnableIncomingWebhooks"] = strconv.FormatBool(c.ServiceSettings.EnableIncomingWebhooks)
	props["EnableOutgoingWebhooks"] = strconv.FormatBool(c.ServiceSettings.EnableOutgoingWebhooks)
	props["EnablePostUsernameOverride"] = strconv.FormatBool(c.ServiceSettings.EnablePostUsernameOverride)
	props["EnablePostIconOverride"] = strconv.FormatBool(c.ServiceSettings.EnablePostIconOverride)

	props["SendEmailNotifications"] = strconv.FormatBool(c.EmailSettings.SendEmailNotifications)
	props["EnableSignUpWithEmail"] = strconv.FormatBool(c.EmailSettings.EnableSignUpWithEmail)
	props["RequireEmailVerification"] = strconv.FormatBool(c.EmailSettings.RequireEmailVerification)
	props["FeedbackEmail"] = c.EmailSettings.FeedbackEmail

	props["EnableSignUpWithGitLab"] = strconv.FormatBool(c.GitLabSettings.Enable)

	props["ShowEmailAddress"] = strconv.FormatBool(c.PrivacySettings.ShowEmailAddress)

	props["EnablePublicLink"] = strconv.FormatBool(c.FileSettings.EnablePublicLink)
	props["ProfileHeight"] = fmt.Sprintf("%v", c.FileSettings.ProfileHeight)
	props["ProfileWidth"] = fmt.Sprintf("%v", c.FileSettings.ProfileWidth)

	return props
}
<|endoftext|>


<commit_before>/*Package seq balabala

This package defines a *Seq* type, and provides some basic operations of sequence,
like validation of DNA/RNA/Protein sequence and getting reverse complement sequence.

This package was inspired by
[biogo](https://code.google.com/p/biogo/source/browse/#git%2Falphabet).

IUPAC nucleotide code: ACGTURYSWKMBDHVN

http://droog.gs.washington.edu/parc/images/iupac.html

	code	base	Complement
	A	A	T
	C	C	G
	G	G	C
	T/U	T	A

	R	A/G	Y
	Y	C/T	R
	S	C/G	S
	W	A/T	W
	K	G/T	M
	M	A/C	K

	B	C/G/T	V
	D	A/G/T	H
	H	A/C/T	D
	V	A/C/G	B

	X/N	A/C/G/T	X
	.	not A/C/G/T
	 or-	gap

IUPAC amino acid code: `ACGTRYSWKMBDHV`

	A	Ala	Alanine
	B	Asx	Aspartic acid or Asparagine [2]
	C	Cys	Cysteine
	D	Asp	Aspartic Acid
	E	Glu	Glutamic Acid
	F	Phe	Phenylalanine
	G	Gly	Glycine
	H	His	Histidine
	I	Ile	Isoleucine
	J		Isoleucine or Leucine [4]
	K	Lys	Lysine
	L	Leu	Leucine
	M	Met	Methionine
	N	Asn	Asparagine
	P	Pro	Proline
	Q	Gln	Glutamine
	R	Arg	Arginine
	S	Ser	Serine
	T	Thr	Threonine
	V	Val	Valine
	W	Trp	Tryptophan
	Y	Tyr	Tyrosine
	Z	Glx	Glutamine or Glutamic acid [2]

Other links:

	1. http://www.bioinformatics.org/sms/iupac.html
	2. http://www.dnabaser.com/articles/IUPAC%20ambiguity%20codes.html
	3. http://www.bioinformatics.org/sms2/iupac.html
	4. http://www.matrixscience.com/blog/non-standard-amino-acid-residues.html

*/
package seq

import (
	"errors"
	"fmt"
	"runtime"
	"sync"

	"github.com/shenwei356/util/byteutil"
)

/*Alphabet could be defined. Attention that,
**the letters are case sensitive**.

For exmaple, DNA:

	DNA, _ = NewAlphabet(
		"DNA",
		[]byte("acgtACGT"),
		[]byte("tgcaTGCA"),
		[]byte(" -"),
		[]byte("nN"))

*/
type Alphabet struct {
	t         string
	isUnlimit bool

	letters   []byte
	pairs     []byte
	gap       []byte
	ambiguous []byte

	allLetters []byte

	pairLetters []byte
}

// NewAlphabet is Constructor for type *Alphabet*
func NewAlphabet(
	t string,
	isUnlimit bool,
	letters []byte,
	pairs []byte,
	gap []byte,
	ambiguous []byte,
) (*Alphabet, error) {

	a := &Alphabet{t, isUnlimit, letters, pairs, gap, ambiguous, []byte{}, []byte{}}

	if isUnlimit {
		return a, nil
	}

	if len(letters) != len(pairs) {
		return a, errors.New("mismatch of length of letters and pairs")
	}

	for i := 0; i < len(letters); i++ {
		a.allLetters = append(a.allLetters, letters[i])
	}
	// add gap and ambiguous code
	for _, v := range gap {
		a.allLetters = append(a.allLetters, v)
	}
	for _, v := range ambiguous {
		a.allLetters = append(a.allLetters, v)
	}

	// construct special slice.
	// index are the integer of a byte, and value is the original byte.
	// it's faster than map!!!!
	max := -1
	for i := 0; i < len(a.allLetters); i++ {
		b := int(a.allLetters[i])
		if max < b {
			max = b
		}
	}
	a.pairLetters = make([]byte, max+1)
	for i := 0; i < len(letters); i++ {
		a.pairLetters[int(letters[i])] = pairs[i]
	}
	for _, v := range gap {
		a.pairLetters[int(v)] = v
	}
	for _, v := range ambiguous {
		a.pairLetters[int(v)] = v
	}

	return a, nil
}

// Type returns type of the alphabet
func (a *Alphabet) Type() string {
	return a.t
}

// Letters returns letters
func (a *Alphabet) Letters() []byte {
	return a.letters
}

// Gaps returns gaps
func (a *Alphabet) Gaps() []byte {
	return a.gap
}

// AmbiguousLetters returns AmbiguousLetters
func (a *Alphabet) AmbiguousLetters() []byte {
	return a.ambiguous
}

// AllLetters return all letters
func (a *Alphabet) AllLetters() []byte {
	return a.allLetters
}

// String returns type of the alphabet
func (a *Alphabet) String() string {
	return a.t
}

// IsValidLetter is used to validate a letter
func (a *Alphabet) IsValidLetter(b byte) bool {
	if a.isUnlimit {
		return true
	}
	return a.pairLetters[int(b)] != 0
}

// ValidSeqLengthThreshold is the threshold of sequence length that
// needed to  parallelly checking sequence
var ValidSeqLengthThreshold = 100000

// ValidSeqThreads is the threads number of parallelly checking sequence
var ValidSeqThreads = runtime.NumCPU()

type seqCheckStatus struct {
	err error
}

// IsValid is used to validate a byte slice
func (a *Alphabet) IsValid(s []byte) error {
	if len(s) == 0 {
		return nil
	}
	if a == nil || a.isUnlimit {
		return nil
	}

	l := len(s)
	if l < ValidSeqLengthThreshold {
		for _, b := range s {
			if !a.IsValidLetter(b) {
				return fmt.Errorf("invalid %s lebtter: %s", a, []byte{b})
			}
		}
	} else {
		chunkSize, start, end := int(l/ValidSeqThreads), 0, 0

		var wg sync.WaitGroup
		tokens := make(chan int, ValidSeqThreads)
		ch := make(chan seqCheckStatus, ValidSeqThreads)
		done := make(chan struct{})
		finished := false
		for i := 0; i < ValidSeqThreads; i++ {
			start = i * chunkSize
			end = (i + 1) * chunkSize
			if end > l {
				end = l
			}
			tokens <- 1
			wg.Add(1)
			go func(start, end int) {
				defer func() {
					<-tokens
					wg.Done()
				}()

				select {
				case <-done:
					if !finished {
						finished = true
						close(ch)
						return
					}
				default:

				}

				for i := start; i < end; i++ {
					if !a.IsValidLetter(s[i]) {
						ch <- seqCheckStatus{fmt.Errorf("invalid %s lebtter: %s at %d", a, []byte{s[i]}, i)}
						close(done)
						return
					}
				}
				ch <- seqCheckStatus{nil}
			}(start, end)
		}
		wg.Wait()
		close(ch)
		for status := range ch {
			if status.err != nil {
				return status.err
			}
		}
	}
	return nil
}

// PairLetter return the Pair Letter
func (a *Alphabet) PairLetter(b byte) (byte, error) {
	if a.isUnlimit {
		return b, nil
	}

	p := a.pairLetters[int(b)]
	if p == 0 {
		return b, fmt.Errorf("invalid letter: %c", b)
	}
	return p, nil
}

/*Four types of alphabets are pre-defined:

  DNA           Deoxyribonucleotide code
  DNAredundant  DNA + Ambiguity Codes
  RNA           Oxyribonucleotide code
  RNAredundant  RNA + Ambiguity Codes
  Protein       Amino Acide single-letter Code
  Unlimit       Self-defined, including all 26 English letters

*/
var (
	DNA          *Alphabet
	DNAredundant *Alphabet
	RNA          *Alphabet
	RNAredundant *Alphabet
	Protein      *Alphabet
	Unlimit      *Alphabet

	abProtein      map[byte]bool
	abDNAredundant map[byte]bool
	abDNA          map[byte]bool
	abRNAredundant map[byte]bool
	abRNA          map[byte]bool
)

func init() {
	DNA, _ = NewAlphabet(
		"DNA",
		false,
		[]byte("acgtACGT."),
		[]byte("tgcaTGCA."),
		[]byte(" -"),
		[]byte("nxNX"))

	DNAredundant, _ = NewAlphabet(
		"DNAredundant",
		false,
		[]byte("acgtryswkmbdhvACGTRYSWKMBDHV."),
		[]byte("tgcayrswmkvhdbTGCAYRSWMKVHDB."),
		[]byte(" -"),
		[]byte("nxNX"))

	RNA, _ = NewAlphabet(
		"RNA",
		false,
		[]byte("acguACGU"),
		[]byte("ugcaUGCA"),
		[]byte(" -"),
		[]byte("nxNX"))

	RNAredundant, _ = NewAlphabet(
		"RNAredundant",
		false,
		[]byte("acguryswkmbdhvACGURYSWKMBDHV."),
		[]byte("ugcayrswmkvhdbUGCAYRSWMKVHDB."),
		[]byte(" -"),
		[]byte("nxNX"))

	Protein, _ = NewAlphabet(
		"Protein",
		false,
		[]byte("abcdefghijklmnpqrstvwyzABCDEFGHIJKLMNPQRSTVWYZ*_."),
		[]byte("abcdefghijklmnpqrstvwyzABCDEFGHIJKLMNPQRSTVWYZ*_."),
		[]byte(" -"),
		[]byte("xX"))

	Unlimit, _ = NewAlphabet(
		"Unlimit",
		true,
		nil,
		nil,
		nil,
		nil)

	abProtein = slice2map(byteutil.Alphabet(Protein.AllLetters()))
	abDNAredundant = slice2map(byteutil.Alphabet(DNAredundant.AllLetters()))
	abDNA = slice2map(byteutil.Alphabet(DNA.AllLetters()))
	abRNAredundant = slice2map(byteutil.Alphabet(RNAredundant.AllLetters()))
	abRNA = slice2map(byteutil.Alphabet(RNA.AllLetters()))
}

// AlphabetGuessSeqLenghtThreshold is the length of sequence prefix of the first FASTA record
// based which FastaRecord guesses the sequence type
var AlphabetGuessSeqLenghtThreshold = 10000

// GuessAlphabet guesses alphabet by given
func GuessAlphabet(seqs []byte) *Alphabet {
	var alphabetMap map[byte]bool
	if len(seqs) > AlphabetGuessSeqLenghtThreshold { // reduce guessing time
		alphabetMap = slice2map(byteutil.Alphabet(seqs[0:AlphabetGuessSeqLenghtThreshold]))
	} else {
		alphabetMap = slice2map(byteutil.Alphabet(seqs))
	}
	if isSubset(alphabetMap, abDNA) {
		return DNA
	}
	if isSubset(alphabetMap, abRNA) {
		return RNA
	}
	if isSubset(alphabetMap, abDNAredundant) {
		return DNAredundant
	}
	if isSubset(alphabetMap, abRNAredundant) {
		return RNAredundant
	}
	if isSubset(alphabetMap, abProtein) {
		return Protein
	}

	return Unlimit
}

// GuessAlphabetLessConservatively change DNA to DNAredundant and RNA to RNAredundant
func GuessAlphabetLessConservatively(seqs []byte) *Alphabet {
	ab := GuessAlphabet(seqs)
	if ab == DNA {
		return DNAredundant
	}
	if ab == RNA {
		return RNAredundant
	}
	return ab
}

func isSubset(query, subject map[byte]bool) bool {
	for b := range query {
		if _, ok := subject[b]; !ok {
			return false
		}
	}
	return true
}

func slice2map(s []byte) map[byte]bool {
	m := make(map[byte]bool)
	for _, b := range s {
		m[b] = true
	}
	return m
}
<commit_msg>fix typo<commit_after>/*Package seq balabala

This package defines a *Seq* type, and provides some basic operations of sequence,
like validation of DNA/RNA/Protein sequence and getting reverse complement sequence.

This package was inspired by
[biogo](https://code.google.com/p/biogo/source/browse/#git%2Falphabet).

IUPAC nucleotide code: ACGTURYSWKMBDHVN

http://droog.gs.washington.edu/parc/images/iupac.html

	code	base	Complement
	A	A	T
	C	C	G
	G	G	C
	T/U	T	A

	R	A/G	Y
	Y	C/T	R
	S	C/G	S
	W	A/T	W
	K	G/T	M
	M	A/C	K

	B	C/G/T	V
	D	A/G/T	H
	H	A/C/T	D
	V	A/C/G	B

	X/N	A/C/G/T	X
	.	not A/C/G/T
	 or-	gap

IUPAC amino acid code: `ACGTRYSWKMBDHV`

	A	Ala	Alanine
	B	Asx	Aspartic acid or Asparagine [2]
	C	Cys	Cysteine
	D	Asp	Aspartic Acid
	E	Glu	Glutamic Acid
	F	Phe	Phenylalanine
	G	Gly	Glycine
	H	His	Histidine
	I	Ile	Isoleucine
	J		Isoleucine or Leucine [4]
	K	Lys	Lysine
	L	Leu	Leucine
	M	Met	Methionine
	N	Asn	Asparagine
	P	Pro	Proline
	Q	Gln	Glutamine
	R	Arg	Arginine
	S	Ser	Serine
	T	Thr	Threonine
	V	Val	Valine
	W	Trp	Tryptophan
	Y	Tyr	Tyrosine
	Z	Glx	Glutamine or Glutamic acid [2]

Other links:

	1. http://www.bioinformatics.org/sms/iupac.html
	2. http://www.dnabaser.com/articles/IUPAC%20ambiguity%20codes.html
	3. http://www.bioinformatics.org/sms2/iupac.html
	4. http://www.matrixscience.com/blog/non-standard-amino-acid-residues.html

*/
package seq

import (
	"errors"
	"fmt"
	"runtime"
	"sync"

	"github.com/shenwei356/util/byteutil"
)

/*Alphabet could be defined. Attention that,
**the letters are case sensitive**.

For exmaple, DNA:

	DNA, _ = NewAlphabet(
		"DNA",
		[]byte("acgtACGT"),
		[]byte("tgcaTGCA"),
		[]byte(" -"),
		[]byte("nN"))

*/
type Alphabet struct {
	t         string
	isUnlimit bool

	letters   []byte
	pairs     []byte
	gap       []byte
	ambiguous []byte

	allLetters []byte

	pairLetters []byte
}

// NewAlphabet is Constructor for type *Alphabet*
func NewAlphabet(
	t string,
	isUnlimit bool,
	letters []byte,
	pairs []byte,
	gap []byte,
	ambiguous []byte,
) (*Alphabet, error) {

	a := &Alphabet{t, isUnlimit, letters, pairs, gap, ambiguous, []byte{}, []byte{}}

	if isUnlimit {
		return a, nil
	}

	if len(letters) != len(pairs) {
		return a, errors.New("mismatch of length of letters and pairs")
	}

	for i := 0; i < len(letters); i++ {
		a.allLetters = append(a.allLetters, letters[i])
	}
	// add gap and ambiguous code
	for _, v := range gap {
		a.allLetters = append(a.allLetters, v)
	}
	for _, v := range ambiguous {
		a.allLetters = append(a.allLetters, v)
	}

	// construct special slice.
	// index are the integer of a byte, and value is the original byte.
	// it's faster than map!!!!
	max := -1
	for i := 0; i < len(a.allLetters); i++ {
		b := int(a.allLetters[i])
		if max < b {
			max = b
		}
	}
	a.pairLetters = make([]byte, max+1)
	for i := 0; i < len(letters); i++ {
		a.pairLetters[int(letters[i])] = pairs[i]
	}
	for _, v := range gap {
		a.pairLetters[int(v)] = v
	}
	for _, v := range ambiguous {
		a.pairLetters[int(v)] = v
	}

	return a, nil
}

// Type returns type of the alphabet
func (a *Alphabet) Type() string {
	return a.t
}

// Letters returns letters
func (a *Alphabet) Letters() []byte {
	return a.letters
}

// Gaps returns gaps
func (a *Alphabet) Gaps() []byte {
	return a.gap
}

// AmbiguousLetters returns AmbiguousLetters
func (a *Alphabet) AmbiguousLetters() []byte {
	return a.ambiguous
}

// AllLetters return all letters
func (a *Alphabet) AllLetters() []byte {
	return a.allLetters
}

// String returns type of the alphabet
func (a *Alphabet) String() string {
	return a.t
}

// IsValidLetter is used to validate a letter
func (a *Alphabet) IsValidLetter(b byte) bool {
	if a.isUnlimit {
		return true
	}
	return a.pairLetters[int(b)] != 0
}

// ValidSeqLengthThreshold is the threshold of sequence length that
// needed to  parallelly checking sequence
var ValidSeqLengthThreshold = 100000

// ValidSeqThreads is the threads number of parallelly checking sequence
var ValidSeqThreads = runtime.NumCPU()

type seqCheckStatus struct {
	err error
}

// IsValid is used to validate a byte slice
func (a *Alphabet) IsValid(s []byte) error {
	if len(s) == 0 {
		return nil
	}
	if a == nil || a.isUnlimit {
		return nil
	}

	l := len(s)
	if l < ValidSeqLengthThreshold {
		for _, b := range s {
			if !a.IsValidLetter(b) {
				return fmt.Errorf("invalid %s letter: %s", a, []byte{b})
			}
		}
	} else {
		chunkSize, start, end := int(l/ValidSeqThreads), 0, 0

		var wg sync.WaitGroup
		tokens := make(chan int, ValidSeqThreads)
		ch := make(chan seqCheckStatus, ValidSeqThreads)
		done := make(chan struct{})
		finished := false
		for i := 0; i < ValidSeqThreads; i++ {
			start = i * chunkSize
			end = (i + 1) * chunkSize
			if end > l {
				end = l
			}
			tokens <- 1
			wg.Add(1)
			go func(start, end int) {
				defer func() {
					<-tokens
					wg.Done()
				}()

				select {
				case <-done:
					if !finished {
						finished = true
						close(ch)
						return
					}
				default:

				}

				for i := start; i < end; i++ {
					if !a.IsValidLetter(s[i]) {
						ch <- seqCheckStatus{fmt.Errorf("invalid %s lebtter: %s at %d", a, []byte{s[i]}, i)}
						close(done)
						return
					}
				}
				ch <- seqCheckStatus{nil}
			}(start, end)
		}
		wg.Wait()
		close(ch)
		for status := range ch {
			if status.err != nil {
				return status.err
			}
		}
	}
	return nil
}

// PairLetter return the Pair Letter
func (a *Alphabet) PairLetter(b byte) (byte, error) {
	if a.isUnlimit {
		return b, nil
	}

	p := a.pairLetters[int(b)]
	if p == 0 {
		return b, fmt.Errorf("invalid letter: %c", b)
	}
	return p, nil
}

/*Four types of alphabets are pre-defined:

  DNA           Deoxyribonucleotide code
  DNAredundant  DNA + Ambiguity Codes
  RNA           Oxyribonucleotide code
  RNAredundant  RNA + Ambiguity Codes
  Protein       Amino Acide single-letter Code
  Unlimit       Self-defined, including all 26 English letters

*/
var (
	DNA          *Alphabet
	DNAredundant *Alphabet
	RNA          *Alphabet
	RNAredundant *Alphabet
	Protein      *Alphabet
	Unlimit      *Alphabet

	abProtein      map[byte]bool
	abDNAredundant map[byte]bool
	abDNA          map[byte]bool
	abRNAredundant map[byte]bool
	abRNA          map[byte]bool
)

func init() {
	DNA, _ = NewAlphabet(
		"DNA",
		false,
		[]byte("acgtACGT."),
		[]byte("tgcaTGCA."),
		[]byte(" -"),
		[]byte("nxNX"))

	DNAredundant, _ = NewAlphabet(
		"DNAredundant",
		false,
		[]byte("acgtryswkmbdhvACGTRYSWKMBDHV."),
		[]byte("tgcayrswmkvhdbTGCAYRSWMKVHDB."),
		[]byte(" -"),
		[]byte("nxNX"))

	RNA, _ = NewAlphabet(
		"RNA",
		false,
		[]byte("acguACGU"),
		[]byte("ugcaUGCA"),
		[]byte(" -"),
		[]byte("nxNX"))

	RNAredundant, _ = NewAlphabet(
		"RNAredundant",
		false,
		[]byte("acguryswkmbdhvACGURYSWKMBDHV."),
		[]byte("ugcayrswmkvhdbUGCAYRSWMKVHDB."),
		[]byte(" -"),
		[]byte("nxNX"))

	Protein, _ = NewAlphabet(
		"Protein",
		false,
		[]byte("abcdefghijklmnpqrstvwyzABCDEFGHIJKLMNPQRSTVWYZ*_."),
		[]byte("abcdefghijklmnpqrstvwyzABCDEFGHIJKLMNPQRSTVWYZ*_."),
		[]byte(" -"),
		[]byte("xX"))

	Unlimit, _ = NewAlphabet(
		"Unlimit",
		true,
		nil,
		nil,
		nil,
		nil)

	abProtein = slice2map(byteutil.Alphabet(Protein.AllLetters()))
	abDNAredundant = slice2map(byteutil.Alphabet(DNAredundant.AllLetters()))
	abDNA = slice2map(byteutil.Alphabet(DNA.AllLetters()))
	abRNAredundant = slice2map(byteutil.Alphabet(RNAredundant.AllLetters()))
	abRNA = slice2map(byteutil.Alphabet(RNA.AllLetters()))
}

// AlphabetGuessSeqLenghtThreshold is the length of sequence prefix of the first FASTA record
// based which FastaRecord guesses the sequence type
var AlphabetGuessSeqLenghtThreshold = 10000

// GuessAlphabet guesses alphabet by given
func GuessAlphabet(seqs []byte) *Alphabet {
	var alphabetMap map[byte]bool
	if len(seqs) > AlphabetGuessSeqLenghtThreshold { // reduce guessing time
		alphabetMap = slice2map(byteutil.Alphabet(seqs[0:AlphabetGuessSeqLenghtThreshold]))
	} else {
		alphabetMap = slice2map(byteutil.Alphabet(seqs))
	}
	if isSubset(alphabetMap, abDNA) {
		return DNA
	}
	if isSubset(alphabetMap, abRNA) {
		return RNA
	}
	if isSubset(alphabetMap, abDNAredundant) {
		return DNAredundant
	}
	if isSubset(alphabetMap, abRNAredundant) {
		return RNAredundant
	}
	if isSubset(alphabetMap, abProtein) {
		return Protein
	}

	return Unlimit
}

// GuessAlphabetLessConservatively change DNA to DNAredundant and RNA to RNAredundant
func GuessAlphabetLessConservatively(seqs []byte) *Alphabet {
	ab := GuessAlphabet(seqs)
	if ab == DNA {
		return DNAredundant
	}
	if ab == RNA {
		return RNAredundant
	}
	return ab
}

func isSubset(query, subject map[byte]bool) bool {
	for b := range query {
		if _, ok := subject[b]; !ok {
			return false
		}
	}
	return true
}

func slice2map(s []byte) map[byte]bool {
	m := make(map[byte]bool)
	for _, b := range s {
		m[b] = true
	}
	return m
}
<|endoftext|>


<commit_before>package yamux

import (
	"bytes"
	"fmt"
	"io"
	"sync"
	"testing"
	"time"
)

type pipeConn struct {
	reader *io.PipeReader
	writer *io.PipeWriter
}

func (p *pipeConn) Read(b []byte) (int, error) {
	return p.reader.Read(b)
}

func (p *pipeConn) Write(b []byte) (int, error) {
	return p.writer.Write(b)
}

func (p *pipeConn) Close() error {
	p.reader.Close()
	return p.writer.Close()
}

func testConn() (io.ReadWriteCloser, io.ReadWriteCloser) {
	read1, write1 := io.Pipe()
	read2, write2 := io.Pipe()
	return &pipeConn{read1, write2}, &pipeConn{read2, write1}
}

func testClientServer() (*Session, *Session) {
	conn1, conn2 := testConn()
	client, _ := Client(conn1, nil)
	server, _ := Server(conn2, nil)
	return client, server
}

func TestPing(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	rtt, err := client.Ping()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if rtt == 0 {
		t.Fatalf("bad: %v", rtt)
	}

	rtt, err = server.Ping()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if rtt == 0 {
		t.Fatalf("bad: %v", rtt)
	}
}

func TestAccept(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	wg := &sync.WaitGroup{}
	wg.Add(4)

	go func() {
		defer wg.Done()
		stream, err := server.AcceptStream()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if id := stream.StreamID(); id != 1 {
			t.Fatalf("bad: %v", id)
		}
		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	go func() {
		defer wg.Done()
		stream, err := client.AcceptStream()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if id := stream.StreamID(); id != 2 {
			t.Fatalf("bad: %v", id)
		}
		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	go func() {
		defer wg.Done()
		stream, err := server.Open()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if id := stream.StreamID(); id != 2 {
			t.Fatalf("bad: %v", id)
		}
		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	go func() {
		defer wg.Done()
		stream, err := client.Open()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if id := stream.StreamID(); id != 1 {
			t.Fatalf("bad: %v", id)
		}
		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	doneCh := make(chan struct{})
	go func() {
		wg.Wait()
		close(doneCh)
	}()

	select {
	case <-doneCh:
	case <-time.After(time.Second):
		panic("timeout")
	}
}

func TestSendData_Small(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	wg := &sync.WaitGroup{}
	wg.Add(2)

	go func() {
		defer wg.Done()
		stream, err := server.AcceptStream()
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		buf := make([]byte, 4)
		for i := 0; i < 1000; i++ {
			n, err := stream.Read(buf)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4 {
				t.Fatalf("short read: %d", n)
			}
			if string(buf) != "test" {
				t.Fatalf("bad: %s", buf)
			}
		}

		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	go func() {
		defer wg.Done()
		stream, err := client.Open()
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		for i := 0; i < 1000; i++ {
			n, err := stream.Write([]byte("test"))
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4 {
				t.Fatalf("short write %d", n)
			}
		}

		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	doneCh := make(chan struct{})
	go func() {
		wg.Wait()
		close(doneCh)
	}()
	select {
	case <-doneCh:
	case <-time.After(time.Second):
		panic("timeout")
	}
}

func TestSendData_Large(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	data := make([]byte, 512*1024)
	for idx := range data {
		data[idx] = byte(idx % 256)
	}

	wg := &sync.WaitGroup{}
	wg.Add(2)

	go func() {
		defer wg.Done()
		stream, err := server.AcceptStream()
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		buf := make([]byte, 4*1024)
		for i := 0; i < 128; i++ {
			n, err := stream.Read(buf)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4*1024 {
				t.Fatalf("short read: %d", n)
			}
			for idx := range buf {
				if buf[idx] != byte(idx%256) {
					t.Fatalf("bad: %v %v %v", i, idx, buf[idx])
				}
			}
		}

		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	go func() {
		defer wg.Done()
		stream, err := client.Open()
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		n, err := stream.Write(data)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if n != len(data) {
			t.Fatalf("short write %d", n)
		}

		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	doneCh := make(chan struct{})
	go func() {
		wg.Wait()
		close(doneCh)
	}()
	select {
	case <-doneCh:
	case <-time.After(time.Second):
		panic("timeout")
	}
}

func TestGoAway(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	if err := server.GoAway(); err != nil {
		t.Fatalf("err: %v", err)
	}

	_, err := client.Open()
	if err != ErrRemoteGoAway {
		t.Fatalf("err: %v", err)
	}
}

func TestManyStreams(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	wg := &sync.WaitGroup{}

	acceptor := func(i int) {
		defer wg.Done()
		stream, err := server.AcceptStream()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		defer stream.Close()

		buf := make([]byte, 512)
		for {
			n, err := stream.Read(buf)
			if err == io.EOF {
				return
			}
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n == 0 {
				t.Fatalf("err: %v", err)
			}
		}
	}
	sender := func(i int) {
		defer wg.Done()
		stream, err := client.Open()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		defer stream.Close()

		msg := fmt.Sprintf("%08d", i)
		for i := 0; i < 1000; i++ {
			n, err := stream.Write([]byte(msg))
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != len(msg) {
				t.Fatalf("short write %d", n)
			}
		}
	}

	for i := 0; i < 50; i++ {
		wg.Add(2)
		go acceptor(i)
		go sender(i)
	}

	wg.Wait()
}

func TestManyStreams_PingPong(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	wg := &sync.WaitGroup{}

	ping := []byte("ping")
	pong := []byte("pong")

	acceptor := func(i int) {
		defer wg.Done()
		stream, err := server.AcceptStream()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		defer stream.Close()

		buf := make([]byte, 4)
		for {
			n, err := stream.Read(buf)
			if err == io.EOF {
				return
			}
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4 {
				t.Fatalf("err: %v", err)
			}
			if !bytes.Equal(buf, ping) {
				t.Fatalf("bad: %s", buf)
			}
			n, err = stream.Write(pong)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4 {
				t.Fatalf("err: %v", err)
			}
		}
	}
	sender := func(i int) {
		defer wg.Done()
		stream, err := client.Open()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		defer stream.Close()

		buf := make([]byte, 4)
		for i := 0; i < 1000; i++ {
			n, err := stream.Write(ping)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4 {
				t.Fatalf("short write %d", n)
			}

			n, err = stream.Read(buf)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4 {
				t.Fatalf("err: %v", err)
			}
			if !bytes.Equal(buf, pong) {
				t.Fatalf("bad: %s", buf)
			}
		}
	}

	for i := 0; i < 50; i++ {
		wg.Add(2)
		go acceptor(i)
		go sender(i)
	}

	wg.Wait()
}

func TestHalfClose(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	stream, err := client.Open()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if _, err := stream.Write([]byte("a")); err != nil {
		t.Fatalf("err: %v", err)
	}

	stream2, err := server.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	stream2.Close() // Half close

	buf := make([]byte, 4)
	n, err := stream2.Read(buf)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if n != 1 {
		t.Fatalf("bad: %v", n)
	}

	// Send more
	if _, err := stream.Write([]byte("bcd")); err != nil {
		t.Fatalf("err: %v", err)
	}
	stream.Close()

	// Read after close
	n, err = stream2.Read(buf)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if n != 3 {
		t.Fatalf("bad: %v", n)
	}

	// EOF after close
	n, err = stream2.Read(buf)
	if err != io.EOF {
		t.Fatalf("err: %v", err)
	}
	if n != 0 {
		t.Fatalf("bad: %v", n)
	}
}
<commit_msg>Testing deadlines<commit_after>package yamux

import (
	"bytes"
	"fmt"
	"io"
	"sync"
	"testing"
	"time"
)

type pipeConn struct {
	reader *io.PipeReader
	writer *io.PipeWriter
}

func (p *pipeConn) Read(b []byte) (int, error) {
	return p.reader.Read(b)
}

func (p *pipeConn) Write(b []byte) (int, error) {
	return p.writer.Write(b)
}

func (p *pipeConn) Close() error {
	p.reader.Close()
	return p.writer.Close()
}

func testConn() (io.ReadWriteCloser, io.ReadWriteCloser) {
	read1, write1 := io.Pipe()
	read2, write2 := io.Pipe()
	return &pipeConn{read1, write2}, &pipeConn{read2, write1}
}

func testClientServer() (*Session, *Session) {
	conn1, conn2 := testConn()
	client, _ := Client(conn1, nil)
	server, _ := Server(conn2, nil)
	return client, server
}

func TestPing(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	rtt, err := client.Ping()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if rtt == 0 {
		t.Fatalf("bad: %v", rtt)
	}

	rtt, err = server.Ping()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if rtt == 0 {
		t.Fatalf("bad: %v", rtt)
	}
}

func TestAccept(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	wg := &sync.WaitGroup{}
	wg.Add(4)

	go func() {
		defer wg.Done()
		stream, err := server.AcceptStream()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if id := stream.StreamID(); id != 1 {
			t.Fatalf("bad: %v", id)
		}
		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	go func() {
		defer wg.Done()
		stream, err := client.AcceptStream()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if id := stream.StreamID(); id != 2 {
			t.Fatalf("bad: %v", id)
		}
		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	go func() {
		defer wg.Done()
		stream, err := server.Open()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if id := stream.StreamID(); id != 2 {
			t.Fatalf("bad: %v", id)
		}
		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	go func() {
		defer wg.Done()
		stream, err := client.Open()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if id := stream.StreamID(); id != 1 {
			t.Fatalf("bad: %v", id)
		}
		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	doneCh := make(chan struct{})
	go func() {
		wg.Wait()
		close(doneCh)
	}()

	select {
	case <-doneCh:
	case <-time.After(time.Second):
		panic("timeout")
	}
}

func TestSendData_Small(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	wg := &sync.WaitGroup{}
	wg.Add(2)

	go func() {
		defer wg.Done()
		stream, err := server.AcceptStream()
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		buf := make([]byte, 4)
		for i := 0; i < 1000; i++ {
			n, err := stream.Read(buf)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4 {
				t.Fatalf("short read: %d", n)
			}
			if string(buf) != "test" {
				t.Fatalf("bad: %s", buf)
			}
		}

		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	go func() {
		defer wg.Done()
		stream, err := client.Open()
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		for i := 0; i < 1000; i++ {
			n, err := stream.Write([]byte("test"))
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4 {
				t.Fatalf("short write %d", n)
			}
		}

		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	doneCh := make(chan struct{})
	go func() {
		wg.Wait()
		close(doneCh)
	}()
	select {
	case <-doneCh:
	case <-time.After(time.Second):
		panic("timeout")
	}
}

func TestSendData_Large(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	data := make([]byte, 512*1024)
	for idx := range data {
		data[idx] = byte(idx % 256)
	}

	wg := &sync.WaitGroup{}
	wg.Add(2)

	go func() {
		defer wg.Done()
		stream, err := server.AcceptStream()
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		buf := make([]byte, 4*1024)
		for i := 0; i < 128; i++ {
			n, err := stream.Read(buf)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4*1024 {
				t.Fatalf("short read: %d", n)
			}
			for idx := range buf {
				if buf[idx] != byte(idx%256) {
					t.Fatalf("bad: %v %v %v", i, idx, buf[idx])
				}
			}
		}

		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	go func() {
		defer wg.Done()
		stream, err := client.Open()
		if err != nil {
			t.Fatalf("err: %v", err)
		}

		n, err := stream.Write(data)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		if n != len(data) {
			t.Fatalf("short write %d", n)
		}

		if err := stream.Close(); err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	doneCh := make(chan struct{})
	go func() {
		wg.Wait()
		close(doneCh)
	}()
	select {
	case <-doneCh:
	case <-time.After(time.Second):
		panic("timeout")
	}
}

func TestGoAway(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	if err := server.GoAway(); err != nil {
		t.Fatalf("err: %v", err)
	}

	_, err := client.Open()
	if err != ErrRemoteGoAway {
		t.Fatalf("err: %v", err)
	}
}

func TestManyStreams(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	wg := &sync.WaitGroup{}

	acceptor := func(i int) {
		defer wg.Done()
		stream, err := server.AcceptStream()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		defer stream.Close()

		buf := make([]byte, 512)
		for {
			n, err := stream.Read(buf)
			if err == io.EOF {
				return
			}
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n == 0 {
				t.Fatalf("err: %v", err)
			}
		}
	}
	sender := func(i int) {
		defer wg.Done()
		stream, err := client.Open()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		defer stream.Close()

		msg := fmt.Sprintf("%08d", i)
		for i := 0; i < 1000; i++ {
			n, err := stream.Write([]byte(msg))
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != len(msg) {
				t.Fatalf("short write %d", n)
			}
		}
	}

	for i := 0; i < 50; i++ {
		wg.Add(2)
		go acceptor(i)
		go sender(i)
	}

	wg.Wait()
}

func TestManyStreams_PingPong(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	wg := &sync.WaitGroup{}

	ping := []byte("ping")
	pong := []byte("pong")

	acceptor := func(i int) {
		defer wg.Done()
		stream, err := server.AcceptStream()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		defer stream.Close()

		buf := make([]byte, 4)
		for {
			n, err := stream.Read(buf)
			if err == io.EOF {
				return
			}
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4 {
				t.Fatalf("err: %v", err)
			}
			if !bytes.Equal(buf, ping) {
				t.Fatalf("bad: %s", buf)
			}
			n, err = stream.Write(pong)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4 {
				t.Fatalf("err: %v", err)
			}
		}
	}
	sender := func(i int) {
		defer wg.Done()
		stream, err := client.Open()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
		defer stream.Close()

		buf := make([]byte, 4)
		for i := 0; i < 1000; i++ {
			n, err := stream.Write(ping)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4 {
				t.Fatalf("short write %d", n)
			}

			n, err = stream.Read(buf)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if n != 4 {
				t.Fatalf("err: %v", err)
			}
			if !bytes.Equal(buf, pong) {
				t.Fatalf("bad: %s", buf)
			}
		}
	}

	for i := 0; i < 50; i++ {
		wg.Add(2)
		go acceptor(i)
		go sender(i)
	}

	wg.Wait()
}

func TestHalfClose(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	stream, err := client.Open()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if _, err := stream.Write([]byte("a")); err != nil {
		t.Fatalf("err: %v", err)
	}

	stream2, err := server.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	stream2.Close() // Half close

	buf := make([]byte, 4)
	n, err := stream2.Read(buf)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if n != 1 {
		t.Fatalf("bad: %v", n)
	}

	// Send more
	if _, err := stream.Write([]byte("bcd")); err != nil {
		t.Fatalf("err: %v", err)
	}
	stream.Close()

	// Read after close
	n, err = stream2.Read(buf)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if n != 3 {
		t.Fatalf("bad: %v", n)
	}

	// EOF after close
	n, err = stream2.Read(buf)
	if err != io.EOF {
		t.Fatalf("err: %v", err)
	}
	if n != 0 {
		t.Fatalf("bad: %v", n)
	}
}

func TestReadDeadline(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	stream, err := client.Open()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer stream.Close()

	stream2, err := server.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer stream2.Close()

	if err := stream.SetReadDeadline(time.Now().Add(5 * time.Millisecond)); err != nil {
		t.Fatalf("err: %v", err)
	}

	buf := make([]byte, 4)
	if _, err := stream.Read(buf); err != ErrTimeout {
		t.Fatalf("err: %v", err)
	}
}

func TestWriteDeadline(t *testing.T) {
	client, server := testClientServer()
	defer client.Close()
	defer server.Close()

	stream, err := client.Open()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer stream.Close()

	stream2, err := server.Accept()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	defer stream2.Close()

	if err := stream.SetWriteDeadline(time.Now().Add(50 * time.Millisecond)); err != nil {
		t.Fatalf("err: %v", err)
	}

	buf := make([]byte, 512)
	for i := 0; i < int(initialStreamWindow); i++ {
		_, err := stream.Write(buf)
		if err != nil && err == ErrTimeout {
			return
		} else if err != nil {
			t.Fatalf("err: %v", err)
		}
	}
	t.Fatalf("Expected timeout")
}
<|endoftext|>