<commit_before>// Copyright 2013 The StudyGolang Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// http://studygolang.com, http://golang.top
// Author: polaris	polaris@studygolang.com

// 解析 http://www.reddit.com/r/golang 最新 Go 信息
package logic

import (
	"errors"
	"math/rand"
	"regexp"
	"strconv"
	"strings"
	"time"

	. "db"
	"model"

	"github.com/PuerkitoBio/goquery"
	"github.com/polaris1119/logger"
)

type RedditLogic struct {
	domain string
	golang string
}

var DefaultReddit = RedditLogic{
	domain: "http://www.reddit.com",
	golang: "/r/golang/new/",
}

// Parse 获取url对应的资源并根据规则进行解析
func (this *RedditLogic) Parse(redditUrl string) error {
	redditUrl = strings.TrimSpace(redditUrl)
	if redditUrl == "" {
		redditUrl = this.domain + this.golang
	} else if !strings.HasPrefix(redditUrl, "http") {
		redditUrl = "http://" + redditUrl
	}

	var (
		doc *goquery.Document
		err error
	)

	if doc, err = goquery.NewDocument(redditUrl); err != nil {
		logger.Errorln("goquery reddit newdocument error:", err)
		return err
	}

	// 最后面的先入库处理
	resourcesSelection := doc.Find("#siteTable .link")

	for i := resourcesSelection.Length() - 1; i >= 0; i-- {
		err = this.dealRedditOneResource(goquery.NewDocumentFromNode(resourcesSelection.Get(i)).Selection)

		if err != nil {
			logger.Errorln(err)
		}
	}

	return err
}

var PresetUids = []int{1, 1747, 1748, 1827}

var resourceRe = regexp.MustCompile(`\n\n`)

// 处理 Reddit 中的一条资源
func (this *RedditLogic) dealRedditOneResource(contentSelection *goquery.Selection) error {
	aSelection := contentSelection.Find(".title a.title")

	title := aSelection.Text()
	if title == "" {
		return errors.New("title is empty")
	}

	resourceUrl, ok := aSelection.Attr("href")
	if !ok || resourceUrl == "" {
		return errors.New("resource url is empty")
	}

	isReddit := false

	resource := &model.Resource{}
	// Reddit 自身的内容
	if contentSelection.HasClass("self") {
		isReddit = true
		resourceUrl = this.domain + resourceUrl
	}

	_, err := MasterDB.Where("url=?", resourceUrl).Get(resource)
	if err != nil {
		return err
	}
	// 已经存在
	if resource.Id != 0 {
		// 如果是 reddit 本身的，可以更新评论信息
		if !isReddit {
			return errors.New("url" + resourceUrl + "has exists!")
		}
	}

	if isReddit {

		resource.Form = model.ContentForm

		var doc *goquery.Document

		if doc, err = goquery.NewDocument(resourceUrl); err != nil {
			return errors.New("goquery reddit.com/r/golang self newdocument error:" + err.Error())
		}

		content, err := doc.Find("#siteTable .usertext .md").Html()
		if err != nil {
			return err
		}

		doc.Find(".commentarea .comment .usertext .md").Each(func(i int, contentSel *goquery.Selection) {
			if i == 0 {
				content += `<hr/>**评论：**<br/><br/>`
			}

			comment, err := contentSel.Html()
			if err != nil {
				return
			}

			comment = strings.TrimSpace(comment)
			comment = resourceRe.ReplaceAllLiteralString(comment, "\n")

			author := contentSel.ParentsFiltered(".usertext").Prev().Find(".author").Text()
			content += author + ":
" + comment + "
"
		})

		if strings.TrimSpace(content) == "" {
			return errors.New("goquery reddit.com/r/golang self newdocument(" + resourceUrl + ") error: content is empty")
		}

		resource.Content = content

		// reddit 本身的，当做其他资源
		resource.Catid = 4
	} else {
		resource.Form = model.LinkForm

		// Github，是开源项目
		if contentSelection.Find(".title .domain a").Text() == "github.com" {
			resource.Catid = 2
		} else {
			resource.Catid = 1
		}
	}

	resource.Title = title
	resource.Url = resourceUrl
	resource.Uid = PresetUids[rand.Intn(4)]

	ctime := time.Now()
	datetime, ok := contentSelection.Find(".tagline time").Attr("datetime")
	if ok {
		dtime, err := time.ParseInLocation(time.RFC3339, datetime, time.UTC)
		if err != nil {
			logger.Errorln("parse ctime error:", err)
		} else {
			ctime = dtime.Local()
		}
	}
	resource.Ctime = model.OftenTime(ctime)

	if resource.Id == 0 {
		session := MasterDB.NewSession()
		defer session.Close()
		session.Begin()

		_, err = session.Insert(resource)
		if err != nil {
			session.Rollback()
			return errors.New("insert into Resource error:" + err.Error())
		}

		// 存扩展信息
		resourceEx := &model.ResourceEx{}
		resourceEx.Id = resource.Id
		if _, err = session.Insert(resourceEx); err != nil {
			session.Rollback()
			return errors.New("insert into ResourceEx error:" + err.Error())
		}
		session.Commit()
	} else {
		if _, err = MasterDB.Id(resource.Id).Update(resource); err != nil {
			return errors.New("update resource:" + strconv.Itoa(resource.Id) + " error:" + err.Error())
		}
	}

	return nil
}
<commit_msg>reddit->https<commit_after>// Copyright 2013 The StudyGolang Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// http://studygolang.com, http://golang.top
// Author: polaris	polaris@studygolang.com

// 解析 http://www.reddit.com/r/golang 最新 Go 信息
package logic

import (
	"errors"
	"math/rand"
	"regexp"
	"strconv"
	"strings"
	"time"

	. "db"
	"model"

	"github.com/PuerkitoBio/goquery"
	"github.com/polaris1119/logger"
)

type RedditLogic struct {
	domain string
	golang string
}

var DefaultReddit = RedditLogic{
	domain: "https://www.reddit.com",
	golang: "/r/golang/new/",
}

// Parse 获取url对应的资源并根据规则进行解析
func (this *RedditLogic) Parse(redditUrl string) error {
	redditUrl = strings.TrimSpace(redditUrl)
	if redditUrl == "" {
		redditUrl = this.domain + this.golang
	} else if !strings.HasPrefix(redditUrl, "https") {
		redditUrl = "https://" + redditUrl
	}

	var (
		doc *goquery.Document
		err error
	)

	if doc, err = goquery.NewDocument(redditUrl); err != nil {
		logger.Errorln("goquery reddit newdocument error:", err)
		return err
	}

	// 最后面的先入库处理
	resourcesSelection := doc.Find("#siteTable .link")

	for i := resourcesSelection.Length() - 1; i >= 0; i-- {
		err = this.dealRedditOneResource(goquery.NewDocumentFromNode(resourcesSelection.Get(i)).Selection)

		if err != nil {
			logger.Errorln(err)
		}
	}

	return err
}

var PresetUids = []int{1, 1747, 1748, 1827}

var resourceRe = regexp.MustCompile(`\n\n`)

// 处理 Reddit 中的一条资源
func (this *RedditLogic) dealRedditOneResource(contentSelection *goquery.Selection) error {
	aSelection := contentSelection.Find(".title a.title")

	title := aSelection.Text()
	if title == "" {
		return errors.New("title is empty")
	}

	resourceUrl, ok := aSelection.Attr("href")
	if !ok || resourceUrl == "" {
		return errors.New("resource url is empty")
	}

	isReddit := false

	resource := &model.Resource{}
	// Reddit 自身的内容
	if contentSelection.HasClass("self") {
		isReddit = true
		resourceUrl = this.domain + resourceUrl
	}

	_, err := MasterDB.Where("url=?", resourceUrl).Get(resource)
	if err != nil {
		return err
	}
	// 已经存在
	if resource.Id != 0 {
		// 如果是 reddit 本身的，可以更新评论信息
		if !isReddit {
			return errors.New("url" + resourceUrl + "has exists!")
		}
	}

	if isReddit {

		resource.Form = model.ContentForm

		var doc *goquery.Document

		if doc, err = goquery.NewDocument(resourceUrl); err != nil {
			return errors.New("goquery reddit.com/r/golang self newdocument error:" + err.Error())
		}

		content, err := doc.Find("#siteTable .usertext .md").Html()
		if err != nil {
			return err
		}

		doc.Find(".commentarea .comment .usertext .md").Each(func(i int, contentSel *goquery.Selection) {
			if i == 0 {
				content += `<hr/>**评论：**<br/><br/>`
			}

			comment, err := contentSel.Html()
			if err != nil {
				return
			}

			comment = strings.TrimSpace(comment)
			comment = resourceRe.ReplaceAllLiteralString(comment, "\n")

			author := contentSel.ParentsFiltered(".usertext").Prev().Find(".author").Text()
			content += author + ":
" + comment + "
"
		})

		if strings.TrimSpace(content) == "" {
			return errors.New("goquery reddit.com/r/golang self newdocument(" + resourceUrl + ") error: content is empty")
		}

		resource.Content = content

		// reddit 本身的，当做其他资源
		resource.Catid = 4
	} else {
		resource.Form = model.LinkForm

		// Github，是开源项目
		if contentSelection.Find(".title .domain a").Text() == "github.com" {
			resource.Catid = 2
		} else {
			resource.Catid = 1
		}
	}

	resource.Title = title
	resource.Url = resourceUrl
	resource.Uid = PresetUids[rand.Intn(4)]

	ctime := time.Now()
	datetime, ok := contentSelection.Find(".tagline time").Attr("datetime")
	if ok {
		dtime, err := time.ParseInLocation(time.RFC3339, datetime, time.UTC)
		if err != nil {
			logger.Errorln("parse ctime error:", err)
		} else {
			ctime = dtime.Local()
		}
	}
	resource.Ctime = model.OftenTime(ctime)

	if resource.Id == 0 {
		session := MasterDB.NewSession()
		defer session.Close()
		session.Begin()

		_, err = session.Insert(resource)
		if err != nil {
			session.Rollback()
			return errors.New("insert into Resource error:" + err.Error())
		}

		// 存扩展信息
		resourceEx := &model.ResourceEx{}
		resourceEx.Id = resource.Id
		if _, err = session.Insert(resourceEx); err != nil {
			session.Rollback()
			return errors.New("insert into ResourceEx error:" + err.Error())
		}
		session.Commit()
	} else {
		if _, err = MasterDB.Id(resource.Id).Update(resource); err != nil {
			return errors.New("update resource:" + strconv.Itoa(resource.Id) + " error:" + err.Error())
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package lint

import (
	"errors"
	"fmt"
	"strings"
)

// Runtime stores a set of rules.
type Runtime struct {
	rules   Rules
	configs RuntimeConfigs
}

// NewRuntime creates a new Runtime.
func NewRuntime(configs RuntimeConfigs) *Runtime {
	return &Runtime{
		rules:   make(Rules),
		configs: configs,
	}
}

// AddRules adds rules, of which the name will be added a prefix to reduce collisions
func (r *Runtime) AddRules(rules ...Rule) error {
	for _, rl := range rules {
		if _, found := r.rules[rl.Info().Name]; found {
			return fmt.Errorf("duplicate repository entry with name %q", rl.Info().Name)
		}
		r.rules[rl.Info().Name] = rl
	}
	return nil
}

// Run executes rules on the request when a config is found for the file path of the request.
//
// If the found config contains rule configs for some rules, the status and
// category of the affected rules will be updated accordingly. In other words,
// rule configs can be used to turn on/off certain rules and change the category
// of the returned problems.
func (r *Runtime) Run(req Request) (Response, error) {
	finalResp := Response{}
	var errMessages []string

	for name, rl := range r.rules {
		config, err := r.configs.getRuleConfig(req.ProtoFile().Path(), name)

		if err != nil {
			errMessages = append(errMessages, err.Error())
			continue
		}

		config = defaultRuleConfig.withOverride(config)

		if config.Status == Enabled {
			if resp, err := rl.Lint(req); err == nil {
				for _, p := range resp.Problems {

					if config.Category != "" {
						p.category = config.Category
					}

					finalResp.Problems = append(finalResp.Problems, p)
				}
			} else {
				errMessages = append(errMessages, err.Error())
			}
		}
	}

	var err error
	if len(errMessages) != 0 {
		err = errors.New(strings.Join(errMessages, "; "))
	}

	return finalResp, err
}
<commit_msg>Reformatted code, removed unnecessary condition<commit_after>package lint

import (
	"errors"
	"fmt"
	"strings"
)

// Runtime stores a set of rules.
type Runtime struct {
	rules   Rules
	configs RuntimeConfigs
}

// NewRuntime creates a new Runtime.
func NewRuntime(configs RuntimeConfigs) *Runtime {
	return &Runtime{
		rules:   make(Rules),
		configs: configs,
	}
}

// AddRules adds rules, of which the name will be added a prefix to reduce collisions
func (r *Runtime) AddRules(rules ...Rule) error {
	for _, rl := range rules {
		if _, found := r.rules[rl.Info().Name]; found {
			return fmt.Errorf("duplicate repository entry with name %q", rl.Info().Name)
		}
		r.rules[rl.Info().Name] = rl
	}
	return nil
}

// Run executes rules on the request when a config is found for the file path of the request.
//
// If the found config contains rule configs for some rules, the status and
// category of the affected rules will be updated accordingly. In other words,
// rule configs can be used to turn on/off certain rules and change the category
// of the returned problems.
func (r *Runtime) Run(req Request) (Response, error) {
	finalResp := Response{}
	var errMessages []string

	for name, rl := range r.rules {
		config := defaultRuleConfig

		if c, err := r.configs.getRuleConfig(req.ProtoFile().Path(), name); err == nil {
			config = config.withOverride(c)
		} else {
			errMessages = append(errMessages, err.Error())
			continue
		}

		if config.Status == Enabled {
			if resp, err := rl.Lint(req); err == nil {
				for _, p := range resp.Problems {
					p.category = config.Category
					finalResp.Problems = append(finalResp.Problems, p)
				}
			} else {
				errMessages = append(errMessages, err.Error())
			}
		}
	}

	var err error
	if len(errMessages) != 0 {
		err = errors.New(strings.Join(errMessages, "; "))
	}

	return finalResp, err
}
<|endoftext|>


<commit_before>package dou

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"reflect"
	"strings"
	"testing"
)

type recorder struct {
	calledTime int
}

func newRecorder() *recorder {
	return &recorder{calledTime: 0}
}

type testAPI struct {
	beforeDispatchCalled bool
	afterDispatchCalled  bool
	recoverCalled        bool
}

func (p *testAPI) BeforeDispatch(w http.ResponseWriter, r *http.Request) (http.ResponseWriter, *http.Request) {
	p.beforeDispatchCalled = true
	return w, r
}

func (p *testAPI) AfterDispatch(w http.ResponseWriter, r *http.Request) (http.ResponseWriter, *http.Request) {
	p.afterDispatchCalled = true
	return w, r
}

func (p *testAPI) OnPanic(w http.ResponseWriter, r *http.Request) {
	p.recoverCalled = true
}

// Enable to stub
var testAPIMarshal = func(v interface{}) ([]byte, error) {
	return nil, nil
}

func (p *testAPI) Marshal(v interface{}) ([]byte, error) {
	return testAPIMarshal(v)
}

// Enable to stub
var testAPIUnmarshal = func(data []byte, v interface{}) error {
	return nil
}

func (p *testAPI) Unmarshal(data []byte, v interface{}) error {
	return testAPIUnmarshal(data, v)
}

func (p *testAPI) APIStatus(w http.ResponseWriter, code int) {
}

func TestCallBeforeDispatchAndAfterDispatch(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	ta := &testAPI{}

	Register("testapi", ta)
	defer delete(plugins, "testapi")

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	if err != nil {
		panic(err)
	}

	a.LogStackTrace = false

	if err != nil {
		panic(err)
	}

	a.ServeHTTP(response, request)

	if !ta.beforeDispatchCalled {
		t.Error("Plugin.BeforeDispatch should be called")
	}

	if !ta.afterDispatchCalled {
		t.Error("Plugin.AfterDispatch should be called")
	}
}

func TestCallOnPanicIfOccurPanicInHandler(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	ta := &testAPI{}

	Register("testapi", ta)
	defer delete(plugins, "testapi")

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		panic("<test panic>")
	})

	if err != nil {
		panic(err)
	}

	a.LogStackTrace = false

	if err != nil {
		panic(err)
	}

	a.ServeHTTP(response, request)

	if !ta.recoverCalled {
		t.Error("Plugin.OnPanic should be called")
	}
}

func TestCallAfterDispatchIfOccurPanicInHandler(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	ta := &testAPI{}

	Register("testapi", ta)
	defer delete(plugins, "testapi")

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		panic("<test panic>")
	})

	if err != nil {
		panic(err)
	}

	a.LogStackTrace = false

	a.ServeHTTP(response, request)

	if !ta.afterDispatchCalled {
		t.Error("Plugin.AfterDispatch should be called")
	}
}

func TestNewSafeWriter(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		sw := NewSafeWriter(w)

		if !reflect.DeepEqual(sw.ResponseWriter, w) {
			t.Error("NewSafeWriter should set given value to .ResponseWriter")
		}

		if sw.Wrote != false {
			t.Error("NewSafeWriter should set false to .wrote")
		}
	}).ServeHTTP(response, request)
}

func TestNewSafeWriterWrite(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		sw := NewSafeWriter(w)
		sw.Write([]byte("hello"))

		if sw.Wrote != true {
			t.Error("NewSafeWriter.Wrote should be true after called Write")
		}
	}).ServeHTTP(response, request)
}

func TestAPIUnmarshal(t *testing.T) {
	request, _ := http.NewRequest("GET", `/?json={"name": "ToQoz"}`, nil)
	response := httptest.NewRecorder()

	Register("testapi", &testAPI{})
	defer delete(plugins, "testapi")

	// stub testAPI.Marshal
	testAPIUnmarshal = func(data []byte, v interface{}) error {
		return json.Unmarshal(data, v)
	}

	a, err := NewAPI("testapi")

	if err != nil {
		panic(err)
	}

	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		v := map[string]string{}

		expected := map[string]string{"name": "ToQoz"}

		a.Unmarshal([]byte(request.FormValue("json")), &v)

		if !reflect.DeepEqual(v, expected) {
			t.Errorf("fail to Unmarshal.\nexpected: %v\ngot: %v\n", expected, v)
		}
	})

	a.ServeHTTP(response, request)

}

func TestAPIOkSetGivenHTTPStatusCodeAndResponseBody(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	Register("testapi", &testAPI{})
	defer delete(plugins, "testapi")

	expectedBodyString := "stubed testAPIMarshal"
	expectedCode := http.StatusCreated

	// stub testAPI.Marshal
	testAPIMarshal = func(v interface{}) ([]byte, error) {
		return []byte(expectedBodyString), nil
	}

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.Ok(w, "", expectedCode)
	})

	if err != nil {
		panic(err)
	}

	a.ServeHTTP(response, request)

	// [Test] http status code
	if response.Code != expectedCode {
		t.Errorf("API.OK should set given status code\nexpected: %v\ngot: %v\n", expectedCode, response.Code)
	}

	// [Test] responseBody
	gotBody, err := ioutil.ReadAll(response.Body)

	if err != nil {
		panic(err)
	}

	gotBodyString := strings.TrimSuffix(string(gotBody), "\n")

	if gotBodyString != expectedBodyString {
		t.Errorf("API.OK should marshal given resource and write it\nexpected: %v\ngot: %v\n", expectedBodyString, gotBodyString)
	}
}

func TestAPIOkSet200IfGiven0AsHTTPStatusCode(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	Register("testapi", &testAPI{})
	defer delete(plugins, "testapi")

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.Ok(w, "", 0)
	})

	if err != nil {
		panic(err)
	}

	a.ServeHTTP(response, request)

	if response.Code != http.StatusOK {
		t.Errorf("API.OK should set 200 if given 0\nexpected: %v\ngot: %v\n", http.StatusOK, response.Code)
	}

}

func TestAPIErrorSetGivenHTTPStatusCodeAndResponseBody(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	Register("testapi", &testAPI{})
	defer delete(plugins, "testapi")

	expectedBodyString := "stubed testAPIMarshal"
	expectedCode := http.StatusNotFound

	// stub testAPI.Marshal
	testAPIMarshal = func(v interface{}) ([]byte, error) {
		return []byte(expectedBodyString), nil
	}

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.Error(w, "", expectedCode)
	})

	if err != nil {
		panic(err)
	}

	a.ServeHTTP(response, request)

	// [Test] http status code
	if response.Code != expectedCode {
		t.Errorf("API.Error should set given status code\nexpected: %v\ngot: %v\n", expectedCode, response.Code)
	}

	// [Test] responseBody
	gotBody, err := ioutil.ReadAll(response.Body)

	if err != nil {
		panic(err)
	}

	gotBodyString := strings.TrimSuffix(string(gotBody), "\n")

	if gotBodyString != expectedBodyString {
		t.Errorf("API.Error should marshal given resource and write it\nexpected: \"%v\"\ngot: \"%v\"\n", expectedBodyString, gotBodyString)
	}
}

func TestAPIErrorSet500IfGiven0AsHTTPStatusCode(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	Register("testapi", &testAPI{})
	defer delete(plugins, "testapi")

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.Error(w, "", 0)
	})

	if err != nil {
		panic(err)
	}

	a.ServeHTTP(response, request)

	if response.Code != http.StatusInternalServerError {
		t.Errorf("API.OK should set 200 if given 0\nexpected: %v\ngot: %v\n", http.StatusOK, response.Code)
	}

}
<commit_msg>Write test for Deregister<commit_after>package dou

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"reflect"
	"strings"
	"testing"
)

type recorder struct {
	calledTime int
}

func newRecorder() *recorder {
	return &recorder{calledTime: 0}
}

type testAPI struct {
	beforeDispatchCalled bool
	afterDispatchCalled  bool
	recoverCalled        bool
}

func (p *testAPI) BeforeDispatch(w http.ResponseWriter, r *http.Request) (http.ResponseWriter, *http.Request) {
	p.beforeDispatchCalled = true
	return w, r
}

func (p *testAPI) AfterDispatch(w http.ResponseWriter, r *http.Request) (http.ResponseWriter, *http.Request) {
	p.afterDispatchCalled = true
	return w, r
}

func (p *testAPI) OnPanic(w http.ResponseWriter, r *http.Request) {
	p.recoverCalled = true
}

// Enable to stub
var testAPIMarshal = func(v interface{}) ([]byte, error) {
	return nil, nil
}

func (p *testAPI) Marshal(v interface{}) ([]byte, error) {
	return testAPIMarshal(v)
}

// Enable to stub
var testAPIUnmarshal = func(data []byte, v interface{}) error {
	return nil
}

func (p *testAPI) Unmarshal(data []byte, v interface{}) error {
	return testAPIUnmarshal(data, v)
}

func (p *testAPI) APIStatus(w http.ResponseWriter, code int) {
}

// - Test

func TestNewAPIWithUnregisteredPluginName(t *testing.T) {
	api, err := NewAPI("unknown")

	if api != nil {
		t.Error("If NewAPI called with unknown plugin name, api should be nil")
	}

	if err == nil {
		t.Error("If NewAPI called with unknown plugin name, api should not be nil")
	}
}

func TestDeregister(t *testing.T) {
	Register("testapi", &testAPI{})
	Deregister("testapi")

	if _, ok := plugins["testapi"]; ok != false {
		t.Error("Deregister should delete registered plugin")
	}
}

func TestCallBeforeDispatchAndAfterDispatch(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	ta := &testAPI{}

	Register("testapi", ta)
	defer delete(plugins, "testapi")

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})

	if err != nil {
		panic(err)
	}

	a.LogStackTrace = false

	if err != nil {
		panic(err)
	}

	a.ServeHTTP(response, request)

	if !ta.beforeDispatchCalled {
		t.Error("Plugin.BeforeDispatch should be called")
	}

	if !ta.afterDispatchCalled {
		t.Error("Plugin.AfterDispatch should be called")
	}
}

func TestCallOnPanicIfOccurPanicInHandler(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	ta := &testAPI{}

	Register("testapi", ta)
	defer delete(plugins, "testapi")

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		panic("<test panic>")
	})

	if err != nil {
		panic(err)
	}

	a.LogStackTrace = false

	if err != nil {
		panic(err)
	}

	a.ServeHTTP(response, request)

	if !ta.recoverCalled {
		t.Error("Plugin.OnPanic should be called")
	}
}

func TestCallAfterDispatchIfOccurPanicInHandler(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	ta := &testAPI{}

	Register("testapi", ta)
	defer delete(plugins, "testapi")

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		panic("<test panic>")
	})

	if err != nil {
		panic(err)
	}

	a.LogStackTrace = false

	a.ServeHTTP(response, request)

	if !ta.afterDispatchCalled {
		t.Error("Plugin.AfterDispatch should be called")
	}
}

func TestNewSafeWriter(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		sw := NewSafeWriter(w)

		if !reflect.DeepEqual(sw.ResponseWriter, w) {
			t.Error("NewSafeWriter should set given value to .ResponseWriter")
		}

		if sw.Wrote != false {
			t.Error("NewSafeWriter should set false to .wrote")
		}
	}).ServeHTTP(response, request)
}

func TestNewSafeWriterWrite(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		sw := NewSafeWriter(w)
		sw.Write([]byte("hello"))

		if sw.Wrote != true {
			t.Error("NewSafeWriter.Wrote should be true after called Write")
		}
	}).ServeHTTP(response, request)
}

func TestAPIUnmarshal(t *testing.T) {
	request, _ := http.NewRequest("GET", `/?json={"name": "ToQoz"}`, nil)
	response := httptest.NewRecorder()

	Register("testapi", &testAPI{})
	defer delete(plugins, "testapi")

	// stub testAPI.Marshal
	testAPIUnmarshal = func(data []byte, v interface{}) error {
		return json.Unmarshal(data, v)
	}

	a, err := NewAPI("testapi")

	if err != nil {
		panic(err)
	}

	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		v := map[string]string{}

		expected := map[string]string{"name": "ToQoz"}

		a.Unmarshal([]byte(request.FormValue("json")), &v)

		if !reflect.DeepEqual(v, expected) {
			t.Errorf("fail to Unmarshal.\nexpected: %v\ngot: %v\n", expected, v)
		}
	})

	a.ServeHTTP(response, request)

}

func TestAPIOkSetGivenHTTPStatusCodeAndResponseBody(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	Register("testapi", &testAPI{})
	defer delete(plugins, "testapi")

	expectedBodyString := "stubed testAPIMarshal"
	expectedCode := http.StatusCreated

	// stub testAPI.Marshal
	testAPIMarshal = func(v interface{}) ([]byte, error) {
		return []byte(expectedBodyString), nil
	}

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.Ok(w, "", expectedCode)
	})

	if err != nil {
		panic(err)
	}

	a.ServeHTTP(response, request)

	// [Test] http status code
	if response.Code != expectedCode {
		t.Errorf("API.OK should set given status code\nexpected: %v\ngot: %v\n", expectedCode, response.Code)
	}

	// [Test] responseBody
	gotBody, err := ioutil.ReadAll(response.Body)

	if err != nil {
		panic(err)
	}

	gotBodyString := strings.TrimSuffix(string(gotBody), "\n")

	if gotBodyString != expectedBodyString {
		t.Errorf("API.OK should marshal given resource and write it\nexpected: %v\ngot: %v\n", expectedBodyString, gotBodyString)
	}
}

func TestAPIOkSet200IfGiven0AsHTTPStatusCode(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	Register("testapi", &testAPI{})
	defer delete(plugins, "testapi")

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.Ok(w, "", 0)
	})

	if err != nil {
		panic(err)
	}

	a.ServeHTTP(response, request)

	if response.Code != http.StatusOK {
		t.Errorf("API.OK should set 200 if given 0\nexpected: %v\ngot: %v\n", http.StatusOK, response.Code)
	}

}

func TestAPIErrorSetGivenHTTPStatusCodeAndResponseBody(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	Register("testapi", &testAPI{})
	defer delete(plugins, "testapi")

	expectedBodyString := "stubed testAPIMarshal"
	expectedCode := http.StatusNotFound

	// stub testAPI.Marshal
	testAPIMarshal = func(v interface{}) ([]byte, error) {
		return []byte(expectedBodyString), nil
	}

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.Error(w, "", expectedCode)
	})

	if err != nil {
		panic(err)
	}

	a.ServeHTTP(response, request)

	// [Test] http status code
	if response.Code != expectedCode {
		t.Errorf("API.Error should set given status code\nexpected: %v\ngot: %v\n", expectedCode, response.Code)
	}

	// [Test] responseBody
	gotBody, err := ioutil.ReadAll(response.Body)

	if err != nil {
		panic(err)
	}

	gotBodyString := strings.TrimSuffix(string(gotBody), "\n")

	if gotBodyString != expectedBodyString {
		t.Errorf("API.Error should marshal given resource and write it\nexpected: \"%v\"\ngot: \"%v\"\n", expectedBodyString, gotBodyString)
	}
}

func TestAPIErrorSet500IfGiven0AsHTTPStatusCode(t *testing.T) {
	request, _ := http.NewRequest("GET", "/", nil)
	response := httptest.NewRecorder()

	Register("testapi", &testAPI{})
	defer delete(plugins, "testapi")

	a, err := NewAPI("testapi")
	a.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.Error(w, "", 0)
	})

	if err != nil {
		panic(err)
	}

	a.ServeHTTP(response, request)

	if response.Code != http.StatusInternalServerError {
		t.Errorf("API.OK should set 200 if given 0\nexpected: %v\ngot: %v\n", http.StatusOK, response.Code)
	}

}
<|endoftext|>


<commit_before>package core

type Registers struct {
    R [16]uint32
    Psr uint32
}

/* Special registers in r13-15 */
const (
    LR = 13
    SP = 14
    PC = 15
)

func (regs *Registers) Lr() uint32 {
    return regs.R[LR]
}

func (regs *Registers) Sp() uint32 {
    return regs.R[SP]
}

func (regs *Registers) Pc() uint32 {
    return regs.R[PC]
}
<commit_msg>Provide type for general purpose registers<commit_after>package core

type GeneralRegs [16]uint32

type Registers struct {
    R GeneralRegs
    Psr uint32
}

/* Special registers in r13-15 */
const (
    LR = 13
    SP = 14
    PC = 15
)

func (regs *Registers) Lr() uint32 {
    return regs.R[LR]
}

func (regs *Registers) Sp() uint32 {
    return regs.R[SP]
}

func (regs *Registers) Pc() uint32 {
    return regs.R[PC]
}
<|endoftext|>


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

package core

import (
	// "fmt"
	"log"
	"sort"
	"sync"

	"github.com/go-ego/riot/types"
	"github.com/go-ego/riot/utils"
)

var (
	gOnlyID bool
)

// Ranker ranker
type Ranker struct {
	lock struct {
		sync.RWMutex
		fields map[uint64]interface{}
		docs   map[uint64]bool
		// new
		content map[uint64]string
		attri   map[uint64]interface{}
	}
	initialized bool
}

// Init init ranker
func (ranker *Ranker) Init(onlyID bool) {
	if ranker.initialized == true {
		log.Fatal("排序器不能初始化两次")
	}
	ranker.initialized = true

	ranker.lock.fields = make(map[uint64]interface{})
	ranker.lock.docs = make(map[uint64]bool)

	gOnlyID = onlyID
	if !onlyID {
		// new
		ranker.lock.content = make(map[uint64]string)
		ranker.lock.attri = make(map[uint64]interface{})
	}
}

// AddDoc add doc
// 给某个文档添加评分字段
func (ranker *Ranker) AddDoc(
	// docId uint64, fields interface{}, content string, attri interface{}) {
	docId uint64, fields interface{}, content ...interface{}) {
	if ranker.initialized == false {
		log.Fatal("排序器尚未初始化")
	}

	ranker.lock.Lock()
	ranker.lock.fields[docId] = fields
	ranker.lock.docs[docId] = true
	if !gOnlyID {
		// new
		if len(content) > 0 {
			ranker.lock.content[docId] = content[0].(string)
		}

		if len(content) > 1 {
			ranker.lock.attri[docId] = content[1]
			// ranker.lock.attri[docId] = attri
		}
	}

	ranker.lock.Unlock()
}

// RemoveDoc 删除某个文档的评分字段
func (ranker *Ranker) RemoveDoc(docId uint64) {
	if ranker.initialized == false {
		log.Fatal("排序器尚未初始化")
	}

	ranker.lock.Lock()
	delete(ranker.lock.fields, docId)
	delete(ranker.lock.docs, docId)

	if !gOnlyID {
		// new
		delete(ranker.lock.content, docId)
		delete(ranker.lock.attri, docId)
	}

	ranker.lock.Unlock()
}

// Rank rank
// 给文档评分并排序
func (ranker *Ranker) Rank(
	docs []types.IndexedDocument, options types.RankOptions,
	countDocsOnly bool) (types.ScoredDocuments, int) {

	if ranker.initialized == false {
		log.Fatal("排序器尚未初始化")
	}

	// 对每个文档评分
	var outputDocs types.ScoredDocuments
	numDocs := 0
	for _, d := range docs {
		ranker.lock.RLock()
		// 判断doc是否存在
		if _, ok := ranker.lock.docs[d.DocId]; ok {

			fs := ranker.lock.fields[d.DocId]
			content := ranker.lock.content[d.DocId]
			attri := ranker.lock.attri[d.DocId]

			ranker.lock.RUnlock()
			// 计算评分并剔除没有分值的文档
			scores := options.ScoringCriteria.Score(d, fs)
			if len(scores) > 0 {
				if !countDocsOnly {
					if !gOnlyID {
						outputDocs = append(outputDocs, types.ScoredDocument{
							DocId: d.DocId,
							// new
							Fields:  fs,
							Content: content,
							Attri:   attri,
							//
							Scores:                scores,
							TokenSnippetLocations: d.TokenSnippetLocations,
							TokenLocations:        d.TokenLocations})
					} else {
						outputDocs = append(outputDocs, types.ScoredDocument{
							DocId:                 d.DocId,
							Scores:                scores,
							TokenSnippetLocations: d.TokenSnippetLocations,
							TokenLocations:        d.TokenLocations})
					}
				}
				numDocs++
			}
		} else {
			ranker.lock.RUnlock()
		}
	}

	// 排序
	if !countDocsOnly {
		if options.ReverseOrder {
			sort.Sort(sort.Reverse(outputDocs))
		} else {
			sort.Sort(outputDocs)
		}
		// 当用户要求只返回部分结果时返回部分结果
		var start, end int
		if options.MaxOutputs != 0 {
			start = utils.MinInt(options.OutputOffset, len(outputDocs))
			end = utils.MinInt(options.OutputOffset+options.MaxOutputs, len(outputDocs))
		} else {
			start = utils.MinInt(options.OutputOffset, len(outputDocs))
			end = len(outputDocs)
		}
		return outputDocs[start:end], numDocs
	}
	return outputDocs, numDocs
}
<commit_msg>update init and fix test<commit_after>// Copyright 2013 Hui Chen
// Copyright 2016 ego authors
//
// Licensed under the Apache License, Version 2.0 (the "License"): you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package core

import (
	// "fmt"
	"log"
	"sort"
	"sync"

	"github.com/go-ego/riot/types"
	"github.com/go-ego/riot/utils"
)

var (
	gOnlyID bool
)

// Ranker ranker
type Ranker struct {
	lock struct {
		sync.RWMutex
		fields map[uint64]interface{}
		docs   map[uint64]bool
		// new
		content map[uint64]string
		attri   map[uint64]interface{}
	}
	initialized bool
}

// Init init ranker
func (ranker *Ranker) Init(onlyID ...bool) {
	if ranker.initialized == true {
		log.Fatal("排序器不能初始化两次")
	}
	ranker.initialized = true

	ranker.lock.fields = make(map[uint64]interface{})
	ranker.lock.docs = make(map[uint64]bool)

	if len(onlyID) > 0 {
		gOnlyID = onlyID[0]
	}

	if !gOnlyID {
		// new
		ranker.lock.content = make(map[uint64]string)
		ranker.lock.attri = make(map[uint64]interface{})
	}
}

// AddDoc add doc
// 给某个文档添加评分字段
func (ranker *Ranker) AddDoc(
	// docId uint64, fields interface{}, content string, attri interface{}) {
	docId uint64, fields interface{}, content ...interface{}) {
	if ranker.initialized == false {
		log.Fatal("排序器尚未初始化")
	}

	ranker.lock.Lock()
	ranker.lock.fields[docId] = fields
	ranker.lock.docs[docId] = true
	if !gOnlyID {
		// new
		if len(content) > 0 {
			ranker.lock.content[docId] = content[0].(string)
		}

		if len(content) > 1 {
			ranker.lock.attri[docId] = content[1]
			// ranker.lock.attri[docId] = attri
		}
	}

	ranker.lock.Unlock()
}

// RemoveDoc 删除某个文档的评分字段
func (ranker *Ranker) RemoveDoc(docId uint64) {
	if ranker.initialized == false {
		log.Fatal("排序器尚未初始化")
	}

	ranker.lock.Lock()
	delete(ranker.lock.fields, docId)
	delete(ranker.lock.docs, docId)

	if !gOnlyID {
		// new
		delete(ranker.lock.content, docId)
		delete(ranker.lock.attri, docId)
	}

	ranker.lock.Unlock()
}

// Rank rank
// 给文档评分并排序
func (ranker *Ranker) Rank(
	docs []types.IndexedDocument, options types.RankOptions,
	countDocsOnly bool) (types.ScoredDocuments, int) {

	if ranker.initialized == false {
		log.Fatal("排序器尚未初始化")
	}

	// 对每个文档评分
	var outputDocs types.ScoredDocuments
	numDocs := 0
	for _, d := range docs {
		ranker.lock.RLock()
		// 判断doc是否存在
		if _, ok := ranker.lock.docs[d.DocId]; ok {

			fs := ranker.lock.fields[d.DocId]
			content := ranker.lock.content[d.DocId]
			attri := ranker.lock.attri[d.DocId]

			ranker.lock.RUnlock()
			// 计算评分并剔除没有分值的文档
			scores := options.ScoringCriteria.Score(d, fs)
			if len(scores) > 0 {
				if !countDocsOnly {
					if !gOnlyID {
						outputDocs = append(outputDocs, types.ScoredDocument{
							DocId: d.DocId,
							// new
							Fields:  fs,
							Content: content,
							Attri:   attri,
							//
							Scores:                scores,
							TokenSnippetLocations: d.TokenSnippetLocations,
							TokenLocations:        d.TokenLocations})
					} else {
						outputDocs = append(outputDocs, types.ScoredDocument{
							DocId:                 d.DocId,
							Scores:                scores,
							TokenSnippetLocations: d.TokenSnippetLocations,
							TokenLocations:        d.TokenLocations})
					}
				}
				numDocs++
			}
		} else {
			ranker.lock.RUnlock()
		}
	}

	// 排序
	if !countDocsOnly {
		if options.ReverseOrder {
			sort.Sort(sort.Reverse(outputDocs))
		} else {
			sort.Sort(outputDocs)
		}
		// 当用户要求只返回部分结果时返回部分结果
		var start, end int
		if options.MaxOutputs != 0 {
			start = utils.MinInt(options.OutputOffset, len(outputDocs))
			end = utils.MinInt(options.OutputOffset+options.MaxOutputs, len(outputDocs))
		} else {
			start = utils.MinInt(options.OutputOffset, len(outputDocs))
			end = len(outputDocs)
		}
		return outputDocs[start:end], numDocs
	}
	return outputDocs, numDocs
}
<|endoftext|>


<commit_before>package core

import (
	"fmt"
	"github.com/blackbass1988/access_logs_stats/core/output"
	_ "github.com/blackbass1988/access_logs_stats/core/output/console"
	_ "github.com/blackbass1988/access_logs_stats/core/output/zabbix"
	"sort"
	"strconv"
	"strings"
	"sync"
)

//Sender sends data to output. omg omg omg
type Sender struct {
	filter *Filter

	//ссылка на конфиг приложения
	config *Config

	//настроенный отправлятор, реализации настраиваются в конфиге "outputs"
	output *output.Output

	//мап флоатов с реализацией агрегирующих фунций
	floatData map[string]*Float64Data

	//закешированное кол-во секунд в периоде, указанному в "period" конфигурации
	periodInSeconds float64

	//здесь хранятся числа по полям, указанные в "aggregates" конфигурации
	floatsForAggregates map[string][]float64

	//здесь хранятся счетчики уникальных значений по полям, указанные в "counts" конфигурации
	//хранится по схеме поле.уник_значение.кол-во
	//вывод происходит по схеме - кол-во в 1 секунду
	counts map[string]map[string]uint64

	globalLock sync.Mutex
}

func (s *Sender) resetData() {
	s.globalLock.Lock()
	s.floatsForAggregates = make(map[string][]float64)
	s.floatData = make(map[string]*Float64Data)
	s.counts = make(map[string]map[string]uint64)
	s.globalLock.Unlock()
}

func (s *Sender) appendIfOk(row *RowEntry) (err error) {

	if s.filter.MatchString(row.Raw) {

		s.globalLock.Lock()

		for field, val := range row.Fields {

			if _, ok := s.config.Aggregates[field]; ok {
				valFloat, err := strconv.ParseFloat(val, 10)
				checkOrFail(err)
				s.floatsForAggregates[field] = append(s.floatsForAggregates[field], valFloat)
			}

			//в конфиге указано поле, как поле, по которому считаются
			// суммы по уникальным значениям
			if _, ok := s.config.Counts[field]; ok {
				if s.counts[field] == nil {
					s.counts[field] = make(map[string]uint64)
				}
				s.counts[field][val]++

			}
		}
		s.globalLock.Unlock()
	}

	return err
}

func (s *Sender) sendStats() (err error) {

	s.globalLock.Lock()
	for _, metricsOfField := range s.filter.Items {

		for _, metric := range metricsOfField.Metrics {
			s.appendToOutput(metricsOfField.Field, metric)
		}
	}
	s.globalLock.Unlock()
	go s.output.Send()

	return err
}

func (s *Sender) appendToOutput(field string, metric string) {
	var (
		cnt             uint64
		key             string
		ok              bool
		result          float64
		periodInSeconds float64
	)

	key = fmt.Sprintf("%s_%v", field, metric)

	switch {
	case metric == "min":
		s.output.AddMessage(key, fmt.Sprintf("%.3f", s.getFloatData(field).Min()))
	case metric == "max":
		s.output.AddMessage(key, fmt.Sprintf("%.3f", s.getFloatData(field).Max()))
	case metric == "len":
		s.output.AddMessage(key, fmt.Sprintf("%d", s.getFloatData(field).Len()))
	case metric == "avg":
		s.output.AddMessage(key, fmt.Sprintf("%.3f", s.getFloatData(field).Avg()))
	case metric == "sum":
		s.output.AddMessage(key, fmt.Sprintf("%.3f", s.getFloatData(field).Sum()))
	case metric == "sum_ps":

		result := s.getFloatData(field).Sum()
		periodInSeconds = s.getPeriodInSeconds()
		if periodInSeconds == 0 {
			result = 0
		} else {
			result = s.getFloatData(field).Sum() / s.getPeriodInSeconds()
		}

		s.output.AddMessage(key, fmt.Sprintf("%.3f", result))
	case metric == "ips":
		s.output.AddMessage(key, fmt.Sprintf("%.3f", s.getFloatData(field).ItemsPerSeconds(s.getPeriodInSeconds())))
	case strings.Contains(metric, "cent_"):
		cent := strings.Split(metric, "_")
		centFloat, err := strconv.ParseFloat(cent[1], 10)
		checkOrFail(err)
		s.output.AddMessage(key, fmt.Sprintf("%.3f", s.getFloatData(field).Percentile(centFloat)))

	case metric == "uniq":
		s.output.AddMessage(key, fmt.Sprintf("%d", s.getUniqCnt(field)))
	case metric == "uniq_ps":
		s.output.AddMessage(key, fmt.Sprintf("%.3f", float64(s.getUniqCnt(field))/s.getPeriodInSeconds()))
	case strings.Contains(metric, "cps_"):
		metrics := strings.Split(metric, "_")
		metric = metrics[1]
		if _, ok = s.counts[field]; !ok {
			cnt = 0
		} else if cnt, ok = s.counts[field][metric]; !ok {
			cnt = 0
		}
		s.output.AddMessage(key, fmt.Sprintf("%.3f", float64(cnt)/s.getPeriodInSeconds()))
	case strings.Contains(metric, "percentage_"):
		metrics := strings.Split(metric, "_")
		metric = metrics[1]

		total := s.getTotalCountByField(field)

		result = 0
		if cnt, ok = s.counts[field][metric]; ok && total > 0 {
			result = float64(cnt * 100 / total)
		}

		s.output.AddMessage(key, fmt.Sprintf("%.3f", result))
	}
}

func (s *Sender) getTotalCountByField(field string) uint64 {
	var (
		ok  bool
		cnt uint64
	)
	if _, ok = s.counts[field]; !ok {
		return 0
	}

	cnt = 0
	for _, c := range s.counts[field] {
		cnt += c
	}

	return cnt
}

func (s *Sender) getUniqCnt(field string) uint64 {
	var (
		cnt uint64
		ok  bool
	)
	cnt = 0
	if _, ok = s.counts[field]; ok {
		cnt = uint64(len(s.counts[field]))
	}
	return cnt
}
func (s *Sender) getFloatData(field string) *Float64Data {
	//кешируем флоатдату
	if _, ok := s.floatData[field]; !ok {
		f := Float64Data(s.floatsForAggregates[field])
		s.floatData[field] = &f
		sort.Sort(s.floatData[field])
	}
	return s.floatData[field]
}

func (s *Sender) getPeriodInSeconds() float64 {
	if s.periodInSeconds == 0 {
		s.periodInSeconds = s.config.Period.Seconds()
	}
	return s.periodInSeconds
}

//NewSender create new sender
func NewSender(filter *Filter, config *Config) (*Sender, error) {
	sender := new(Sender)
	sender.filter = filter
	sender.config = config

	sender.output = new(output.Output)

	if len(filter.Prefix) > 0 {
		sender.output.SetPrefix(filter.Prefix)
	}

	for _, s := range config.Outputs {
		sender.output.Init(s.Type, s.Settings)
	}

	return sender, nil
}

//SenderCollection is a collection of Senders
type SenderCollection struct {
	procs  []*Sender
	config *Config
}

//NewSenderCollection create SenderCollection of Senders
func NewSenderCollection(config *Config) *SenderCollection {
	subProcesses := new(SenderCollection)

	processes := []*Sender{}
	for _, f := range config.Filters {
		sp, err := NewSender(f, config)
		checkOrFail(err)
		processes = append(processes, sp)
	}

	subProcesses.procs = processes
	subProcesses.config = config
	subProcesses.resetData()
	return subProcesses
}

func (s *SenderCollection) resetData() {
	for _, proc := range s.procs {
		proc.resetData()
	}
}

//appendData appends RowEntry to every filter instance from config
func (s *SenderCollection) appendData(row *RowEntry) error {
	var err error

	for _, proc := range s.procs {
		go proc.appendIfOk(row)
	}

	return err
}

func (s *SenderCollection) sendStats() {
	var wg sync.WaitGroup
	for _, proc := range s.procs {
		wg.Add(1)
		go func(proc *Sender) {
			defer wg.Done()
			proc.sendStats()
		}(proc)
	}
	wg.Wait()
}
<commit_msg>cyclo in core/sender.go<commit_after>package core

import (
	"fmt"
	"github.com/blackbass1988/access_logs_stats/core/output"
	_ "github.com/blackbass1988/access_logs_stats/core/output/console"
	_ "github.com/blackbass1988/access_logs_stats/core/output/zabbix"
	"sort"
	"strconv"
	"strings"
	"sync"
)

//Sender sends data to output. omg omg omg
type Sender struct {
	filter *Filter

	//ссылка на конфиг приложения
	config *Config

	//настроенный отправлятор, реализации настраиваются в конфиге "outputs"
	output *output.Output

	//мап флоатов с реализацией агрегирующих фунций
	floatData map[string]*Float64Data

	//закешированное кол-во секунд в периоде, указанному в "period" конфигурации
	periodInSeconds float64

	//здесь хранятся числа по полям, указанные в "aggregates" конфигурации
	floatsForAggregates map[string][]float64

	//здесь хранятся счетчики уникальных значений по полям, указанные в "counts" конфигурации
	//хранится по схеме поле.уник_значение.кол-во
	//вывод происходит по схеме - кол-во в 1 секунду
	counts map[string]map[string]uint64

	globalLock sync.Mutex
}

func (s *Sender) resetData() {
	s.globalLock.Lock()
	s.floatsForAggregates = make(map[string][]float64)
	s.floatData = make(map[string]*Float64Data)
	s.counts = make(map[string]map[string]uint64)
	s.globalLock.Unlock()
}

func (s *Sender) appendIfOk(row *RowEntry) (err error) {

	if s.filter.MatchString(row.Raw) {

		s.globalLock.Lock()

		for field, val := range row.Fields {

			if _, ok := s.config.Aggregates[field]; ok {
				valFloat, err := strconv.ParseFloat(val, 10)
				checkOrFail(err)
				s.floatsForAggregates[field] = append(s.floatsForAggregates[field], valFloat)
			}

			//в конфиге указано поле, как поле, по которому считаются
			// суммы по уникальным значениям
			if _, ok := s.config.Counts[field]; ok {
				if s.counts[field] == nil {
					s.counts[field] = make(map[string]uint64)
				}
				s.counts[field][val]++

			}
		}
		s.globalLock.Unlock()
	}

	return err
}

func (s *Sender) sendStats() (err error) {

	s.globalLock.Lock()
	for _, metricsOfField := range s.filter.Items {

		for _, metric := range metricsOfField.Metrics {
			s.appendToOutput(metricsOfField.Field, metric)
		}
	}
	s.globalLock.Unlock()
	go s.output.Send()

	return err
}

func (s *Sender) appendToOutput(field string, metric string) {
	var (
		key             string
		periodInSeconds float64
		message         string
	)

	key = fmt.Sprintf("%s_%v", field, metric)

	switch {
	case metric == "min":
		message = fmt.Sprintf("%.3f", s.getFloatData(field).Min())
	case metric == "max":
		message = fmt.Sprintf("%.3f", s.getFloatData(field).Max())
	case metric == "len":
		message = fmt.Sprintf("%d", s.getFloatData(field).Len())
	case metric == "avg":
		message = fmt.Sprintf("%.3f", s.getFloatData(field).Avg())
	case metric == "sum":
		message = fmt.Sprintf("%.3f", s.getFloatData(field).Sum())
	case metric == "sum_ps":
		result := s.getFloatData(field).Sum()
		periodInSeconds = s.getPeriodInSeconds()
		if periodInSeconds == 0 {
			result = 0
		} else {
			result = s.getFloatData(field).Sum() / s.getPeriodInSeconds()
		}

		message = fmt.Sprintf("%.3f", result)
	case metric == "ips":
		message = fmt.Sprintf("%.3f", s.getFloatData(field).ItemsPerSeconds(s.getPeriodInSeconds()))
	case strings.Contains(metric, "cent_"):
		cent := strings.Split(metric, "_")
		centFloat, err := strconv.ParseFloat(cent[1], 10)
		checkOrFail(err)
		message = fmt.Sprintf("%.3f", s.getFloatData(field).Percentile(centFloat))
	case metric == "uniq":
		message = fmt.Sprintf("%d", s.getUniqCnt(field))
	case metric == "uniq_ps":
		message = fmt.Sprintf("%.3f", float64(s.getUniqCnt(field))/s.getPeriodInSeconds())
	case strings.Contains(metric, "cps_"):
		message = s.processCps(metric, field)
	case strings.Contains(metric, "percentage_"):
		message = s.processPercentage(metric, field)
	}
	s.output.AddMessage(key, message)
}

func (s *Sender) processCps(metric string, field string) string {
	var (
		ok  bool
		cnt uint64
	)

	metrics := strings.Split(metric, "_")
	metric = metrics[1]
	if _, ok = s.counts[field]; !ok {
		cnt = 0
	} else if cnt, ok = s.counts[field][metric]; !ok {
		cnt = 0
	}
	return fmt.Sprintf("%.3f", float64(cnt)/s.getPeriodInSeconds())
}

func (s *Sender) processPercentage(metric string, field string) string {
	var (
		ok     bool
		cnt    uint64
		result float64
	)

	metrics := strings.Split(metric, "_")
	metric = metrics[1]

	total := s.getTotalCountByField(field)

	result = 0
	if cnt, ok = s.counts[field][metric]; ok && total > 0 {
		result = float64(cnt * 100 / total)
	}
	return fmt.Sprintf("%.3f", result)
}

func (s *Sender) getTotalCountByField(field string) uint64 {
	var (
		ok  bool
		cnt uint64
	)
	if _, ok = s.counts[field]; !ok {
		return 0
	}

	cnt = 0
	for _, c := range s.counts[field] {
		cnt += c
	}

	return cnt
}

func (s *Sender) getUniqCnt(field string) uint64 {
	var (
		cnt uint64
		ok  bool
	)
	cnt = 0
	if _, ok = s.counts[field]; ok {
		cnt = uint64(len(s.counts[field]))
	}
	return cnt
}
func (s *Sender) getFloatData(field string) *Float64Data {
	//кешируем флоатдату
	if _, ok := s.floatData[field]; !ok {
		f := Float64Data(s.floatsForAggregates[field])
		s.floatData[field] = &f
		sort.Sort(s.floatData[field])
	}
	return s.floatData[field]
}

func (s *Sender) getPeriodInSeconds() float64 {
	if s.periodInSeconds == 0 {
		s.periodInSeconds = s.config.Period.Seconds()
	}
	return s.periodInSeconds
}

//NewSender create new sender
func NewSender(filter *Filter, config *Config) (*Sender, error) {
	sender := new(Sender)
	sender.filter = filter
	sender.config = config

	sender.output = new(output.Output)

	if len(filter.Prefix) > 0 {
		sender.output.SetPrefix(filter.Prefix)
	}

	for _, s := range config.Outputs {
		sender.output.Init(s.Type, s.Settings)
	}

	return sender, nil
}

//SenderCollection is a collection of Senders
type SenderCollection struct {
	procs  []*Sender
	config *Config
}

//NewSenderCollection create SenderCollection of Senders
func NewSenderCollection(config *Config) *SenderCollection {
	subProcesses := new(SenderCollection)

	processes := []*Sender{}
	for _, f := range config.Filters {
		sp, err := NewSender(f, config)
		checkOrFail(err)
		processes = append(processes, sp)
	}

	subProcesses.procs = processes
	subProcesses.config = config
	subProcesses.resetData()
	return subProcesses
}

func (s *SenderCollection) resetData() {
	for _, proc := range s.procs {
		proc.resetData()
	}
}

//appendData appends RowEntry to every filter instance from config
func (s *SenderCollection) appendData(row *RowEntry) error {
	var err error

	for _, proc := range s.procs {
		go proc.appendIfOk(row)
	}

	return err
}

func (s *SenderCollection) sendStats() {
	var wg sync.WaitGroup
	for _, proc := range s.procs {
		wg.Add(1)
		go func(proc *Sender) {
			defer wg.Done()
			proc.sendStats()
		}(proc)
	}
	wg.Wait()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"debug/pe"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/bnagy/gapstone"
	"github.com/codegangsta/cli"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"
	"io"
	"log"
	"os"
	"unicode/utf16"
)

var PAGE_SIZE uint64 = 0x1000

var inputFlag = cli.StringFlag{
	Name:  "input_file",
	Usage: "file to decode",
}

var outputFlag = cli.StringFlag{
	Name:  "output_file",
	Usage: "where to write decoded data",
}

var verboseFlag = cli.BoolFlag{
	Name:  "verbose",
	Usage: "print debugging output",
}

var RequiredFlagNotProvidedError = errors.New("Required flag not provided.")

func checkRequiredArgs(c *cli.Context, requiredFlags []cli.StringFlag) error {
	for _, flag := range requiredFlags {
		if c.GlobalString(flag.Name) != "" {
			continue
		}
		if c.String(flag.Name) != "" {
			continue
		}
		log.Printf(fmt.Sprintf("Error: '%s' value required", flag.Name))
		return RequiredFlagNotProvidedError
	}
	return nil
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func roundUp(i uint64, base uint64) uint64 {
	if i%base == 0x0 {
		return i
	} else {
		return i + base - (i % base)
	}
}

func roundUpToPage(i uint64) uint64 {
	return roundUp(i, PAGE_SIZE)
}

type Arch string
type Mode string

var ARCH_X86 Arch = "x86"
var MODE_32 Mode = "32"

var InvalidArchError = errors.New("Invalid ARCH provided.")
var InvalidModeError = errors.New("Invalid MODE provided.")

type LoadedModule struct {
	Name        string
	BaseAddress uint64 // VA (*not* RVA)
	EntryPoint  uint64 // VA (*not* RVA)
}

func (m LoadedModule) VA(rva uint64) uint64 {
	return m.BaseAddress + rva
}

func (m LoadedModule) MemRead(env *Environment, rva uint64, length uint64) ([]byte, error) {
	return env.u.MemRead(m.VA(rva), length)
}

func (m LoadedModule) MemWrite(env *Environment, rva uint64, data []byte) error {
	return env.u.MemWrite(m.VA(rva), data)
}

type MemoryRegion struct {
	Address uint64 // VA
	Length  uint64
}

type Environment struct {
	u             uc.Unicorn
	Arch          Arch
	Mode          Mode
	loadedModules []LoadedModule
	memoryRegions []MemoryRegion
}

func NewEnvironment(arch Arch, mode Mode) (*Environment, error) {

	if arch != ARCH_X86 {
		return nil, InvalidArchError
	}
	if mode != MODE_32 {
		return nil, InvalidModeError
	}

	u, e := uc.NewUnicorn(uc.ARCH_X86, uc.MODE_32)
	if e != nil {
		return nil, e
	}

	return &Environment{
		u:             u,
		Arch:          arch,
		Mode:          mode,
		loadedModules: make([]LoadedModule, 1),
		memoryRegions: make([]MemoryRegion, 5),
	}, nil
}

type ImageImportDirectory struct {
	rvaImportLookupTable  uint32
	TimeDateStamp         uint32
	ForwarderChain        uint32
	rvaModuleName         uint32
	rvaImportAddressTable uint32
}

func (env *Environment) LoadPE(name string, f *pe.File) (*LoadedModule, error) {
	var imageBase uint64
	var addressOfEntryPoint uint64
	var dataDirectory [16]pe.DataDirectory

	if optionalHeader, ok := f.OptionalHeader.(*pe.OptionalHeader32); ok {
		imageBase = uint64(optionalHeader.ImageBase)
		addressOfEntryPoint = uint64(optionalHeader.AddressOfEntryPoint)
		dataDirectory = optionalHeader.DataDirectory
	} else {
		return nil, InvalidModeError
	}

	mod := LoadedModule{
		Name:        name,
		BaseAddress: imageBase,
		EntryPoint:  imageBase + addressOfEntryPoint,
	}

	for _, section := range f.Sections {
		h := section.SectionHeader

		fmt.Printf("section: %s\n", section.SectionHeader.Name)
		fmt.Printf("  virtual address: 0x%x\n", section.SectionHeader.VirtualAddress)
		fmt.Printf("  virtual size: 0x%x\n", section.SectionHeader.VirtualSize)
		fmt.Printf("  file offset: 0x%x\n", section.SectionHeader.Offset)
		fmt.Printf("  file size: 0x%x\n", section.SectionHeader.Size)

		secStart := mod.VA(uint64(h.VirtualAddress))
		secLength := roundUpToPage(uint64(h.VirtualSize))
		e := env.u.MemMap(secStart, secLength)
		check(e)

		d, e := section.Data()
		check(e)

		e = mod.MemWrite(env, uint64(h.VirtualAddress), d)
		check(e)

		region := MemoryRegion{
			Address: secStart,
			Length:  secLength,
		}
		env.memoryRegions = append(env.memoryRegions, region)
	}

	// since we always map at ImageBase, we don't need to apply (32bit) relocs
	// TODO: check 64bit reloc types

	importDirectory := dataDirectory[1]
	importRva := importDirectory.VirtualAddress
	importSize := importDirectory.Size
	fmt.Printf("import rva: 0x%x\n", importRva)
	fmt.Printf("import size: 0x%x\n", importSize)

	// TODO: check ranges
	d, e := mod.MemRead(env, uint64(importDirectory.VirtualAddress), uint64(importDirectory.Size))
	check(e)

	p := bytes.NewBuffer(d)
	for {
		var dir ImageImportDirectory
		binary.Read(p, binary.LittleEndian, &dir.rvaImportLookupTable)
		fmt.Printf("rva import lookup table: 0x%x\n", dir.rvaImportLookupTable)
		if dir.rvaImportLookupTable == 0 {
			break
		}
		binary.Read(p, binary.LittleEndian, &dir.TimeDateStamp)
		fmt.Printf("time date stamp: 0x%x\n", dir.TimeDateStamp)

		binary.Read(p, binary.LittleEndian, &dir.ForwarderChain)
		fmt.Printf("forwarder chain: 0x%x\n", dir.ForwarderChain)

		binary.Read(p, binary.LittleEndian, &dir.rvaModuleName)
		fmt.Printf("rva module name: 0x%x\n", dir.rvaModuleName)

		moduleNameBuf, e := mod.MemRead(env, uint64(dir.rvaModuleName), 0x100)
		check(e)
		moduleName, e := readAscii(moduleNameBuf)
		check(e)

		fmt.Printf("module name: %s\n", string(moduleName))

		binary.Read(p, binary.LittleEndian, &dir.rvaImportAddressTable)
	}

	env.loadedModules = append(env.loadedModules, mod)
	return &mod, nil
}

func readAscii(buf []byte) (string, error) {
	br := bufio.NewReader(bytes.NewReader(buf))
	bytez, e := br.ReadBytes(byte(0x00))
	check(e)
	return string(bytez), nil
}

func readUtf16le(buf []byte) ([]rune, error) {
	start := 0
	end := 0

	for i := 0; i < len(buf)/2; i++ {
		if !(buf[i*2] == 0 && buf[i*2+1] == 0) {
			end = i * 2
			break
		}
	}

	numChars := (end - start) / 2
	d := make([]uint16, numChars)
	for i := 0; i < numChars; i++ {
		d[i] = uint16(buf[i*2]) | (uint16(buf[i*2+1]) << 8)
	}
	return utf16.Decode(d), nil
}

var GAPSTONE_ARCH_MAP = map[Arch]int{
	ARCH_X86: gapstone.CS_ARCH_X86,
}

var GAPSTONE_MODE_MAP = map[Mode]uint{
	MODE_32: gapstone.CS_MODE_32,
}

func (env *Environment) disassembleBytes(data []byte, address uint64, w io.Writer) error {
	// TODO: cache the engine on the Environment?

	engine, e := gapstone.New(
		GAPSTONE_ARCH_MAP[env.Arch],
		GAPSTONE_MODE_MAP[env.Mode],
	)
	check(e)
	defer engine.Close()

	insns, e := engine.Disasm([]byte(data), address, 0 /* all instructions */)
	check(e)

	w.Write([]byte(fmt.Sprintf("Disasm:\n")))
	for _, insn := range insns {
		w.Write([]byte(fmt.Sprintf("0x%x:\t%s\t\t%s\n", insn.Address, insn.Mnemonic, insn.OpStr)))
	}

	return nil
}

func (env *Environment) Disassemble(address uint64, length uint64, w io.Writer) error {
	d, e := env.u.MemRead(address, length)
	check(e)
	return env.disassembleBytes(d, address, w)
}

func doit(path string) error {
	f, e := pe.Open(path)
	check(e)

	env, e := NewEnvironment(ARCH_X86, MODE_32)
	check(e)

	m, e := env.LoadPE(path, f)
	check(e)

	e = env.Disassemble(m.EntryPoint, 0x20, os.Stdout)
	check(e)

	return nil
}

func main() {
	app := cli.NewApp()
	app.Version = "0.1"
	app.Name = "crystal-tiger-pe-loader"
	app.Usage = "Load PE files."
	app.Commands = []cli.Command{
		{
			Name:    "load",
			Aliases: []string{"l"},
			Usage:   "load a pe file",
			Flags:   []cli.Flag{inputFlag, verboseFlag},
			Action: func(c *cli.Context) {
				if checkRequiredArgs(c, []cli.StringFlag{inputFlag}) != nil {
					return
				}
				check(doit(c.String("input_file")))
			},
		},
	}
	app.Run(os.Args)
}
<commit_msg>pe: refactor out loadPESection<commit_after>package main

import (
	"bufio"
	"bytes"
	"debug/pe"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/bnagy/gapstone"
	"github.com/codegangsta/cli"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"
	"io"
	"log"
	"os"
	"unicode/utf16"
)

var PAGE_SIZE uint64 = 0x1000

var inputFlag = cli.StringFlag{
	Name:  "input_file",
	Usage: "file to decode",
}

var outputFlag = cli.StringFlag{
	Name:  "output_file",
	Usage: "where to write decoded data",
}

var verboseFlag = cli.BoolFlag{
	Name:  "verbose",
	Usage: "print debugging output",
}

var RequiredFlagNotProvidedError = errors.New("Required flag not provided.")

func checkRequiredArgs(c *cli.Context, requiredFlags []cli.StringFlag) error {
	for _, flag := range requiredFlags {
		if c.GlobalString(flag.Name) != "" {
			continue
		}
		if c.String(flag.Name) != "" {
			continue
		}
		log.Printf(fmt.Sprintf("Error: '%s' value required", flag.Name))
		return RequiredFlagNotProvidedError
	}
	return nil
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func roundUp(i uint64, base uint64) uint64 {
	if i%base == 0x0 {
		return i
	} else {
		return i + base - (i % base)
	}
}

func roundUpToPage(i uint64) uint64 {
	return roundUp(i, PAGE_SIZE)
}

type Arch string
type Mode string

var ARCH_X86 Arch = "x86"
var MODE_32 Mode = "32"

var InvalidArchError = errors.New("Invalid ARCH provided.")
var InvalidModeError = errors.New("Invalid MODE provided.")

type LoadedModule struct {
	Name        string
	BaseAddress uint64 // VA (*not* RVA)
	EntryPoint  uint64 // VA (*not* RVA)
}

func (m LoadedModule) VA(rva uint64) uint64 {
	return m.BaseAddress + rva
}

func (m LoadedModule) MemRead(env *Environment, rva uint64, length uint64) ([]byte, error) {
	return env.u.MemRead(m.VA(rva), length)
}

func (m LoadedModule) MemWrite(env *Environment, rva uint64, data []byte) error {
	return env.u.MemWrite(m.VA(rva), data)
}

type MemoryRegion struct {
	Address uint64 // VA
	Length  uint64
}

type Environment struct {
	u             uc.Unicorn
	Arch          Arch
	Mode          Mode
	loadedModules []LoadedModule
	memoryRegions []MemoryRegion
}

func NewEnvironment(arch Arch, mode Mode) (*Environment, error) {

	if arch != ARCH_X86 {
		return nil, InvalidArchError
	}
	if mode != MODE_32 {
		return nil, InvalidModeError
	}

	u, e := uc.NewUnicorn(uc.ARCH_X86, uc.MODE_32)
	if e != nil {
		return nil, e
	}

	return &Environment{
		u:             u,
		Arch:          arch,
		Mode:          mode,
		loadedModules: make([]LoadedModule, 1),
		memoryRegions: make([]MemoryRegion, 5),
	}, nil
}

type ImageImportDirectory struct {
	rvaImportLookupTable  uint32
	TimeDateStamp         uint32
	ForwarderChain        uint32
	rvaModuleName         uint32
	rvaImportAddressTable uint32
}

func (env *Environment) loadPESection(mod *LoadedModule, section *pe.Section) error {
	h := section.SectionHeader

	fmt.Printf("section: %s\n", section.SectionHeader.Name)
	fmt.Printf("  virtual address: 0x%x\n", section.SectionHeader.VirtualAddress)
	fmt.Printf("  virtual size: 0x%x\n", section.SectionHeader.VirtualSize)
	fmt.Printf("  file offset: 0x%x\n", section.SectionHeader.Offset)
	fmt.Printf("  file size: 0x%x\n", section.SectionHeader.Size)

	secStart := mod.VA(uint64(h.VirtualAddress))
	secLength := roundUpToPage(uint64(h.VirtualSize))
	e := env.u.MemMap(secStart, secLength)
	check(e)

	d, e := section.Data()
	check(e)

	e = mod.MemWrite(env, uint64(h.VirtualAddress), d)
	check(e)

	// TODO: apply permissions

	region := MemoryRegion{
		Address: secStart,
		Length:  secLength,
	}
	env.memoryRegions = append(env.memoryRegions, region)

	return nil
}

func (env *Environment) LoadPE(name string, f *pe.File) (*LoadedModule, error) {
	var imageBase uint64
	var addressOfEntryPoint uint64
	var dataDirectory [16]pe.DataDirectory

	if optionalHeader, ok := f.OptionalHeader.(*pe.OptionalHeader32); ok {
		imageBase = uint64(optionalHeader.ImageBase)
		addressOfEntryPoint = uint64(optionalHeader.AddressOfEntryPoint)
		dataDirectory = optionalHeader.DataDirectory
	} else {
		return nil, InvalidModeError
	}

	mod := LoadedModule{
		Name:        name,
		BaseAddress: imageBase,
		EntryPoint:  imageBase + addressOfEntryPoint,
	}

	for _, section := range f.Sections {
		e := env.loadPESection(&mod, section)
		check(e)
	}

	// since we always map at ImageBase, we don't need to apply (32bit) relocs
	// TODO: check 64bit reloc types

	importDirectory := dataDirectory[1]
	importRva := importDirectory.VirtualAddress
	importSize := importDirectory.Size
	fmt.Printf("import rva: 0x%x\n", importRva)
	fmt.Printf("import size: 0x%x\n", importSize)

	// TODO: check ranges
	d, e := mod.MemRead(env, uint64(importDirectory.VirtualAddress), uint64(importDirectory.Size))
	check(e)

	p := bytes.NewBuffer(d)
	for {
		var dir ImageImportDirectory
		binary.Read(p, binary.LittleEndian, &dir.rvaImportLookupTable)
		fmt.Printf("rva import lookup table: 0x%x\n", dir.rvaImportLookupTable)
		if dir.rvaImportLookupTable == 0 {
			break
		}
		binary.Read(p, binary.LittleEndian, &dir.TimeDateStamp)
		fmt.Printf("time date stamp: 0x%x\n", dir.TimeDateStamp)

		binary.Read(p, binary.LittleEndian, &dir.ForwarderChain)
		fmt.Printf("forwarder chain: 0x%x\n", dir.ForwarderChain)

		binary.Read(p, binary.LittleEndian, &dir.rvaModuleName)
		fmt.Printf("rva module name: 0x%x\n", dir.rvaModuleName)

		moduleNameBuf, e := mod.MemRead(env, uint64(dir.rvaModuleName), 0x100)
		check(e)
		moduleName, e := readAscii(moduleNameBuf)
		check(e)

		fmt.Printf("module name: %s\n", string(moduleName))

		binary.Read(p, binary.LittleEndian, &dir.rvaImportAddressTable)
	}

	env.loadedModules = append(env.loadedModules, mod)
	return &mod, nil
}

func readAscii(buf []byte) (string, error) {
	br := bufio.NewReader(bytes.NewReader(buf))
	bytez, e := br.ReadBytes(byte(0x00))
	check(e)
	return string(bytez), nil
}

func readUtf16le(buf []byte) ([]rune, error) {
	start := 0
	end := 0

	for i := 0; i < len(buf)/2; i++ {
		if !(buf[i*2] == 0 && buf[i*2+1] == 0) {
			end = i * 2
			break
		}
	}

	numChars := (end - start) / 2
	d := make([]uint16, numChars)
	for i := 0; i < numChars; i++ {
		d[i] = uint16(buf[i*2]) | (uint16(buf[i*2+1]) << 8)
	}
	return utf16.Decode(d), nil
}

var GAPSTONE_ARCH_MAP = map[Arch]int{
	ARCH_X86: gapstone.CS_ARCH_X86,
}

var GAPSTONE_MODE_MAP = map[Mode]uint{
	MODE_32: gapstone.CS_MODE_32,
}

func (env *Environment) disassembleBytes(data []byte, address uint64, w io.Writer) error {
	// TODO: cache the engine on the Environment?

	engine, e := gapstone.New(
		GAPSTONE_ARCH_MAP[env.Arch],
		GAPSTONE_MODE_MAP[env.Mode],
	)
	check(e)
	defer engine.Close()

	insns, e := engine.Disasm([]byte(data), address, 0 /* all instructions */)
	check(e)

	w.Write([]byte(fmt.Sprintf("Disasm:\n")))
	for _, insn := range insns {
		w.Write([]byte(fmt.Sprintf("0x%x:\t%s\t\t%s\n", insn.Address, insn.Mnemonic, insn.OpStr)))
	}

	return nil
}

func (env *Environment) Disassemble(address uint64, length uint64, w io.Writer) error {
	d, e := env.u.MemRead(address, length)
	check(e)
	return env.disassembleBytes(d, address, w)
}

func doit(path string) error {
	f, e := pe.Open(path)
	check(e)

	env, e := NewEnvironment(ARCH_X86, MODE_32)
	check(e)

	m, e := env.LoadPE(path, f)
	check(e)

	e = env.Disassemble(m.EntryPoint, 0x20, os.Stdout)
	check(e)

	return nil
}

func main() {
	app := cli.NewApp()
	app.Version = "0.1"
	app.Name = "crystal-tiger-pe-loader"
	app.Usage = "Load PE files."
	app.Commands = []cli.Command{
		{
			Name:    "load",
			Aliases: []string{"l"},
			Usage:   "load a pe file",
			Flags:   []cli.Flag{inputFlag, verboseFlag},
			Action: func(c *cli.Context) {
				if checkRequiredArgs(c, []cli.StringFlag{inputFlag}) != nil {
					return
				}
				check(doit(c.String("input_file")))
			},
		},
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package e2e

import (
	"testing"

	. "github.com/argoproj/argo-cd/pkg/apis/application/v1alpha1"
	. "github.com/argoproj/argo-cd/test/e2e/fixture"
	. "github.com/argoproj/argo-cd/test/e2e/fixture/app"
	"github.com/argoproj/argo-cd/test/fixture/test"
)

// when a app gets stuck in sync, and we try to delete it, it won't delete, instead we must then terminate it
// and deletion will then just happen
func TestDeletingAppStuckInSync(t *testing.T) {

	test.Flaky(t)

	Given(t).
		Async(true).
		Path("hook").
		When().
		PatchFile("hook.yaml", `[{"op": "replace", "path": "/spec/containers/0/command", "value": ["sh", "-c", "until ls /tmp/done; do sleep 0.1; done"]}]`).
		Create().
		Sync().
		Then().
		// stuck in running state
		Expect(OperationPhaseIs(OperationRunning)).
		When().
		Delete(true).
		Then().
		// delete is ignored, still stuck in running state
		Expect(OperationPhaseIs(OperationRunning)).
		When().
		TerminateOp().
		And(func() {
			// force delete the resource
			_, _ = Run("", "kubectl", "-n", DeploymentNamespace(), "exec", "-i", "hook", "touch", "/tmp/done")
		}).
		Then().
		// delete is successful
		Expect(DoesNotExist())
}
<commit_msg>Fixes TestDeletingAppStuckInSync. Closes #2087 (#2093)<commit_after>package e2e

import (
	"testing"

	"github.com/argoproj/argo-cd/errors"
	. "github.com/argoproj/argo-cd/pkg/apis/application/v1alpha1"
	. "github.com/argoproj/argo-cd/test/e2e/fixture"
	. "github.com/argoproj/argo-cd/test/e2e/fixture/app"
)

// when a app gets stuck in sync, and we try to delete it, it won't delete, instead we must then terminate it
// and deletion will then just happen
func TestDeletingAppStuckInSync(t *testing.T) {
	Given(t).
		Async(true).
		Path("hook").
		When().
		PatchFile("hook.yaml", `[{"op": "replace", "path": "/spec/containers/0/command", "value": ["sh", "-c", "until ls /tmp/done; do sleep 0.1; done"]}]`).
		Create().
		Sync().
		Then().
		// stuck in running state
		Expect(OperationPhaseIs(OperationRunning)).
		Expect(SyncStatusIs(SyncStatusCodeSynced)).
		Expect(ResourceResultNumbering(2)).
		When().
		Delete(true).
		Then().
		// delete is ignored, still stuck in running state
		Expect(OperationPhaseIs(OperationRunning)).
		When().
		TerminateOp().
		And(func() {
			// force delete the resource
			errors.FailOnErr(Run("", "kubectl", "-n", DeploymentNamespace(), "exec", "-i", "hook", "touch", "/tmp/done"))
		}).
		Then().
		// delete is successful
		Expect(DoesNotExist())
}
<|endoftext|>


<commit_before>package local

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/yinqiwen/gsnova/common/channel"
	"github.com/yinqiwen/gsnova/common/dns"
	"github.com/yinqiwen/gsnova/common/helper"
	"github.com/yinqiwen/gsnova/common/hosts"
	"github.com/yinqiwen/gsnova/common/logger"
	"github.com/yinqiwen/gsnova/common/mux"
	"github.com/yinqiwen/gsnova/common/socks"
	"github.com/yinqiwen/pmux"
)

var proxyServerRunning = true
var runningProxyStreamCount int64
var activeStreams sync.Map

func isTimeoutErr(err error) bool {
	if err == pmux.ErrTimeout {
		return true
	}
	if err, ok := err.(net.Error); ok && err.Timeout() {
		return true
	}
	return false
}

type proxyStreamContext struct {
	stream mux.MuxStream
	c      io.ReadWriteCloser
}

func serveProxyConn(conn net.Conn, remoteHost, remotePort string, proxy *ProxyConfig) {
	var proxyChannelName string
	protocol := "tcp"
	localConn := conn
	atomic.AddInt64(&runningProxyStreamCount, 1)
	//runningProxyConns.Store(conn, true)
	defer localConn.Close()
	defer atomic.AddInt64(&runningProxyStreamCount, -1)
	//defer runningProxyConns.Delete(conn)

	isSocksProxy := false
	isHttpsProxy := false
	isHttp11Proto := false
	isTransparentProxy := len(remoteHost) > 0
	var initialHTTPReq *http.Request

	var bufconn *bufio.Reader
	if !isTransparentProxy {
		socksConn, sbufconn, err := socks.NewSocksConn(conn)
		if nil == err {
			isSocksProxy = true
			logger.Debug("Local proxy recv %s proxy conn to %s", socksConn.Version(), socksConn.Req.Target)
			socksConn.Grant(&net.TCPAddr{
				IP: net.ParseIP("0.0.0.0"), Port: 0})
			localConn = socksConn
			if socksConn.Req.Target == GConf.UDPGW.Addr {
				logger.Debug("Handle udpgw conn for %v", socksConn.Req.Target)
				handleUDPGatewayConn(localConn, proxy)
				return
			}

			remoteHost, remotePort, err = net.SplitHostPort(socksConn.Req.Target)
			if nil != err {
				logger.Error("Invalid socks target addresss:%s with reason %v", socksConn.Req.Target, err)
				return
			}
			bufconn = sbufconn
		} else {
			if nil == sbufconn {
				localConn.Close()
				return
			}
			bufconn = sbufconn
		}
	}

	if nil == bufconn {
		bufconn = bufio.NewReader(localConn)
	}

	trySniffDomain := false
	if len(remoteHost) == 0 || (net.ParseIP(remoteHost) != nil && !helper.IsPrivateIP(remoteHost)) {
		//this is a ip from local dns query
		trySniffDomain = true
	}

	//1. sniff SNI first
	if (isSocksProxy || isTransparentProxy) && trySniffDomain {
		if remotePort == "80" {
			conn.SetReadDeadline(time.Now().Add(10 * time.Millisecond))
		} else {
			conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
		}

		if nil != dns.CNIPSet {
			remoteIP := net.ParseIP(remoteHost)
			logger.Debug("Recv proxy request to IP:%v CNIP:%v", remoteIP, dns.CNIPSet.IsInCountry(remoteIP, "CN"))
		}
		sni, err := helper.PeekTLSServerName(bufconn)
		if nil != err {
			//logger.Debug("##Failed to sniff SNI with error:%v", err)
		} else {
			if redirect, ok := GConf.SNI.redirect(sni); ok {
				sni = redirect
			}
			logger.Debug("Sniffed SNI:%s:%s for IP:%s:%s", sni, remotePort, remoteHost, remotePort)
			remoteHost = sni
			trySniffDomain = false
		}
	}

	//2. sniff domain via http
	if trySniffDomain {
		localConn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
		headChunk, err := bufconn.Peek(7)
		if len(headChunk) != 7 {
			if err != io.EOF {
				logger.Error("Peek:%s %d %v to %s:%s", string(headChunk), len(headChunk), err, remoteHost, remotePort)
			}
			goto START
		}
		method := string(headChunk)
		if tmp := strings.Fields(method); len(tmp) > 0 {
			method = tmp[0]
		}
		method = strings.ToUpper(method)
		switch method {
		case "GET":
			fallthrough
		case "POST":
			fallthrough
		case "HEAD":
			fallthrough
		case "PUT":
			fallthrough
		case "DELETE":
			fallthrough
		case "CONNECT":
			fallthrough
		case "OPTIONS":
			fallthrough
		case "TRACE":
			fallthrough
		case "PATCH":
			isHttp11Proto = true
		default:
			isHttp11Proto = false
		}
		//log.Printf("[%d]]Method:%s", sid, method)
		if isHttp11Proto {
			initialHTTPReq, err = http.ReadRequest(bufconn)
			if nil != err {
				logger.Error("Read first request failed from proxy connection for reason:%v", err)
				return
			}
			//log.Printf("Host:%s %v", initialHTTPReq.Host, initialHTTPReq.URL)
			if strings.Contains(initialHTTPReq.Host, ":") {
				remoteHost, remotePort, _ = net.SplitHostPort(initialHTTPReq.Host)
			} else {
				remoteHost = initialHTTPReq.Host
				remotePort = "80"
				if strings.EqualFold(initialHTTPReq.Method, "CONNECT") {
					remotePort = "443"
				}
			}
			if strings.EqualFold(initialHTTPReq.Method, "CONNECT") {
				protocol = "https"
				if !isSocksProxy {
					localConn.Write([]byte("HTTP/1.0 200 Connection established\r\n\r\n"))
					isHttpsProxy = true
					initialHTTPReq = nil
				}
			} else {
				protocol = "http"
			}
		} else {
			if !isHttp11Proto && !isSocksProxy && !isTransparentProxy {
				logger.Error("[ERROR]Can NOT handle non HTTP1.1 proto in non socks proxy mode %v:%v:%v  %v to %v", isHttp11Proto, isSocksProxy, isTransparentProxy, conn.LocalAddr(), conn.RemoteAddr())
				return
			}
		}
	}
START:
	if len(remoteHost) == 0 || len(remotePort) == 0 {
		logger.Error("Can NOT resolve remote host or port %s:%s %v", remoteHost, remotePort, initialHTTPReq)
		return
	}
	proxyChannelName = proxy.getProxyChannelByHost(protocol, remoteHost)

	if len(proxyChannelName) == 0 {
		logger.Error("[ERROR]No proxy found for %s:%s", protocol, remoteHost)
		return
	}
	stream, conf, err := channel.GetMuxStreamByChannel(proxyChannelName)
	if nil != err || nil == stream {
		logger.Error("Failed to open stream for reason:%v by proxy:%s", err, proxyChannelName)
		return
	}
	defer stream.Close()

	ssid := stream.StreamID()
	opt := mux.StreamOptions{
		DialTimeout: conf.RemoteDialMSTimeout,
		Hops:        conf.Hops,
	}

	if remotePort == "443" && nil == net.ParseIP(remoteHost) {
		remoteSNI := conf.GetRemoteSNI(remoteHost)
		if len(remoteSNI) > 0 {
			sniHost := hosts.GetHost(remoteSNI)
			logger.Notice("Proxy stream[%d] select remote SNI host %s for proxy to %s:%s", ssid, sniHost, remoteHost, remotePort)
			remoteHost = sniHost
		}
	}

	logger.Notice("Proxy stream[%d] select %s for proxy to %s:%s", ssid, proxyChannelName, remoteHost, remotePort)
	err = stream.Connect("tcp", net.JoinHostPort(remoteHost, remotePort), opt)
	if nil != err {
		logger.Error("Connect failed from proxy connection for reason:%v", err)
		return
	}

	//clear read timeout
	var zero time.Time
	localConn.SetReadDeadline(zero)
	streamReader, streamWriter := mux.GetCompressStreamReaderWriter(stream, conf.Compressor)

	streamCtx := &proxyStreamContext{}
	streamCtx.stream = stream
	streamCtx.c = localConn
	activeStreams.Store(streamCtx, true)

	maxIdleTime := time.Duration(GConf.Mux.StreamIdleTimeout) * time.Second
	if maxIdleTime == 0 {
		maxIdleTime = 10 * time.Second
	}
	closeCh := make(chan int, 1)
	go func() {
		buf := make([]byte, 128*1024)
		for {
			localConn.SetReadDeadline(time.Now().Add(maxIdleTime))
			_, err := io.CopyBuffer(localConn, streamReader, buf)
			if isTimeoutErr(err) && time.Now().Sub(stream.LatestIOTime()) < maxIdleTime {
				continue
			}
		}
		closeCh <- 1
	}()

	//start task to check stream timeout(if the stream has no read&write action more than 10s)

	if (isSocksProxy || isHttpsProxy || isTransparentProxy) && nil == initialHTTPReq {
		buf := make([]byte, 128*1024)
		io.CopyBuffer(streamWriter, bufconn, buf)
		if close, ok := streamWriter.(io.Closer); ok {
			close.Close()
		}
	} else {
		proxyReq := initialHTTPReq
		initialHTTPReq = nil
		for {
			if nil != proxyReq {
				proxyReq.Header.Del("Proxy-Connection")
				proxyReq.Header.Del("Proxy-Authorization")
				err = proxyReq.Write(streamWriter)
				if nil != err {
					logger.Error("Failed to write http request for reason:%v", err)
					return
				}
			}
			prevReq := proxyReq
			//localConn.SetReadDeadline(time.Now().Add(5 * time.Second))
			proxyReq, err = http.ReadRequest(bufconn)
			if nil != err {
				if err, ok := err.(net.Error); ok && err.Timeout() {

				}
				if err != io.EOF && !strings.Contains(err.Error(), "use of closed network connection") {
					logger.Notice("Failed to read proxy http request to %s:%s for reason:%v", remoteHost, remotePort, err)
				}
				return
			}
			if nil != prevReq && prevReq.Host != proxyReq.Host {
				logger.Debug("Switch to next stream since target host change from %s to %s", prevReq.Host, proxyReq.Host)
				stream.Close()
				goto START
			}
		}
	}
	<-closeCh
	activeStreams.Delete(streamCtx)
}

func startLocalProxyServer(proxyIdx int) (*net.TCPListener, error) {
	proxyConf := &GConf.Proxy[proxyIdx]
	if supportTransparentProxy() {
		go startTransparentUDProxy(proxyConf.Local, proxyConf)
	}
	tcpaddr, err := net.ResolveTCPAddr("tcp", proxyConf.Local)
	if nil != err {
		logger.Fatal("[ERROR]Local server address:%s error:%v", proxyConf.Local, err)
		return nil, err
	}
	var lp *net.TCPListener
	lp, err = net.ListenTCP("tcp", tcpaddr)
	if nil != err {
		logger.Fatal("Can NOT listen on address:%s", proxyConf.Local)
		return nil, err
	}
	logger.Info("Listen on address %s", proxyConf.Local)
	runningServers[proxyIdx] = lp
	go func() {
		for proxyServerRunning {
			var conn net.Conn
			conn, err = lp.AcceptTCP()
			if nil != err {
				continue
			}
			isTransparent := false
			if supportTransparentProxy() {
				tcpAddr, ok := conn.RemoteAddr().(*net.TCPAddr)
				if ok {
					_, exist := helper.GetLocalIPSet()[tcpAddr.IP.String()]
					if !exist {
						isTransparent = true
					}
				}
			}
			var originalHost, originalPort string
			if isTransparent {
				newConn, remoteIP, remotePort, err := getOrinalTCPRemoteAddr(conn)
				if nil != err {
					logger.Error("Failed to get original address for transparent proxy:%v", err)
					continue
				}
				conn = newConn
				originalHost = remoteIP.String()
				originalPort = fmt.Sprintf("%d", remotePort)
			}
			go serveProxyConn(conn, originalHost, originalPort, &GConf.Proxy[proxyIdx])
		}
		lp.Close()
	}()
	return lp, nil
}

var runningServers []*net.TCPListener

func startLocalServers() error {
	proxyServerRunning = true
	runningServers = make([]*net.TCPListener, len(GConf.Proxy))
	for i := range GConf.Proxy {
		startLocalProxyServer(i)
	}
	return nil
}

func stopLocalServers() {
	proxyServerRunning = false
	for _, l := range runningServers {
		if nil != l {
			l.Close()
		}
	}
	//closeAllProxySession()
	closeAllUDPSession()
	activeStreams.Range(func(key, value interface{}) bool {
		ctx := key.(proxyStreamContext)
		if nil != ctx.c {
			ctx.c.Close()
		}
		return true
	})
}
<commit_msg>use less tiker/timer<commit_after>package local

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/yinqiwen/gsnova/common/channel"
	"github.com/yinqiwen/gsnova/common/dns"
	"github.com/yinqiwen/gsnova/common/helper"
	"github.com/yinqiwen/gsnova/common/hosts"
	"github.com/yinqiwen/gsnova/common/logger"
	"github.com/yinqiwen/gsnova/common/mux"
	"github.com/yinqiwen/gsnova/common/socks"
	"github.com/yinqiwen/pmux"
)

var proxyServerRunning = true
var runningProxyStreamCount int64
var activeStreams sync.Map

func isTimeoutErr(err error) bool {
	if err == pmux.ErrTimeout {
		return true
	}
	if err, ok := err.(net.Error); ok && err.Timeout() {
		return true
	}
	return false
}

type proxyStreamContext struct {
	stream mux.MuxStream
	c      io.ReadWriteCloser
}

func serveProxyConn(conn net.Conn, remoteHost, remotePort string, proxy *ProxyConfig) {
	var proxyChannelName string
	protocol := "tcp"
	localConn := conn
	atomic.AddInt64(&runningProxyStreamCount, 1)
	//runningProxyConns.Store(conn, true)
	defer localConn.Close()
	defer atomic.AddInt64(&runningProxyStreamCount, -1)
	//defer runningProxyConns.Delete(conn)

	isSocksProxy := false
	isHttpsProxy := false
	isHttp11Proto := false
	isTransparentProxy := len(remoteHost) > 0
	var initialHTTPReq *http.Request

	var bufconn *bufio.Reader
	if !isTransparentProxy {
		socksConn, sbufconn, err := socks.NewSocksConn(conn)
		if nil == err {
			isSocksProxy = true
			logger.Debug("Local proxy recv %s proxy conn to %s", socksConn.Version(), socksConn.Req.Target)
			socksConn.Grant(&net.TCPAddr{
				IP: net.ParseIP("0.0.0.0"), Port: 0})
			localConn = socksConn
			if socksConn.Req.Target == GConf.UDPGW.Addr {
				logger.Debug("Handle udpgw conn for %v", socksConn.Req.Target)
				handleUDPGatewayConn(localConn, proxy)
				return
			}

			remoteHost, remotePort, err = net.SplitHostPort(socksConn.Req.Target)
			if nil != err {
				logger.Error("Invalid socks target addresss:%s with reason %v", socksConn.Req.Target, err)
				return
			}
			bufconn = sbufconn
		} else {
			if nil == sbufconn {
				localConn.Close()
				return
			}
			bufconn = sbufconn
		}
	}

	if nil == bufconn {
		bufconn = bufio.NewReader(localConn)
	}

	trySniffDomain := false
	if len(remoteHost) == 0 || (net.ParseIP(remoteHost) != nil && !helper.IsPrivateIP(remoteHost)) {
		//this is a ip from local dns query
		trySniffDomain = true
	}

	//1. sniff SNI first
	if (isSocksProxy || isTransparentProxy) && trySniffDomain {
		if remotePort == "80" {
			conn.SetReadDeadline(time.Now().Add(10 * time.Millisecond))
		} else {
			conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
		}

		if nil != dns.CNIPSet {
			remoteIP := net.ParseIP(remoteHost)
			logger.Debug("Recv proxy request to IP:%v CNIP:%v", remoteIP, dns.CNIPSet.IsInCountry(remoteIP, "CN"))
		}
		sni, err := helper.PeekTLSServerName(bufconn)
		if nil != err {
			//logger.Debug("##Failed to sniff SNI with error:%v", err)
		} else {
			if redirect, ok := GConf.SNI.redirect(sni); ok {
				sni = redirect
			}
			logger.Debug("Sniffed SNI:%s:%s for IP:%s:%s", sni, remotePort, remoteHost, remotePort)
			remoteHost = sni
			trySniffDomain = false
		}
	}

	//2. sniff domain via http
	if trySniffDomain {
		localConn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
		headChunk, err := bufconn.Peek(7)
		if len(headChunk) != 7 {
			if err != io.EOF {
				logger.Error("Peek:%s %d %v to %s:%s", string(headChunk), len(headChunk), err, remoteHost, remotePort)
			}
			goto START
		}
		method := string(headChunk)
		if tmp := strings.Fields(method); len(tmp) > 0 {
			method = tmp[0]
		}
		method = strings.ToUpper(method)
		switch method {
		case "GET":
			fallthrough
		case "POST":
			fallthrough
		case "HEAD":
			fallthrough
		case "PUT":
			fallthrough
		case "DELETE":
			fallthrough
		case "CONNECT":
			fallthrough
		case "OPTIONS":
			fallthrough
		case "TRACE":
			fallthrough
		case "PATCH":
			isHttp11Proto = true
		default:
			isHttp11Proto = false
		}
		//log.Printf("[%d]]Method:%s", sid, method)
		if isHttp11Proto {
			initialHTTPReq, err = http.ReadRequest(bufconn)
			if nil != err {
				logger.Error("Read first request failed from proxy connection for reason:%v", err)
				return
			}
			//log.Printf("Host:%s %v", initialHTTPReq.Host, initialHTTPReq.URL)
			if strings.Contains(initialHTTPReq.Host, ":") {
				remoteHost, remotePort, _ = net.SplitHostPort(initialHTTPReq.Host)
			} else {
				remoteHost = initialHTTPReq.Host
				remotePort = "80"
				if strings.EqualFold(initialHTTPReq.Method, "CONNECT") {
					remotePort = "443"
				}
			}
			if strings.EqualFold(initialHTTPReq.Method, "CONNECT") {
				protocol = "https"
				if !isSocksProxy {
					localConn.Write([]byte("HTTP/1.0 200 Connection established\r\n\r\n"))
					isHttpsProxy = true
					initialHTTPReq = nil
				}
			} else {
				protocol = "http"
			}
		} else {
			if !isHttp11Proto && !isSocksProxy && !isTransparentProxy {
				logger.Error("[ERROR]Can NOT handle non HTTP1.1 proto in non socks proxy mode %v:%v:%v  %v to %v", isHttp11Proto, isSocksProxy, isTransparentProxy, conn.LocalAddr(), conn.RemoteAddr())
				return
			}
		}
	}
START:
	if len(remoteHost) == 0 || len(remotePort) == 0 {
		logger.Error("Can NOT resolve remote host or port %s:%s %v", remoteHost, remotePort, initialHTTPReq)
		return
	}
	proxyChannelName = proxy.getProxyChannelByHost(protocol, remoteHost)

	if len(proxyChannelName) == 0 {
		logger.Error("[ERROR]No proxy found for %s:%s", protocol, remoteHost)
		return
	}
	stream, conf, err := channel.GetMuxStreamByChannel(proxyChannelName)
	if nil != err || nil == stream {
		logger.Error("Failed to open stream for reason:%v by proxy:%s", err, proxyChannelName)
		return
	}
	defer stream.Close()

	ssid := stream.StreamID()
	opt := mux.StreamOptions{
		DialTimeout: conf.RemoteDialMSTimeout,
		Hops:        conf.Hops,
	}

	if remotePort == "443" && nil == net.ParseIP(remoteHost) {
		remoteSNI := conf.GetRemoteSNI(remoteHost)
		if len(remoteSNI) > 0 {
			sniHost := hosts.GetHost(remoteSNI)
			logger.Notice("Proxy stream[%d] select remote SNI host %s for proxy to %s:%s", ssid, sniHost, remoteHost, remotePort)
			remoteHost = sniHost
		}
	}

	logger.Notice("Proxy stream[%d] select %s for proxy to %s:%s", ssid, proxyChannelName, remoteHost, remotePort)
	err = stream.Connect("tcp", net.JoinHostPort(remoteHost, remotePort), opt)
	if nil != err {
		logger.Error("Connect failed from proxy connection for reason:%v", err)
		return
	}

	//clear read timeout
	var zero time.Time
	localConn.SetReadDeadline(zero)
	streamReader, streamWriter := mux.GetCompressStreamReaderWriter(stream, conf.Compressor)

	streamCtx := &proxyStreamContext{}
	streamCtx.stream = stream
	streamCtx.c = localConn
	activeStreams.Store(streamCtx, true)

	maxIdleTime := time.Duration(GConf.Mux.StreamIdleTimeout) * time.Second
	if maxIdleTime == 0 {
		maxIdleTime = 10 * time.Second
	}
	closeCh := make(chan int, 1)
	go func() {
		buf := make([]byte, 128*1024)
		io.CopyBuffer(localConn, streamReader, buf)
		closeCh <- 1
	}()

	//start task to check stream timeout(if the stream has no read&write action more than 10s)

	if (isSocksProxy || isHttpsProxy || isTransparentProxy) && nil == initialHTTPReq {
		buf := make([]byte, 128*1024)
		for {
			localConn.SetReadDeadline(time.Now().Add(maxIdleTime))
			_, cerr := io.CopyBuffer(streamWriter, bufconn, buf)
			if isTimeoutErr(cerr) && time.Now().Sub(stream.LatestIOTime()) < maxIdleTime {
				continue
			}
		}

		if close, ok := streamWriter.(io.Closer); ok {
			close.Close()
		}
	} else {
		proxyReq := initialHTTPReq
		initialHTTPReq = nil
		for {
			if nil != proxyReq {
				proxyReq.Header.Del("Proxy-Connection")
				proxyReq.Header.Del("Proxy-Authorization")
				err = proxyReq.Write(streamWriter)
				if nil != err {
					logger.Error("Failed to write http request for reason:%v", err)
					return
				}
			}
			prevReq := proxyReq
			for {
				localConn.SetReadDeadline(time.Now().Add(maxIdleTime))
				proxyReq, err = http.ReadRequest(bufconn)
				if isTimeoutErr(err) && time.Now().Sub(stream.LatestIOTime()) < maxIdleTime {
					continue
				}
				break
			}

			if nil != err {
				if err, ok := err.(net.Error); ok && err.Timeout() {

				}
				if err != io.EOF && !strings.Contains(err.Error(), "use of closed network connection") {
					logger.Notice("Failed to read proxy http request to %s:%s for reason:%v", remoteHost, remotePort, err)
				}
				return
			}
			if nil != prevReq && prevReq.Host != proxyReq.Host {
				logger.Debug("Switch to next stream since target host change from %s to %s", prevReq.Host, proxyReq.Host)
				stream.Close()
				goto START
			}
		}
	}
	<-closeCh
	activeStreams.Delete(streamCtx)
}

func startLocalProxyServer(proxyIdx int) (*net.TCPListener, error) {
	proxyConf := &GConf.Proxy[proxyIdx]
	if supportTransparentProxy() {
		go startTransparentUDProxy(proxyConf.Local, proxyConf)
	}
	tcpaddr, err := net.ResolveTCPAddr("tcp", proxyConf.Local)
	if nil != err {
		logger.Fatal("[ERROR]Local server address:%s error:%v", proxyConf.Local, err)
		return nil, err
	}
	var lp *net.TCPListener
	lp, err = net.ListenTCP("tcp", tcpaddr)
	if nil != err {
		logger.Fatal("Can NOT listen on address:%s", proxyConf.Local)
		return nil, err
	}
	logger.Info("Listen on address %s", proxyConf.Local)
	runningServers[proxyIdx] = lp
	go func() {
		for proxyServerRunning {
			var conn net.Conn
			conn, err = lp.AcceptTCP()
			if nil != err {
				continue
			}
			isTransparent := false
			if supportTransparentProxy() {
				tcpAddr, ok := conn.RemoteAddr().(*net.TCPAddr)
				if ok {
					_, exist := helper.GetLocalIPSet()[tcpAddr.IP.String()]
					if !exist {
						isTransparent = true
					}
				}
			}
			var originalHost, originalPort string
			if isTransparent {
				newConn, remoteIP, remotePort, err := getOrinalTCPRemoteAddr(conn)
				if nil != err {
					logger.Error("Failed to get original address for transparent proxy:%v", err)
					continue
				}
				conn = newConn
				originalHost = remoteIP.String()
				originalPort = fmt.Sprintf("%d", remotePort)
			}
			go serveProxyConn(conn, originalHost, originalPort, &GConf.Proxy[proxyIdx])
		}
		lp.Close()
	}()
	return lp, nil
}

var runningServers []*net.TCPListener

func startLocalServers() error {
	proxyServerRunning = true
	runningServers = make([]*net.TCPListener, len(GConf.Proxy))
	for i := range GConf.Proxy {
		startLocalProxyServer(i)
	}
	return nil
}

func stopLocalServers() {
	proxyServerRunning = false
	for _, l := range runningServers {
		if nil != l {
			l.Close()
		}
	}
	//closeAllProxySession()
	closeAllUDPSession()
	activeStreams.Range(func(key, value interface{}) bool {
		ctx := key.(proxyStreamContext)
		if nil != ctx.c {
			ctx.c.Close()
		}
		return true
	})
}
<|endoftext|>