<commit_before>package clui

import (
	term "github.com/nsf/termbox-go"
)

// ThumbPosition returns a scrollbar thumb position depending
// on currently active item(itemNo), total number of items
// (itemCount), and length/height of the scrollbar(length)
// including arrows. Returns position in interval of (1..lenght-2)
// or -1 if the thumb is not visible
func ThumbPosition(itemNo, itemCount, length int) int {
	length -= 2
	if itemNo < 0 {
		return -1
	}
	if itemNo >= itemCount-1 {
		return length - 1
	}

	if length < 4 {
		return 0
	}

	ydiff := int(float32(itemNo) / float32(itemCount-1.0) * float32(length-1))
	return ydiff
}

// ItemByThumbPosition calculates item number by scrollbar
// thumb position. Position - thumb position inside scrollbar,
// itemCount - total number of items, lenght - lenght or heigth
// of scrollbar. Return -1 if it is not possible to calculate:
// e.g, itemCount equals zero
func ItemByThumbPosition(position, itemCount, length int) int {
	length -= 2
	if position < 1 {
		return -1
	}
	if itemCount < 1 {
		return -1
	}
	if itemCount == 1 {
		return 0
	}

	newPos := int(float32(itemCount-1)*float32(position-1)/float32(length-1) + 0.9)

	if newPos < 0 {
		newPos = 0
	} else if newPos >= itemCount {
		newPos = itemCount - 1
	}

	return newPos
}

// ChildAt returns the children of parent control that is at absolute
// coordinates x, y. Returns nil if x, y are outside parent control and
// returns parent if no child is at x, y
func ChildAt(parent Control, x, y int) Control {
	px, py := parent.Pos()
	pw, ph := parent.Size()
	if px > x || py > y || px+pw <= x || py+ph <= y {
		return nil
	}

	if len(parent.Children()) == 0 {
		return parent
	}

	var ctrl Control
	ctrl = parent
	for _, child := range parent.Children() {
		check := ChildAt(child, x, y)
		if check != nil {
			ctrl = check
			break
		}
	}

	return ctrl
}

// DeactivateControls makes all children of parent inactive
func DeactivateControls(parent Control) {
	for _, ctrl := range parent.Children() {
		if ctrl.Active() {
			ctrl.SetActive(false)
			ctrl.ProcessEvent(Event{Type: EventActivate, X: 0})
		}

		DeactivateControls(ctrl)
	}
}

// ActivateControl makes control active and disables all other children of
// the parent. Returns true if control was found and activated
func ActivateControl(parent, control Control) bool {
	DeactivateControls(parent)
	res := false
	ctrl := FindChild(parent, control)
	if ctrl != nil {
		res = true
		if !ctrl.Active() {
			ctrl.ProcessEvent(Event{Type: EventActivate, X: 1})
			ctrl.SetActive(true)
		}
	}

	return res
}

// FindChild returns control if it is a child of the parent and nil otherwise
func FindChild(parent, control Control) Control {
	var res Control

	if parent == control {
		return parent
	}

	for _, ctrl := range parent.Children() {
		if ctrl == control {
			res = ctrl
			break
		}

		res = FindChild(ctrl, control)
		if res != nil {
			break
		}
	}

	return res
}

// IsMouseClickEvent returns if a user action can be treated as mouse click.
func IsMouseClickEvent(ev Event) bool {
	if ev.Type == EventClick {
		return true
	}
	if ev.Type == EventMouse && ev.Key == term.MouseLeft {
		return true
	}

	return false
}

// FindFirstControl returns the first child for that fn returns true.
// The function is used to find active or tab-stop control
func FindFirstControl(parent Control, fn func(Control) bool) Control {
	linear := getLinearControlList(parent, fn)
	if len(linear) == 0 {
		return nil
	}

	return linear[0]
}

// FindLastControl returns the first child for that fn returns true.
// The function is used by TAB processing method if a user goes backwards
// with TAB key - not supported now
func FindLastControl(parent Control, fn func(Control) bool) Control {
	linear := getLinearControlList(parent, fn)

	if len(linear) == 0 {
		return nil
	}

	return linear[len(linear)-1]
}

// ActiveControl returns the active child of the parent or nil if no child is
// active
func ActiveControl(parent Control) Control {
	fnActive := func(c Control) bool {
		return c.Active()
	}
	return FindFirstControl(parent, fnActive)
}

// FindFirstActiveControl returns the first active control of a parent
func FindFirstActiveControl(parent Control) Control {
	for _, curr := range getLinearControlList(parent, nil) {
		if curr.Active() {
			return curr
		}
	}
	return nil
}

func getLinearControlList(parent Control, fn func(Control) bool) []Control {
	result := []Control{}

	for _, curr := range parent.Children() {
		if fn != nil && fn(curr) {
			result = append(result, curr)
		}

		if len(curr.Children()) == 0 {
			continue
		}

		ch := getLinearControlList(curr, fn)
		if len(ch) != 0 {
			result = append(result, ch...)
		}
	}

	return result
}

// NextControl returns the next or previous child (depends on next parameter)
// that has tab-stop feature on. Used by library when processing TAB key
func NextControl(parent Control, curr Control, next bool) Control {
	fnTab := func(c Control) bool {
		return c.TabStop() && c.Visible() && c.Enabled()
	}

	linear := getLinearControlList(parent, fnTab)

	if len(linear) == 0 {
		return nil
	}

	var pIndex, nIndex int

	for i, ch := range linear {
		if ch != curr {
			continue
		}

		pIndex = i - 1
		nIndex = i + 1
		break
	}

	if nIndex > len(linear)-1 {
		nIndex = 0
	}

	if pIndex < 0 {
		pIndex = len(linear) - 1
	}

	if next {
		return linear[nIndex]
	} else {
		return linear[pIndex]
	}
}

// SendEventToChild tries to find a child control that should recieve the evetn
// For mouse click events it looks for a control at coordinates of event,
// makes it active, and then sends the event to it.
// If it is not mouse click event then it looks for the first active child and
// sends the event to it if it is not nil
func SendEventToChild(parent Control, ev Event) bool {
	var child Control
	if IsMouseClickEvent(ev) {
		child = ChildAt(parent, ev.X, ev.Y)
		if child != nil && !child.Active() {
			ActivateControl(parent, child)
		}
	} else {
		child = ActiveControl(parent)
	}

	if child != nil && child != parent {
		ev.Target = child
		res := child.ProcessEvent(ev)

		if cparent := ClippedParent(child); cparent != nil && cparent != child {
			cparent.ProcessEvent(ev)
		}

		return res
	}

	return false
}

// CalcClipper calculates the clipper size based on the control's size, position
// and paddings
func CalcClipper(c Control) (int, int, int, int) {
	w, h := c.Size()
	x, y := c.Pos()
	px, py := c.Paddings()

	x = x + px
	y = y + py
	w = w - 2*px
	h = h - 2*py

	return x, y, w, h
}

// ClippedParent finds the first c parent with clipped flag
func ClippedParent(c Control) Control {
	var clipped Control

	ctrl := c.Parent()
	clipped = c

	for ctrl != nil {
		if ctrl.Clipped() {
			clipped = ctrl
			break
		}

		ctrl = ctrl.Parent()
	}

	return clipped
}

// ControlInRect returns true if c is within a given rect
func ControlInRect(c Control, x int, y int, w int, h int) bool {
	xx, yy := c.Pos()
	ww, hh := c.Size()

	return xx >= x && ww <= x+w && yy <= y+h &&
		yy+hh <= y+h && yy >= y && yy+h >= y
}
<commit_msg>Don't consider hidden children on ChildAt<commit_after>package clui

import (
	term "github.com/nsf/termbox-go"
)

// ThumbPosition returns a scrollbar thumb position depending
// on currently active item(itemNo), total number of items
// (itemCount), and length/height of the scrollbar(length)
// including arrows. Returns position in interval of (1..lenght-2)
// or -1 if the thumb is not visible
func ThumbPosition(itemNo, itemCount, length int) int {
	length -= 2
	if itemNo < 0 {
		return -1
	}
	if itemNo >= itemCount-1 {
		return length - 1
	}

	if length < 4 {
		return 0
	}

	ydiff := int(float32(itemNo) / float32(itemCount-1.0) * float32(length-1))
	return ydiff
}

// ItemByThumbPosition calculates item number by scrollbar
// thumb position. Position - thumb position inside scrollbar,
// itemCount - total number of items, lenght - lenght or heigth
// of scrollbar. Return -1 if it is not possible to calculate:
// e.g, itemCount equals zero
func ItemByThumbPosition(position, itemCount, length int) int {
	length -= 2
	if position < 1 {
		return -1
	}
	if itemCount < 1 {
		return -1
	}
	if itemCount == 1 {
		return 0
	}

	newPos := int(float32(itemCount-1)*float32(position-1)/float32(length-1) + 0.9)

	if newPos < 0 {
		newPos = 0
	} else if newPos >= itemCount {
		newPos = itemCount - 1
	}

	return newPos
}

// ChildAt returns the children of parent control that is at absolute
// coordinates x, y. Returns nil if x, y are outside parent control and
// returns parent if no child is at x, y
func ChildAt(parent Control, x, y int) Control {
	px, py := parent.Pos()
	pw, ph := parent.Size()
	if px > x || py > y || px+pw <= x || py+ph <= y {
		return nil
	}

	if len(parent.Children()) == 0 {
		return parent
	}

	var ctrl Control
	ctrl = parent
	for _, child := range parent.Children() {
		if !child.Visible() {
			continue
		}

		check := ChildAt(child, x, y)
		if check != nil {
			ctrl = check
			break
		}
	}

	return ctrl
}

// DeactivateControls makes all children of parent inactive
func DeactivateControls(parent Control) {
	for _, ctrl := range parent.Children() {
		if ctrl.Active() {
			ctrl.SetActive(false)
			ctrl.ProcessEvent(Event{Type: EventActivate, X: 0})
		}

		DeactivateControls(ctrl)
	}
}

// ActivateControl makes control active and disables all other children of
// the parent. Returns true if control was found and activated
func ActivateControl(parent, control Control) bool {
	DeactivateControls(parent)
	res := false
	ctrl := FindChild(parent, control)
	if ctrl != nil {
		res = true
		if !ctrl.Active() {
			ctrl.ProcessEvent(Event{Type: EventActivate, X: 1})
			ctrl.SetActive(true)
		}
	}

	return res
}

// FindChild returns control if it is a child of the parent and nil otherwise
func FindChild(parent, control Control) Control {
	var res Control

	if parent == control {
		return parent
	}

	for _, ctrl := range parent.Children() {
		if ctrl == control {
			res = ctrl
			break
		}

		res = FindChild(ctrl, control)
		if res != nil {
			break
		}
	}

	return res
}

// IsMouseClickEvent returns if a user action can be treated as mouse click.
func IsMouseClickEvent(ev Event) bool {
	if ev.Type == EventClick {
		return true
	}
	if ev.Type == EventMouse && ev.Key == term.MouseLeft {
		return true
	}

	return false
}

// FindFirstControl returns the first child for that fn returns true.
// The function is used to find active or tab-stop control
func FindFirstControl(parent Control, fn func(Control) bool) Control {
	linear := getLinearControlList(parent, fn)
	if len(linear) == 0 {
		return nil
	}

	return linear[0]
}

// FindLastControl returns the first child for that fn returns true.
// The function is used by TAB processing method if a user goes backwards
// with TAB key - not supported now
func FindLastControl(parent Control, fn func(Control) bool) Control {
	linear := getLinearControlList(parent, fn)

	if len(linear) == 0 {
		return nil
	}

	return linear[len(linear)-1]
}

// ActiveControl returns the active child of the parent or nil if no child is
// active
func ActiveControl(parent Control) Control {
	fnActive := func(c Control) bool {
		return c.Active()
	}
	return FindFirstControl(parent, fnActive)
}

// FindFirstActiveControl returns the first active control of a parent
func FindFirstActiveControl(parent Control) Control {
	for _, curr := range getLinearControlList(parent, nil) {
		if curr.Active() {
			return curr
		}
	}
	return nil
}

func getLinearControlList(parent Control, fn func(Control) bool) []Control {
	result := []Control{}

	for _, curr := range parent.Children() {
		if fn != nil && fn(curr) {
			result = append(result, curr)
		}

		if len(curr.Children()) == 0 {
			continue
		}

		ch := getLinearControlList(curr, fn)
		if len(ch) != 0 {
			result = append(result, ch...)
		}
	}

	return result
}

// NextControl returns the next or previous child (depends on next parameter)
// that has tab-stop feature on. Used by library when processing TAB key
func NextControl(parent Control, curr Control, next bool) Control {
	fnTab := func(c Control) bool {
		return c.TabStop() && c.Visible() && c.Enabled()
	}

	linear := getLinearControlList(parent, fnTab)

	if len(linear) == 0 {
		return nil
	}

	var pIndex, nIndex int

	for i, ch := range linear {
		if ch != curr {
			continue
		}

		pIndex = i - 1
		nIndex = i + 1
		break
	}

	if nIndex > len(linear)-1 {
		nIndex = 0
	}

	if pIndex < 0 {
		pIndex = len(linear) - 1
	}

	if next {
		return linear[nIndex]
	} else {
		return linear[pIndex]
	}
}

// SendEventToChild tries to find a child control that should recieve the evetn
// For mouse click events it looks for a control at coordinates of event,
// makes it active, and then sends the event to it.
// If it is not mouse click event then it looks for the first active child and
// sends the event to it if it is not nil
func SendEventToChild(parent Control, ev Event) bool {
	var child Control
	if IsMouseClickEvent(ev) {
		child = ChildAt(parent, ev.X, ev.Y)
		if child != nil && !child.Active() {
			ActivateControl(parent, child)
		}
	} else {
		child = ActiveControl(parent)
	}

	if child != nil && child != parent {
		ev.Target = child
		res := child.ProcessEvent(ev)

		if cparent := ClippedParent(child); cparent != nil && cparent != child {
			cparent.ProcessEvent(ev)
		}

		return res
	}

	return false
}

// CalcClipper calculates the clipper size based on the control's size, position
// and paddings
func CalcClipper(c Control) (int, int, int, int) {
	w, h := c.Size()
	x, y := c.Pos()
	px, py := c.Paddings()

	x = x + px
	y = y + py
	w = w - 2*px
	h = h - 2*py

	return x, y, w, h
}

// ClippedParent finds the first c parent with clipped flag
func ClippedParent(c Control) Control {
	var clipped Control

	ctrl := c.Parent()
	clipped = c

	for ctrl != nil {
		if ctrl.Clipped() {
			clipped = ctrl
			break
		}

		ctrl = ctrl.Parent()
	}

	return clipped
}

// ControlInRect returns true if c is within a given rect
func ControlInRect(c Control, x int, y int, w int, h int) bool {
	xx, yy := c.Pos()
	ww, hh := c.Size()

	return xx >= x && ww <= x+w && yy <= y+h &&
		yy+hh <= y+h && yy >= y && yy+h >= y
}
<|endoftext|>


<commit_before>package scrud

import (
	"errors"
	"reflect"
	"time"

	"github.com/cxr29/scrud/format"
	"github.com/cxr29/scrud/internal/table"
	. "github.com/cxr29/scrud/query"
)

func (db *DB) Snapshot() *Snapshot {
	return &Snapshot{xr: db}
}

func (tx *Tx) Snapshot() *Snapshot {
	return &Snapshot{xr: tx}
}

// snapshot manager
type Snapshot struct {
	xr faker
}

func (s *Snapshot) Insert(data interface{}) (int64, time.Time, error) {
	var zt time.Time

	v := reflect.ValueOf(data)
	t := v.Type()
	if v.Kind() == reflect.Ptr {
		if v.IsNil() {
			return 0, zt, errors.New("scrud: snapshot insert nil")
		}
		v = v.Elem()
		t = v.Type()
	}

	x, err := table.TableOf(t)
	if err != nil {
		return 0, zt, err
	}

	tableName, _, timeName := format.SnapshotName(x.Type.Name(), x.Name)

	i := Insert(tableName)

	st := time.Now()
	i.Columns(timeName)
	values := []interface{}{st}

	for _, c := range x.Columns {
		if c.IsManyRelation() {
			continue
		}
		if v, err := c.GetValue(v); err != nil {
			return 0, zt, err
		} else {
			i.Columns(c.Name)
			values = append(values, v)
		}
	}
	i.Values(values)

	q, a, err := i.Expand(s.xr.Starter())
	if err != nil {
		return 0, zt, err
	}

	r, err := s.xr.Exec(q, a...)
	if err != nil {
		return 0, zt, err
	}

	ai, err := r.LastInsertId()
	if err != nil {
		return 0, zt, err
	}

	return ai, st, nil
}

func (s *Snapshot) Select(id int64, data interface{}, columns ...string) (time.Time, error) {
	var zt time.Time

	v := reflect.ValueOf(data)
	t := v.Type()
	if v.Kind() == reflect.Ptr {
		if v.IsNil() {
			return zt, errors.New("scrud: snapshot select nil")
		}
		v = v.Elem()
		t = v.Type()
	}

	x, err := table.TableOf(t)
	if err != nil {
		return zt, err
	}

	tableName, idName, timeName := format.SnapshotName(x.Type.Name(), x.Name)

	columnMap, exclude, err := tidyColumns("snapshot select", x, columns...)
	if err != nil {
		return zt, err
	}
	count := len(columnMap)

	var st time.Time

	elect := []interface{}{timeName}
	set := make(map[int]*table.Column)
	scans := []interface{}{&st}
	for _, c := range x.Columns {
		if c.IsManyRelation() {
			continue
		}
		if count > 0 {
			if _, ok := columnMap[c.Index]; (ok && exclude) || (!ok && !exclude) {
				continue
			}
		}
		elect = append(elect, c.Name)
		if c.HasEncoding() || c.HasSetter() {
			set[len(scans)] = c
		}
		scans = append(scans, c.Scan(v))
	}

	q, a, err := Select(elect...).From(tableName).Where(Eq(idName, id)).Expand(s.xr.Starter())
	if err != nil {
		return zt, err
	}

	if err := s.xr.QueryRow(q, a...).Scan(scans...); err != nil {
		return zt, err
	}

	for i, c := range set {
		if err := c.SetValue(v, reflect.ValueOf(scans[i]).Elem().Interface()); err != nil {
			return zt, err
		}
	}

	return st, nil
}

func (s *Snapshot) Delete(id int64, data interface{}) error {
	v := reflect.ValueOf(data)
	t := v.Type()
	if v.Kind() == reflect.Ptr {
		if v.IsNil() {
			return errors.New("scrud: snapshot delete nil")
		}
		v = v.Elem()
		t = v.Type()
	}

	x, err := table.TableOf(t)
	if err != nil {
		return err
	}

	tableName, idName, _ := format.SnapshotName(x.Type.Name(), x.Name)

	q, a, err := Delete(tableName).Where(Eq(idName, id)).Limit(1).Expand(s.xr.Starter())
	if err == nil {
		_, err = s.xr.Exec(q, a...)
	}

	return err
}
<commit_msg>Fix snapshot insert only one value<commit_after>package scrud

import (
	"errors"
	"reflect"
	"time"

	"github.com/cxr29/scrud/format"
	"github.com/cxr29/scrud/internal/table"
	. "github.com/cxr29/scrud/query"
)

func (db *DB) Snapshot() *Snapshot {
	return &Snapshot{xr: db}
}

func (tx *Tx) Snapshot() *Snapshot {
	return &Snapshot{xr: tx}
}

// snapshot manager
type Snapshot struct {
	xr faker
}

func (s *Snapshot) Insert(data interface{}) (int64, time.Time, error) {
	var zt time.Time

	v := reflect.ValueOf(data)
	t := v.Type()
	if v.Kind() == reflect.Ptr {
		if v.IsNil() {
			return 0, zt, errors.New("scrud: snapshot insert nil")
		}
		v = v.Elem()
		t = v.Type()
	}

	x, err := table.TableOf(t)
	if err != nil {
		return 0, zt, err
	}

	tableName, _, timeName := format.SnapshotName(x.Type.Name(), x.Name)

	i := Insert(tableName)

	st := time.Now()
	i.Columns(timeName)
	values := []interface{}{st}

	for _, c := range x.Columns {
		if c.IsManyRelation() {
			continue
		}
		if v, err := c.GetValue(v); err != nil {
			return 0, zt, err
		} else {
			i.Columns(c.Name)
			values = append(values, v)
		}
	}
	i.Values(values...)

	q, a, err := i.Expand(s.xr.Starter())
	if err != nil {
		return 0, zt, err
	}

	r, err := s.xr.Exec(q, a...)
	if err != nil {
		return 0, zt, err
	}

	ai, err := r.LastInsertId()
	if err != nil {
		return 0, zt, err
	}

	return ai, st, nil
}

func (s *Snapshot) Select(id int64, data interface{}, columns ...string) (time.Time, error) {
	var zt time.Time

	v := reflect.ValueOf(data)
	t := v.Type()
	if v.Kind() == reflect.Ptr {
		if v.IsNil() {
			return zt, errors.New("scrud: snapshot select nil")
		}
		v = v.Elem()
		t = v.Type()
	}

	x, err := table.TableOf(t)
	if err != nil {
		return zt, err
	}

	tableName, idName, timeName := format.SnapshotName(x.Type.Name(), x.Name)

	columnMap, exclude, err := tidyColumns("snapshot select", x, columns...)
	if err != nil {
		return zt, err
	}
	count := len(columnMap)

	var st time.Time

	elect := []interface{}{timeName}
	set := make(map[int]*table.Column)
	scans := []interface{}{&st}
	for _, c := range x.Columns {
		if c.IsManyRelation() {
			continue
		}
		if count > 0 {
			if _, ok := columnMap[c.Index]; (ok && exclude) || (!ok && !exclude) {
				continue
			}
		}
		elect = append(elect, c.Name)
		if c.HasEncoding() || c.HasSetter() {
			set[len(scans)] = c
		}
		scans = append(scans, c.Scan(v))
	}

	q, a, err := Select(elect...).From(tableName).Where(Eq(idName, id)).Expand(s.xr.Starter())
	if err != nil {
		return zt, err
	}

	if err := s.xr.QueryRow(q, a...).Scan(scans...); err != nil {
		return zt, err
	}

	for i, c := range set {
		if err := c.SetValue(v, reflect.ValueOf(scans[i]).Elem().Interface()); err != nil {
			return zt, err
		}
	}

	return st, nil
}

func (s *Snapshot) Delete(id int64, data interface{}) error {
	v := reflect.ValueOf(data)
	t := v.Type()
	if v.Kind() == reflect.Ptr {
		if v.IsNil() {
			return errors.New("scrud: snapshot delete nil")
		}
		v = v.Elem()
		t = v.Type()
	}

	x, err := table.TableOf(t)
	if err != nil {
		return err
	}

	tableName, idName, _ := format.SnapshotName(x.Type.Name(), x.Name)

	q, a, err := Delete(tableName).Where(Eq(idName, id)).Limit(1).Expand(s.xr.Starter())
	if err == nil {
		_, err = s.xr.Exec(q, a...)
	}

	return err
}
<|endoftext|>


<commit_before>package postgis

import (
	"fmt"
	"github.com/geodan/gost/sensorthings/entities"
	"strconv"

	"encoding/json"
	gostErrors "github.com/geodan/gost/errors"
)

// GetLocation retrieves the location for the given id from the database
func (gdb *GostDatabase) GetLocation(id string) (*entities.Location, error) {
	intID, err := strconv.Atoi(id)
	if err != nil {
		return nil, err
	}

	var sensorID, encodingtype int
	var description, location string
	sql := fmt.Sprintf("select id, description, encodingtype, ST_AsGeoJSON(location) AS location from %s.location where id = $1", gdb.Schema)
	err = gdb.Db.QueryRow(sql, intID).Scan(&sensorID, &description, &encodingtype, &location)

	if err != nil {
		return nil, gostErrors.NewRequestNotFound(fmt.Errorf("Locations(%s) does not exist", id))
	}

	locationMap, err := JSONToMap(location)
	if err != nil {
		return nil, err
	}

	l := entities.Location{
		ID:          strconv.Itoa(sensorID),
		Description: description,
		Location:    locationMap,
	}

	return &l, nil
}

// GetLocations todo
func (gdb *GostDatabase) GetLocations() ([]*entities.Location, error) {
	return nil, nil
}

// PostLocation receives a posted location entity and adds it to the database
// returns the created Location including the generated id
//TODO: ENCODINGTYPE
func (gdb *GostDatabase) PostLocation(location entities.Location) (*entities.Location, error) {
	var locationID int
	locationBytes, _ := json.Marshal(location.Location)
	json := string(locationBytes[:])
	jsonToGeom := fmt.Sprintf("ST_GeomFromGeoJSON('%s')", string(locationBytes[:]))
	sql := fmt.Sprintf("INSERT INTO %s.location (description, encodingtype, location) VALUES ($1, $2, %s) RETURNING id", gdb.Schema, jsonToGeom)
	err := gdb.Db.QueryRow(sql, location.Description, 1, json).Scan(&locationID)
	if err != nil {
		return nil, err
	}

	location.ID = strconv.Itoa(locationID)
	return &location, nil
}

// LocationExists checks if a location is present in the database based on a given id
func (gdb *GostDatabase) LocationExists(locationID int) bool {
	var result bool
	sql := fmt.Sprintf("SELECT exists (SELECT 1 FROM  %s.location WHERE id = $1 LIMIT 1)", gdb.Schema)
	err := gdb.Db.QueryRow(sql, locationID).Scan(&result)
	if err != nil {
		return false
	}

	return result
}

// LinkLocation links a thing with a location
// fails when a thing or location cannot be found for the given id's
func (gdb *GostDatabase) LinkLocation(thingID string, locationID string) error {
	tid, err := strconv.Atoi(thingID)
	if !gdb.ThingExists(tid) || err != nil {
		return fmt.Errorf("Thing(%v) does not exist", thingID)
	}

	lid, err2 := strconv.Atoi(thingID)
	if !gdb.ThingExists(lid) || err2 != nil {
		return fmt.Errorf("Location(%v) does not exist", locationID)
	}

	sql := fmt.Sprintf("INSERT INTO %s.thing_to_location (thing_id, location_id) VALUES ($1, $2)", gdb.Schema)
	_, err3 := gdb.Db.Exec(sql, tid, lid)
	if err3 != nil {
		return err3
	}

	return nil
}
<commit_msg>fixed geojson to postgis geometry on location post<commit_after>package postgis

import (
	"fmt"
	"github.com/geodan/gost/sensorthings/entities"
	"strconv"

	"encoding/json"
	gostErrors "github.com/geodan/gost/errors"
)

// GetLocation retrieves the location for the given id from the database
func (gdb *GostDatabase) GetLocation(id string) (*entities.Location, error) {
	intID, err := strconv.Atoi(id)
	if err != nil {
		return nil, err
	}

	var sensorID, encodingtype int
	var description, location string
	sql := fmt.Sprintf("select id, description, encodingtype, ST_AsGeoJSON(location) AS location from %s.location where id = $1", gdb.Schema)
	err = gdb.Db.QueryRow(sql, intID).Scan(&sensorID, &description, &encodingtype, &location)

	if err != nil {
		return nil, gostErrors.NewRequestNotFound(fmt.Errorf("Locations(%s) does not exist", id))
	}

	locationMap, err := JSONToMap(location)
	if err != nil {
		return nil, err
	}

	l := entities.Location{
		ID:          strconv.Itoa(sensorID),
		Description: description,
		Location:    locationMap,
	}

	return &l, nil
}

// GetLocations todo
func (gdb *GostDatabase) GetLocations() ([]*entities.Location, error) {
	return nil, nil
}

// PostLocation receives a posted location entity and adds it to the database
// returns the created Location including the generated id
//TODO: ENCODINGTYPE
func (gdb *GostDatabase) PostLocation(location entities.Location) (*entities.Location, error) {
	var locationID int
	locationBytes, _ := json.Marshal(location.Location)
	jsonToGeom := fmt.Sprintf("ST_GeomFromGeoJSON('%s')", string(locationBytes[:]))
	sql := fmt.Sprintf("INSERT INTO %s.location (description, encodingtype, location) VALUES ($1, $2, %s) RETURNING id", gdb.Schema, jsonToGeom)
	err := gdb.Db.QueryRow(sql, location.Description, 1).Scan(&locationID)
	if err != nil {
		return nil, err
	}

	location.ID = strconv.Itoa(locationID)
	return &location, nil
}

// LocationExists checks if a location is present in the database based on a given id
func (gdb *GostDatabase) LocationExists(locationID int) bool {
	var result bool
	sql := fmt.Sprintf("SELECT exists (SELECT 1 FROM  %s.location WHERE id = $1 LIMIT 1)", gdb.Schema)
	err := gdb.Db.QueryRow(sql, locationID).Scan(&result)
	if err != nil {
		return false
	}

	return result
}

// LinkLocation links a thing with a location
// fails when a thing or location cannot be found for the given id's
func (gdb *GostDatabase) LinkLocation(thingID string, locationID string) error {
	tid, err := strconv.Atoi(thingID)
	if !gdb.ThingExists(tid) || err != nil {
		return fmt.Errorf("Thing(%v) does not exist", thingID)
	}

	lid, err2 := strconv.Atoi(thingID)
	if !gdb.ThingExists(lid) || err2 != nil {
		return fmt.Errorf("Location(%v) does not exist", locationID)
	}

	sql := fmt.Sprintf("INSERT INTO %s.thing_to_location (thing_id, location_id) VALUES ($1, $2)", gdb.Schema)
	_, err3 := gdb.Db.Exec(sql, tid, lid)
	if err3 != nil {
		return err3
	}

	return nil
}
<|endoftext|>


<commit_before>// The sum of the squares of the first ten natural numbers is,
//
// 1^2 + 2^2 + ... + 102 = 385
//
// The square of the sum of the first ten natural numbers is,
//
// (1 + 2 + ... + 10)^2 = 55^2 = 3025
//
// Hence the difference between the sum of the squares of the first ten natural numbers and the square of the sum is 3025 − 385 = 2640.
//
// Find the difference between the sum of the squares of the first one hundred natural numbers and the square of the sum.
package main

import (
	"fmt"
)

func sumOfTheSquares(n int) int64 {
	var (
		sum int64 = 0
	)

	return sum
}

func squareOfTheSums(n int) int64 {
	var (
		sum int64 = 0
	)

	return sum
}

func main() {
	var (
		diff int64
		// first n natural numbers
		n = 10
	)

	diff = sumOfTheSquares(n) - squareOfTheSums(n)

	fmt.Println(diff)
}
<commit_msg>solved 6<commit_after>// The sum of the squares of the first ten natural numbers is,
//
// 1^2 + 2^2 + ... + 102 = 385
//
// The square of the sum of the first ten natural numbers is,
//
// (1 + 2 + ... + 10)^2 = 55^2 = 3025
//
// Hence the difference between the sum of the squares of the first ten natural numbers and the square of the sum is 3025 − 385 = 2640.
//
// Find the difference between the sum of the squares of the first one hundred natural numbers and the square of the sum.
package main

import (
	"fmt"
)

func sumOfTheSquares(n int) int64 {
	var (
		sum int64 = 0
	)

	for i := 0; i <= n; i++ {
		sum += int64(i * i)
	}

	return sum
}

func squareOfTheSums(n int) int64 {
	var (
		sum int64 = 0
	)

	for i := 0; i <= n; i++ {
		sum += int64(i)

	}

	return sum * sum
}

func main() {
	var (
		diff int64
		// first n natural numbers
		n = 100
	)

	diff = squareOfTheSums(n) - sumOfTheSquares(n)

	fmt.Println(diff)
}
<|endoftext|>


<commit_before>package chatblast

import (
	"log"
	"sync"
)

type Room struct {
	sync.RWMutex
	Name        string           `json:"name,omitempty"`
	Id          string           `json:"id,omitempty"`
	Private     bool             `json:"private,omitempty"`
	Owner       *User            `json:"owner,omitempty"`
	Subscribers map[string]*User `json:"subscribers,omitempty"`
	Channel     chan *Message    `json:"-"`
}

func (r *Room) GetSubscriber(userId string) (*User, bool) {
	r.RLock()
	defer r.RUnlock()
	user, ok := r.Subscribers[userId]
	return user, ok
}

func (r *Room) SetSubscriber(userId string, u *User) {
	r.Lock()
	defer r.Unlock()
	r.Subscribers[userId] = u
}

func (r *Room) RemoveSubscriber(userId string) {
	r.Lock()
	defer r.Unlock()
	delete(r.Subscribers, userId)
}

func NewRoom(name string, owner *User) *Room {
	newRoom := &Room{
		Name:        name,
		Id:          GUID(),
		Owner:       owner,
		Subscribers: map[string]*User{},
		Channel:     make(chan *Message),
	}
	go newRoom.Listen()
	log.Println("New room", name, "created")
	return newRoom
}

func (r *Room) Listen() {
	for incoming := range r.Channel {
		log.Println("Message sent in room", r.Name)
		for _, user := range r.Subscribers {
			user.Tell(incoming)
		}
	}
}

func (r *Room) Broadcast(msg *Message) {
	r.Channel <- msg
}

func (r *Room) Whisper(msg *Message) {
	if whisperee, ok := r.Subscribers[msg.UserId]; ok {
		whisperee.Tell(msg)
	}
	if self, ok := r.Subscribers[msg.User.Id]; ok {
		self.Tell(msg)
	}
}

func (r *Room) RequestInvite(u *User) {
	msg := &Message{
		Cmd:    "inviterequest",
		RoomId: r.Id,
		UserId: u.Id,
	}
	r.Owner.Channel <- msg
}

func (r *Room) Join(u *User) {
	// TODO Add re-join check
	r.SetSubscriber(u.Id, u)
	msg := &Message{
		Cmd:    "sub",
		User:   u,
		RoomId: r.Id,
		UserId: u.Id,
	}
	r.Broadcast(msg)
}

func (r *Room) Leave(u *User) {
	// TODO Add bad leave request check (i.e. user tries to leave room they
	// aren't subscribed to)
	r.RemoveSubscriber(u.Id)
	msg := &Message{
		Cmd:    "unsub",
		RoomId: r.Id,
		UserId: u.Id,
	}
	r.Broadcast(msg)
	u.Tell(msg)
}

func (r *Room) CloseRoom(u *User) (closed bool, err error) {
	closed = false
	err = nil
	if r.Owner != nil && u.Id == r.Owner.Id {
		closed = true
	}
	return closed, err

}
<commit_msg>Prevent whisper sending if whisperee has logged off<commit_after>package chatblast

import (
	"log"
	"sync"
)

type Room struct {
	sync.RWMutex
	Name        string           `json:"name,omitempty"`
	Id          string           `json:"id,omitempty"`
	Private     bool             `json:"private,omitempty"`
	Owner       *User            `json:"owner,omitempty"`
	Subscribers map[string]*User `json:"subscribers,omitempty"`
	Channel     chan *Message    `json:"-"`
}

func (r *Room) GetSubscriber(userId string) (*User, bool) {
	r.RLock()
	defer r.RUnlock()
	user, ok := r.Subscribers[userId]
	return user, ok
}

func (r *Room) SetSubscriber(userId string, u *User) {
	r.Lock()
	defer r.Unlock()
	r.Subscribers[userId] = u
}

func (r *Room) RemoveSubscriber(userId string) {
	r.Lock()
	defer r.Unlock()
	delete(r.Subscribers, userId)
}

func NewRoom(name string, owner *User) *Room {
	newRoom := &Room{
		Name:        name,
		Id:          GUID(),
		Owner:       owner,
		Subscribers: map[string]*User{},
		Channel:     make(chan *Message),
	}
	go newRoom.Listen()
	log.Println("New room", name, "created")
	return newRoom
}

func (r *Room) Listen() {
	for incoming := range r.Channel {
		log.Println("Message sent in room", r.Name)
		for _, user := range r.Subscribers {
			user.Tell(incoming)
		}
	}
}

func (r *Room) Broadcast(msg *Message) {
	r.Channel <- msg
}

func (r *Room) Whisper(msg *Message) {
	if whisperee, ok := r.Subscribers[msg.UserId]; ok {
		msg.Whisperee = whisperee
		whisperee.Tell(msg)
		if self, ok := r.Subscribers[msg.User.Id]; ok {
			self.Tell(msg)
		}
	}
}

func (r *Room) RequestInvite(u *User) {
	msg := &Message{
		Cmd:    "inviterequest",
		RoomId: r.Id,
		UserId: u.Id,
	}
	r.Owner.Channel <- msg
}

func (r *Room) Join(u *User) {
	// TODO Add re-join check
	r.SetSubscriber(u.Id, u)
	msg := &Message{
		Cmd:    "sub",
		User:   u,
		RoomId: r.Id,
		UserId: u.Id,
	}
	r.Broadcast(msg)
}

func (r *Room) Leave(u *User) {
	// TODO Add bad leave request check (i.e. user tries to leave room they
	// aren't subscribed to)
	r.RemoveSubscriber(u.Id)
	msg := &Message{
		Cmd:    "unsub",
		RoomId: r.Id,
		UserId: u.Id,
	}
	r.Broadcast(msg)
	u.Tell(msg)
}

func (r *Room) CloseRoom(u *User) (closed bool, err error) {
	closed = false
	err = nil
	if r.Owner != nil && u.Id == r.Owner.Id {
		closed = true
	}
	return closed, err

}
<|endoftext|>


<commit_before>package mp

import (
	"bufio"
	"bytes"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base64"
	"encoding/hex"
	"encoding/xml"
	"github.com/ridewindx/mel"
	"github.com/ridewindx/mel/binding"
	"sort"
	"sync"
	"sync/atomic"
	"unsafe"
	"go.uber.org/zap"
	"github.com/jiudaoyun/wechat"
	"net/http"
	"fmt"
	"encoding/json"
	"strings"
	"github.com/ridewindx/melware"
)

type Server struct {
	*mel.Mel
	urlPrefix string

	appID string // App ID
	ID    string // Wechat ID

	tokenMutex sync.Mutex
	token      unsafe.Pointer

	aesKeyMutex sync.Mutex
	aesKey      unsafe.Pointer

	client            *Client
	middlewares       []Handler
	messageHandlerMap map[string]Handler
	eventHandlerMap   map[string]Handler

	logger *zap.SugaredLogger
}

func (srv *Server) setURLPrefix(urlPrefix string) {
	if !strings.HasPrefix(urlPrefix, "/") {
		urlPrefix = "/" + urlPrefix
	}
	urlPrefix = strings.TrimRight(urlPrefix, "/")
	srv.urlPrefix = urlPrefix
}

func (srv *Server) SetID(id string) {
	srv.ID = id
}

func (srv *Server) SetAppID(appID string) {
	srv.appID = appID
}

func (srv *Server) SetClient(client *Client) {
	srv.client = client
}

type Token struct {
	current string
	last    string
}

type AESKey struct {
	current string
	last    string
}

func (srv *Server) GetToken() (string, string) {
	p := (*Token)(atomic.LoadPointer(&srv.token))
	if p != nil {
		return p.current, p.last
	}
	return "", ""
}

func (srv *Server) SetToken(token string) {
	if token == "" {
		return
	}

	srv.tokenMutex.Lock()
	defer srv.tokenMutex.Unlock()

	current, _ := srv.GetToken()
	if token == current {
		return
	}

	t := Token{
		current: token,
		last:    current,
	}
	atomic.StorePointer(&srv.token, unsafe.Pointer(&t))
}

func (srv *Server) deleteLastToken() {
	srv.tokenMutex.Lock()
	defer srv.tokenMutex.Unlock()

	current, last := srv.GetToken()
	if last == "" {
		return
	}

	t := Token{
		current: current,
	}
	atomic.StorePointer(&srv.token, unsafe.Pointer(&t))
}

func (srv *Server) GetAESKey() (string, string) {
	p := (*AESKey)(atomic.LoadPointer(&srv.aesKey))
	if p != nil {
		return p.current, p.last
	}
	return "", ""
}

func (srv *Server) SetAESKey(base64AESKey string) {
	if len(base64AESKey) != 43 {
		return
	}
	aesKey, err := base64.StdEncoding.DecodeString(base64AESKey + "=")
	if err != nil {
		return
	}

	srv.aesKeyMutex.Lock()
	defer srv.aesKeyMutex.Unlock()

	current, _ := srv.GetAESKey()
	if bytes.Equal(aesKey, []byte(current)) {
		return
	}

	k := AESKey{
		current: string(aesKey),
		last:    current,
	}
	atomic.StorePointer(&srv.aesKey, unsafe.Pointer(&k))
}

func (srv *Server) deleteLastAESKey() {
	srv.aesKeyMutex.Lock()
	defer srv.aesKeyMutex.Unlock()

	current, last := srv.GetAESKey()
	if last == "" {
		return
	}

	k := AESKey{
		current: current,
	}
	atomic.StorePointer(&srv.aesKey, unsafe.Pointer(&k))
}

func (srv *Server) Use(middlewares ...Handler) {
	srv.middlewares = append(srv.middlewares, middlewares...)
	if len(srv.middlewares)+1 > int(abortIndex) {
		panic("too many middlewares")
	}
}

func (srv *Server) HandleMessage(msgType string, handler Handler) {
	srv.messageHandlerMap[msgType] = handler
}

func (srv *Server) HandleEvent(eventType string, handler Handler) {
	srv.eventHandlerMap[eventType] = handler
}

func NewServer(token, aesKey string, urlPrefix ...string) *Server {
	srv := &Server{
		Mel:               mel.New(),
		messageHandlerMap: make(map[string]Handler),
		eventHandlerMap:   make(map[string]Handler),
		logger:            wechat.Sugar,
	}

	srv.SetToken(token)
	srv.SetAESKey(aesKey)

	if len(urlPrefix) > 0 {
		srv.setURLPrefix(urlPrefix[0])
	}

	equal := func(a, b string) bool {
		return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
	}

	verifySignReturnToken := func(signature, timestamp, nonce string) string {
		currentToken, lastToken := srv.GetToken()
		token := currentToken

		isValid := func() bool {
			computedSignature := computeSign(token, timestamp, nonce)
			return equal(signature, computedSignature)
		}

		if isValid() {
			srv.deleteLastToken()
			return token
		}

		if lastToken != "" {
			token = lastToken
			if isValid() {
				return token
			}
		}

		return ""
	}

	verifySign := func(c *mel.Context) bool {
		signature := c.Query("signature")
		timestamp := c.Query("timestamp")
		nonce := c.Query("nonce")

		return verifySignReturnToken(signature, timestamp, nonce) != ""
	}

	type EncryptMsg struct {
		ToUserName string `xml:"ToUserName"`
		Encrypt    string `xml:"Encrypt"`
	}

	cors := melware.CorsAllowAll()
	cors.AllowCredentials = false
	srv.Mel.Use(cors.Middleware())

	srv.Head("/", func(c *mel.Context) { // health check
		c.Status(200)
	})

	srv.Get(srv.urlPrefix+"/", func(c *mel.Context) {
		if verifySign(c) {
			echostr := c.Query("echostr")
			c.Text(200, echostr)
		}
	})

	handleMessage := func(event *Event) interface{} {
		var handler Handler
		var ok bool
		if event.Type == MessageEvent {
			handler, ok = srv.eventHandlerMap[event.Event]
		} else {
			handler, ok = srv.messageHandlerMap[event.Type]
		}
		if !ok {
			return nil // no registered handler, just respond with empty string
		}

		ctx := &Context{
			Client:   srv.client,
			index:    preStartIndex,
			handlers: append(srv.middlewares, handler),
			Event:    event,
		}

		ctx.Next()

		return ctx.response
	}

	srv.Post(srv.urlPrefix+"/", func(c *mel.Context) {
		encryptType := c.Query("encrypt_type")
		signature := c.Query("signature")
		timestamp := c.Query("timestamp")
		nonce := c.Query("nonce")

		switch encryptType {
		case "aes":
			token := verifySignReturnToken(signature, timestamp, nonce)
			if token == "" {
				srv.logger.Error("Verify sign empty token")
				return
			}

			msgSign := c.Query("msg_signature")

			var obj EncryptMsg
			err := c.BindWith(&obj, binding.XML)
			if err != nil {
				srv.logger.Errorw("Bind with XML failed", "error", err)
				return
			}

			if srv.ID != "" && !equal(obj.ToUserName, srv.ID) {
				srv.logger.Errorw("Wechat ID inconsistent", "id", srv.ID, "ToUserName", obj.ToUserName)
				return
			}

			computedSign := computeSign(token, timestamp, nonce, obj.Encrypt)
			if !equal(computedSign, msgSign) {
				srv.logger.Errorw("Signature inconsistent")
				return
			}

			encryptedMsg, err := base64.StdEncoding.DecodeString(obj.Encrypt)
			if err != nil {
				srv.logger.Errorw("Decode base64 string failed", "error", err)
				return
			}

			current, last := srv.GetAESKey()
			aesKey := current
			random, msg, appId, err := decryptMsg(encryptedMsg, []byte(aesKey))
			if err != nil {
				if last == "" {
					srv.logger.Errorw("Decrypt AES msg failed", "error", err)
					return
				}
				aesKey = last
				random, msg, appId, err = decryptMsg(encryptedMsg, []byte(aesKey))
				if err != nil {
					srv.logger.Errorw("Decrypt AES msg failed", "error", err)
					return
				}
			} else {
				srv.deleteLastAESKey()
			}
			if srv.appID != "" && string(appId) != srv.appID {
				srv.logger.Errorw("AppID inconsistent", "AppID", appId)
				return
			}

			var event Event
			if err = xml.Unmarshal(msg, &event); err != nil {
				srv.logger.Errorw("Unmarshal msg failed", "error", err)
				return
			}

			repBytes, err := xml.Marshal(handleMessage(&event))
			if err != nil {
				srv.logger.Errorw("Marshal msg failed", "error", err)
				return
			}

			encryptedRepBytes := encryptMsg(random, repBytes, appId, []byte(aesKey))
			encryptedRepStr := base64.StdEncoding.EncodeToString(encryptedRepBytes)
			repSignature := computeSign(token, timestamp, nonce, encryptedRepStr)

			type EncryptRepMsg struct {
				Encrypt      string
				MsgSignature string
				TimeStamp    string
				Nonce        string
			}

			err = c.XML(200, &EncryptRepMsg{encryptedRepStr, repSignature, timestamp, nonce})
			if err != nil {
				srv.logger.Errorw("Reply msg failed", "error", err)
			}

		case "", "raw":
			if !verifySign(c) {
				return
			}

			var event Event
			err := c.BindWith(&event, binding.XML)
			if err != nil {
				return
			}

			c.XML(200, handleMessage(&event))

		default:
			return
		}
	})

	handleAuthorize := func(c *mel.Context, url string, state string) {
		rep, err := srv.client.Client.Get(url)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
		defer rep.Body.Close()

		if rep.StatusCode != http.StatusOK {
			c.AbortWithError(http.StatusUnauthorized, fmt.Errorf("http.Status: %s", rep.Status))
			return
		}

		type Result struct {
			AccessToken  string `json:"access_token"`
			ExpiresIn    string `json:"expires_in"`
			RefreshToken string `json:"refresh_token"`
			OpenID       string `json:"openid"`
			Scope        string `json:"scope"`
			State        string `json:"state,omitempty"`
		}

		type ResultWithErr struct {
			Result
			Err
		}

		var result ResultWithErr
		err = json.NewDecoder(rep.Body).Decode(&result)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}

		if result.Code() != OK {
			c.AbortWithError(http.StatusUnauthorized, &result)
			return
		}

		result.State = state
		c.JSON(http.StatusOK, &result.Result)
	}

	srv.Get(srv.urlPrefix+"/token", func(c *mel.Context) {
		code := c.Query("code")
		state := c.Query("state")

		url := fmt.Sprintf("https://api.weixin.qq.com/sns/oauth2/access_token?appid=%s&secret=%s&code=%s&grant_type=authorization_code", srv.client.appId, srv.client.appSecret, code)

		handleAuthorize(c, url, state)
	})

	srv.Get(srv.urlPrefix+"/refresh-token", func(c *mel.Context) {
		refreshToken := c.Query("refresh_token")

		url := fmt.Sprintf("https://api.weixin.qq.com/sns/oauth2/refresh_token?appid=%s&grant_type=refresh_token&refresh_token=%s", srv.client.appId, refreshToken)

		handleAuthorize(c, url, "")
	})

	srv.Get(srv.urlPrefix+"/signature", func(c *mel.Context) {
		timestamp := c.Query("timestamp")
		nonceStr := c.Query("nonceStr")
		url := c.Query("url")
		refresh := c.Query("refresh")

		var ticket string
		var err error
		if refresh != "" && (refresh == "true" || refresh == "True" || refresh == "1") {
			ticket, err = srv.client.RefreshTicket("")
		} else {
			ticket, err = srv.client.Ticket()
		}
		if err != nil {
			c.AbortWithError(http.StatusInternalServerError, err)
			return
		}

		strs := sort.StringSlice{
			"timestamp=" + timestamp,
			"nonceStr=" + nonceStr,
			"url=" + url,
			"jsapi_ticket=" + ticket,
		}
		strs.Sort()
		h := sha1.New()
		buf := bufio.NewWriterSize(h, 1024)
		for i, s := range strs {
			buf.WriteString(s)
			if i < len(strs)-1 {
				buf.WriteByte('&')
			}
		}
		buf.Flush()
		sign := hex.EncodeToString(h.Sum(nil))
		c.JSON(http.StatusOK, map[string]string{
			"signature": sign,
		})
		srv.logger.Infof("signature", "strs", strs, "sign", sign)
	})

	return srv
}

func computeSign(elements ...string) string {
	strs := sort.StringSlice(elements)
	strs.Sort()

	h := sha1.New()

	buf := bufio.NewWriterSize(h, 1024)
	for _, s := range strs {
		buf.WriteString(s)
	}
	buf.Flush()

	return hex.EncodeToString(h.Sum(nil))
}
<commit_msg>Update files<commit_after>package mp

import (
	"bufio"
	"bytes"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base64"
	"encoding/hex"
	"encoding/xml"
	"github.com/ridewindx/mel"
	"github.com/ridewindx/mel/binding"
	"sort"
	"sync"
	"sync/atomic"
	"unsafe"
	"go.uber.org/zap"
	"github.com/jiudaoyun/wechat"
	"net/http"
	"fmt"
	"encoding/json"
	"strings"
	"github.com/ridewindx/melware"
)

type Server struct {
	*mel.Mel
	urlPrefix string

	appID string // App ID
	ID    string // Wechat ID

	tokenMutex sync.Mutex
	token      unsafe.Pointer

	aesKeyMutex sync.Mutex
	aesKey      unsafe.Pointer

	client            *Client
	middlewares       []Handler
	messageHandlerMap map[string]Handler
	eventHandlerMap   map[string]Handler

	logger *zap.SugaredLogger
}

func (srv *Server) setURLPrefix(urlPrefix string) {
	if !strings.HasPrefix(urlPrefix, "/") {
		urlPrefix = "/" + urlPrefix
	}
	urlPrefix = strings.TrimRight(urlPrefix, "/")
	srv.urlPrefix = urlPrefix
}

func (srv *Server) SetID(id string) {
	srv.ID = id
}

func (srv *Server) SetAppID(appID string) {
	srv.appID = appID
}

func (srv *Server) SetClient(client *Client) {
	srv.client = client
}

type Token struct {
	current string
	last    string
}

type AESKey struct {
	current string
	last    string
}

func (srv *Server) GetToken() (string, string) {
	p := (*Token)(atomic.LoadPointer(&srv.token))
	if p != nil {
		return p.current, p.last
	}
	return "", ""
}

func (srv *Server) SetToken(token string) {
	if token == "" {
		return
	}

	srv.tokenMutex.Lock()
	defer srv.tokenMutex.Unlock()

	current, _ := srv.GetToken()
	if token == current {
		return
	}

	t := Token{
		current: token,
		last:    current,
	}
	atomic.StorePointer(&srv.token, unsafe.Pointer(&t))
}

func (srv *Server) deleteLastToken() {
	srv.tokenMutex.Lock()
	defer srv.tokenMutex.Unlock()

	current, last := srv.GetToken()
	if last == "" {
		return
	}

	t := Token{
		current: current,
	}
	atomic.StorePointer(&srv.token, unsafe.Pointer(&t))
}

func (srv *Server) GetAESKey() (string, string) {
	p := (*AESKey)(atomic.LoadPointer(&srv.aesKey))
	if p != nil {
		return p.current, p.last
	}
	return "", ""
}

func (srv *Server) SetAESKey(base64AESKey string) {
	if len(base64AESKey) != 43 {
		return
	}
	aesKey, err := base64.StdEncoding.DecodeString(base64AESKey + "=")
	if err != nil {
		return
	}

	srv.aesKeyMutex.Lock()
	defer srv.aesKeyMutex.Unlock()

	current, _ := srv.GetAESKey()
	if bytes.Equal(aesKey, []byte(current)) {
		return
	}

	k := AESKey{
		current: string(aesKey),
		last:    current,
	}
	atomic.StorePointer(&srv.aesKey, unsafe.Pointer(&k))
}

func (srv *Server) deleteLastAESKey() {
	srv.aesKeyMutex.Lock()
	defer srv.aesKeyMutex.Unlock()

	current, last := srv.GetAESKey()
	if last == "" {
		return
	}

	k := AESKey{
		current: current,
	}
	atomic.StorePointer(&srv.aesKey, unsafe.Pointer(&k))
}

func (srv *Server) Use(middlewares ...Handler) {
	srv.middlewares = append(srv.middlewares, middlewares...)
	if len(srv.middlewares)+1 > int(abortIndex) {
		panic("too many middlewares")
	}
}

func (srv *Server) HandleMessage(msgType string, handler Handler) {
	srv.messageHandlerMap[msgType] = handler
}

func (srv *Server) HandleEvent(eventType string, handler Handler) {
	srv.eventHandlerMap[eventType] = handler
}

func NewServer(token, aesKey string, urlPrefix ...string) *Server {
	srv := &Server{
		Mel:               mel.New(),
		messageHandlerMap: make(map[string]Handler),
		eventHandlerMap:   make(map[string]Handler),
		logger:            wechat.Sugar,
	}

	srv.SetToken(token)
	srv.SetAESKey(aesKey)

	if len(urlPrefix) > 0 {
		srv.setURLPrefix(urlPrefix[0])
	}

	equal := func(a, b string) bool {
		return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
	}

	verifySignReturnToken := func(signature, timestamp, nonce string) string {
		currentToken, lastToken := srv.GetToken()
		token := currentToken

		isValid := func() bool {
			computedSignature := computeSign(token, timestamp, nonce)
			return equal(signature, computedSignature)
		}

		if isValid() {
			srv.deleteLastToken()
			return token
		}

		if lastToken != "" {
			token = lastToken
			if isValid() {
				return token
			}
		}

		return ""
	}

	verifySign := func(c *mel.Context) bool {
		signature := c.Query("signature")
		timestamp := c.Query("timestamp")
		nonce := c.Query("nonce")

		return verifySignReturnToken(signature, timestamp, nonce) != ""
	}

	type EncryptMsg struct {
		ToUserName string `xml:"ToUserName"`
		Encrypt    string `xml:"Encrypt"`
	}

	cors := melware.CorsAllowAll()
	cors.AllowCredentials = false
	srv.Mel.Use(cors.Middleware())

	srv.Head("/", func(c *mel.Context) { // health check
		c.Status(200)
	})

	srv.Get(srv.urlPrefix+"/", func(c *mel.Context) {
		if verifySign(c) {
			echostr := c.Query("echostr")
			c.Text(200, echostr)
		}
	})

	handleMessage := func(event *Event) interface{} {
		var handler Handler
		var ok bool
		if event.Type == MessageEvent {
			handler, ok = srv.eventHandlerMap[event.Event]
		} else {
			handler, ok = srv.messageHandlerMap[event.Type]
		}
		if !ok {
			return nil // no registered handler, just respond with empty string
		}

		ctx := &Context{
			Client:   srv.client,
			index:    preStartIndex,
			handlers: append(srv.middlewares, handler),
			Event:    event,
		}

		ctx.Next()

		return ctx.response
	}

	srv.Post(srv.urlPrefix+"/", func(c *mel.Context) {
		encryptType := c.Query("encrypt_type")
		signature := c.Query("signature")
		timestamp := c.Query("timestamp")
		nonce := c.Query("nonce")

		switch encryptType {
		case "aes":
			token := verifySignReturnToken(signature, timestamp, nonce)
			if token == "" {
				srv.logger.Error("Verify sign empty token")
				return
			}

			msgSign := c.Query("msg_signature")

			var obj EncryptMsg
			err := c.BindWith(&obj, binding.XML)
			if err != nil {
				srv.logger.Errorw("Bind with XML failed", "error", err)
				return
			}

			if srv.ID != "" && !equal(obj.ToUserName, srv.ID) {
				srv.logger.Errorw("Wechat ID inconsistent", "id", srv.ID, "ToUserName", obj.ToUserName)
				return
			}

			computedSign := computeSign(token, timestamp, nonce, obj.Encrypt)
			if !equal(computedSign, msgSign) {
				srv.logger.Errorw("Signature inconsistent")
				return
			}

			encryptedMsg, err := base64.StdEncoding.DecodeString(obj.Encrypt)
			if err != nil {
				srv.logger.Errorw("Decode base64 string failed", "error", err)
				return
			}

			current, last := srv.GetAESKey()
			aesKey := current
			random, msg, appId, err := decryptMsg(encryptedMsg, []byte(aesKey))
			if err != nil {
				if last == "" {
					srv.logger.Errorw("Decrypt AES msg failed", "error", err)
					return
				}
				aesKey = last
				random, msg, appId, err = decryptMsg(encryptedMsg, []byte(aesKey))
				if err != nil {
					srv.logger.Errorw("Decrypt AES msg failed", "error", err)
					return
				}
			} else {
				srv.deleteLastAESKey()
			}
			if srv.appID != "" && string(appId) != srv.appID {
				srv.logger.Errorw("AppID inconsistent", "AppID", appId)
				return
			}

			var event Event
			if err = xml.Unmarshal(msg, &event); err != nil {
				srv.logger.Errorw("Unmarshal msg failed", "error", err)
				return
			}

			repBytes, err := xml.Marshal(handleMessage(&event))
			if err != nil {
				srv.logger.Errorw("Marshal msg failed", "error", err)
				return
			}

			encryptedRepBytes := encryptMsg(random, repBytes, appId, []byte(aesKey))
			encryptedRepStr := base64.StdEncoding.EncodeToString(encryptedRepBytes)
			repSignature := computeSign(token, timestamp, nonce, encryptedRepStr)

			type EncryptRepMsg struct {
				Encrypt      string
				MsgSignature string
				TimeStamp    string
				Nonce        string
			}

			err = c.XML(200, &EncryptRepMsg{encryptedRepStr, repSignature, timestamp, nonce})
			if err != nil {
				srv.logger.Errorw("Reply msg failed", "error", err)
			}

		case "", "raw":
			if !verifySign(c) {
				return
			}

			var event Event
			err := c.BindWith(&event, binding.XML)
			if err != nil {
				return
			}

			c.XML(200, handleMessage(&event))

		default:
			return
		}
	})

	handleAuthorize := func(c *mel.Context, url string, state string) {
		rep, err := srv.client.Client.Get(url)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
		defer rep.Body.Close()

		if rep.StatusCode != http.StatusOK {
			c.AbortWithError(http.StatusUnauthorized, fmt.Errorf("http.Status: %s", rep.Status))
			return
		}

		type Result struct {
			AccessToken  string `json:"access_token"`
			ExpiresIn    string `json:"expires_in"`
			RefreshToken string `json:"refresh_token"`
			OpenID       string `json:"openid"`
			Scope        string `json:"scope"`
			State        string `json:"state,omitempty"`
		}

		type ResultWithErr struct {
			Result
			Err
		}

		var result ResultWithErr
		err = json.NewDecoder(rep.Body).Decode(&result)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}

		if result.Code() != OK {
			c.AbortWithError(http.StatusUnauthorized, &result)
			return
		}

		result.State = state
		c.JSON(http.StatusOK, &result.Result)
	}

	srv.Get(srv.urlPrefix+"/token", func(c *mel.Context) {
		code := c.Query("code")
		state := c.Query("state")

		url := fmt.Sprintf("https://api.weixin.qq.com/sns/oauth2/access_token?appid=%s&secret=%s&code=%s&grant_type=authorization_code", srv.client.appId, srv.client.appSecret, code)

		handleAuthorize(c, url, state)
	})

	srv.Get(srv.urlPrefix+"/refresh-token", func(c *mel.Context) {
		refreshToken := c.Query("refresh_token")

		url := fmt.Sprintf("https://api.weixin.qq.com/sns/oauth2/refresh_token?appid=%s&grant_type=refresh_token&refresh_token=%s", srv.client.appId, refreshToken)

		handleAuthorize(c, url, "")
	})

	srv.Get(srv.urlPrefix+"/signature", func(c *mel.Context) {
		timestamp := c.Query("timestamp")
		nonceStr := c.Query("nonceStr")
		url := c.Query("url")
		refresh := c.Query("refresh")

		var ticket string
		var err error
		if refresh != "" && (refresh == "true" || refresh == "True" || refresh == "1") {
			ticket, err = srv.client.RefreshTicket("")
		} else {
			ticket, err = srv.client.Ticket()
		}
		if err != nil {
			c.AbortWithError(http.StatusInternalServerError, err)
			return
		}

		strs := sort.StringSlice{
			"timestamp=" + timestamp,
			"nonceStr=" + nonceStr,
			"url=" + url,
			"jsapi_ticket=" + ticket,
		}
		strs.Sort()
		h := sha1.New()
		b := bytes.Buffer{}
		for i, s := range strs {
			b.WriteString(s)
			if i < len(strs)-1 {
				b.WriteByte('&')
			}
		}
		srv.logger.Infow("buf", "b", b.String())

		buf := bufio.NewWriterSize(h, 1024)
		for i, s := range strs {
			buf.WriteString(s)
			if i < len(strs)-1 {
				buf.WriteByte('&')
			}
		}
		buf.Flush()
		sign := hex.EncodeToString(h.Sum(nil))
		c.JSON(http.StatusOK, map[string]string{
			"signature": sign,
		})
		srv.logger.Infow("signature", "strs", strs, "sign", sign)
	})

	return srv
}

func computeSign(elements ...string) string {
	strs := sort.StringSlice(elements)
	strs.Sort()

	h := sha1.New()

	buf := bufio.NewWriterSize(h, 1024)
	for _, s := range strs {
		buf.WriteString(s)
	}
	buf.Flush()

	return hex.EncodeToString(h.Sum(nil))
}
<|endoftext|>


<commit_before>package mp

import (
	"bufio"
	"bytes"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base64"
	"encoding/hex"
	"encoding/xml"
	"github.com/ridewindx/mel"
	"github.com/ridewindx/mel/binding"
	"sort"
	"sync"
	"sync/atomic"
	"unsafe"
	"go.uber.org/zap"
	"github.com/jiudaoyun/wechat"
	"net/http"
	"fmt"
	"encoding/json"
	"strings"
	"github.com/ridewindx/melware"
	"path"
)

type Server struct {
	*mel.Mel
	urlPrefix string

	appID string // App ID
	ID    string // Wechat ID

	tokenMutex sync.Mutex
	token      unsafe.Pointer

	aesKeyMutex sync.Mutex
	aesKey      unsafe.Pointer

	client            *Client
	middlewares       []Handler
	messageHandlerMap map[string]Handler
	eventHandlerMap   map[string]Handler

	logger *zap.SugaredLogger
}

func (srv *Server) setURLPrefix(urlPrefix string) {
	if !strings.HasPrefix(urlPrefix, "/") {
		urlPrefix = "/" + urlPrefix
	}
	urlPrefix = strings.TrimRight(urlPrefix, "/")
	srv.urlPrefix = urlPrefix
}

func (srv *Server) SetID(id string) {
	srv.ID = id
}

func (srv *Server) SetAppID(appID string) {
	srv.appID = appID
}

func (srv *Server) SetClient(client *Client) {
	srv.client = client
}

type Token struct {
	current string
	last    string
}

type AESKey struct {
	current string
	last    string
}

func (srv *Server) GetToken() (string, string) {
	p := (*Token)(atomic.LoadPointer(&srv.token))
	if p != nil {
		return p.current, p.last
	}
	return "", ""
}

func (srv *Server) SetToken(token string) {
	if token == "" {
		return
	}

	srv.tokenMutex.Lock()
	defer srv.tokenMutex.Unlock()

	current, _ := srv.GetToken()
	if token == current {
		return
	}

	t := Token{
		current: token,
		last:    current,
	}
	atomic.StorePointer(&srv.token, unsafe.Pointer(&t))
}

func (srv *Server) deleteLastToken() {
	srv.tokenMutex.Lock()
	defer srv.tokenMutex.Unlock()

	current, last := srv.GetToken()
	if last == "" {
		return
	}

	t := Token{
		current: current,
	}
	atomic.StorePointer(&srv.token, unsafe.Pointer(&t))
}

func (srv *Server) GetAESKey() (string, string) {
	p := (*AESKey)(atomic.LoadPointer(&srv.aesKey))
	if p != nil {
		return p.current, p.last
	}
	return "", ""
}

func (srv *Server) SetAESKey(base64AESKey string) {
	if len(base64AESKey) != 43 {
		return
	}
	aesKey, err := base64.StdEncoding.DecodeString(base64AESKey + "=")
	if err != nil {
		return
	}

	srv.aesKeyMutex.Lock()
	defer srv.aesKeyMutex.Unlock()

	current, _ := srv.GetAESKey()
	if bytes.Equal(aesKey, []byte(current)) {
		return
	}

	k := AESKey{
		current: string(aesKey),
		last:    current,
	}
	atomic.StorePointer(&srv.aesKey, unsafe.Pointer(&k))
}

func (srv *Server) deleteLastAESKey() {
	srv.aesKeyMutex.Lock()
	defer srv.aesKeyMutex.Unlock()

	current, last := srv.GetAESKey()
	if last == "" {
		return
	}

	k := AESKey{
		current: current,
	}
	atomic.StorePointer(&srv.aesKey, unsafe.Pointer(&k))
}

func (srv *Server) Use(middlewares ...Handler) {
	srv.middlewares = append(srv.middlewares, middlewares...)
	if len(srv.middlewares)+1 > int(abortIndex) {
		panic("too many middlewares")
	}
}

func (srv *Server) HandleMessage(msgType string, handler Handler) {
	srv.messageHandlerMap[msgType] = handler
}

func (srv *Server) HandleEvent(eventType string, handler Handler) {
	srv.eventHandlerMap[eventType] = handler
}

func (srv *Server) GetVerifyFile(filename string, content []byte) {
	srv.Get(srv.urlPrefix+filename, func(c *mel.Context) {
		c.Data(200, "text/plain", content)
	})
}

func NewServer(token, aesKey string, urlPrefix ...string) *Server {
	srv := &Server{
		Mel:               mel.New(),
		messageHandlerMap: make(map[string]Handler),
		eventHandlerMap:   make(map[string]Handler),
		logger:            wechat.Sugar,
	}

	srv.SetToken(token)
	srv.SetAESKey(aesKey)

	if len(urlPrefix) > 0 {
		srv.setURLPrefix(urlPrefix[0])
	}

	equal := func(a, b string) bool {
		return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
	}

	verifySignReturnToken := func(signature, timestamp, nonce string) string {
		currentToken, lastToken := srv.GetToken()
		token := currentToken

		isValid := func() bool {
			computedSignature := computeSign(token, timestamp, nonce)
			return equal(signature, computedSignature)
		}

		if isValid() {
			srv.deleteLastToken()
			return token
		}

		if lastToken != "" {
			token = lastToken
			if isValid() {
				return token
			}
		}

		return ""
	}

	verifySign := func(c *mel.Context) bool {
		signature := c.Query("signature")
		timestamp := c.Query("timestamp")
		nonce := c.Query("nonce")

		return verifySignReturnToken(signature, timestamp, nonce) != ""
	}

	type EncryptMsg struct {
		ToUserName string `xml:"ToUserName"`
		Encrypt    string `xml:"Encrypt"`
	}

	cors := melware.CorsAllowAll()
	cors.AllowCredentials = false
	srv.Mel.Use(cors.Middleware())

	srv.Head("/", func(c *mel.Context) { // health check
		c.Status(200)
	})

	srv.Get(srv.urlPrefix+"/", func(c *mel.Context) {
		if verifySign(c) {
			echostr := c.Query("echostr")
			c.Text(200, echostr)
		}
	})

	handleMessage := func(event *Event) interface{} {
		var handler Handler
		var ok bool
		if event.Type == MessageEvent {
			handler, ok = srv.eventHandlerMap[event.Event]
		} else {
			handler, ok = srv.messageHandlerMap[event.Type]
		}
		if !ok {
			return nil // no registered handler, just respond with empty string
		}

		ctx := &Context{
			Client:   srv.client,
			index:    preStartIndex,
			handlers: append(srv.middlewares, handler),
			Event:    event,
		}

		ctx.Next()

		return ctx.response
	}

	srv.Post(srv.urlPrefix+"/", func(c *mel.Context) {
		encryptType := c.Query("encrypt_type")
		signature := c.Query("signature")
		timestamp := c.Query("timestamp")
		nonce := c.Query("nonce")

		switch encryptType {
		case "aes":
			token := verifySignReturnToken(signature, timestamp, nonce)
			if token == "" {
				srv.logger.Error("Verify sign empty token")
				return
			}

			msgSign := c.Query("msg_signature")

			var obj EncryptMsg
			err := c.BindWith(&obj, binding.XML)
			if err != nil {
				srv.logger.Errorw("Bind with XML failed", "error", err)
				return
			}

			if srv.ID != "" && !equal(obj.ToUserName, srv.ID) {
				srv.logger.Errorw("Wechat ID inconsistent", "id", srv.ID, "ToUserName", obj.ToUserName)
				return
			}

			computedSign := computeSign(token, timestamp, nonce, obj.Encrypt)
			if !equal(computedSign, msgSign) {
				srv.logger.Errorw("Signature inconsistent")
				return
			}

			encryptedMsg, err := base64.StdEncoding.DecodeString(obj.Encrypt)
			if err != nil {
				srv.logger.Errorw("Decode base64 string failed", "error", err)
				return
			}

			current, last := srv.GetAESKey()
			aesKey := current
			random, msg, appId, err := decryptMsg(encryptedMsg, []byte(aesKey))
			if err != nil {
				if last == "" {
					srv.logger.Errorw("Decrypt AES msg failed", "error", err)
					return
				}
				aesKey = last
				random, msg, appId, err = decryptMsg(encryptedMsg, []byte(aesKey))
				if err != nil {
					srv.logger.Errorw("Decrypt AES msg failed", "error", err)
					return
				}
			} else {
				srv.deleteLastAESKey()
			}
			if srv.appID != "" && string(appId) != srv.appID {
				srv.logger.Errorw("AppID inconsistent", "AppID", appId)
				return
			}

			var event Event
			if err = xml.Unmarshal(msg, &event); err != nil {
				srv.logger.Errorw("Unmarshal msg failed", "error", err)
				return
			}

			repBytes, err := xml.Marshal(handleMessage(&event))
			if err != nil {
				srv.logger.Errorw("Marshal msg failed", "error", err)
				return
			}

			encryptedRepBytes := encryptMsg(random, repBytes, appId, []byte(aesKey))
			encryptedRepStr := base64.StdEncoding.EncodeToString(encryptedRepBytes)
			repSignature := computeSign(token, timestamp, nonce, encryptedRepStr)

			type EncryptRepMsg struct {
				Encrypt      string
				MsgSignature string
				TimeStamp    string
				Nonce        string
			}

			err = c.XML(200, &EncryptRepMsg{encryptedRepStr, repSignature, timestamp, nonce})
			if err != nil {
				srv.logger.Errorw("Reply msg failed", "error", err)
			}

		case "", "raw":
			if !verifySign(c) {
				return
			}

			var event Event
			err := c.BindWith(&event, binding.XML)
			if err != nil {
				return
			}

			c.XML(200, handleMessage(&event))

		default:
			return
		}
	})

	handleAuthorize := func(c *mel.Context, url string, state string) {
		rep, err := srv.client.Client.Get(url)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
		defer rep.Body.Close()

		if rep.StatusCode != http.StatusOK {
			c.AbortWithError(http.StatusUnauthorized, fmt.Errorf("http.Status: %s", rep.Status))
			return
		}

		type Result struct {
			AccessToken  string `json:"access_token"`
			ExpiresIn    string `json:"expires_in"`
			RefreshToken string `json:"refresh_token"`
			OpenID       string `json:"openid"`
			Scope        string `json:"scope"`
			State        string `json:"state,omitempty"`
		}

		type ResultWithErr struct {
			Result
			Err
		}

		var result ResultWithErr
		err = json.NewDecoder(rep.Body).Decode(&result)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}

		if result.Code() != OK {
			c.AbortWithError(http.StatusUnauthorized, &result)
			return
		}

		result.State = state
		c.JSON(http.StatusOK, &result.Result)
		srv.logger.Infof("/token", "result", result.Result)
	}

	srv.Get(srv.urlPrefix+"/token", func(c *mel.Context) {
		code := c.Query("code")
		state := c.Query("state")

		url := fmt.Sprintf("https://api.weixin.qq.com/sns/oauth2/access_token?appid=%s&secret=%s&code=%s&grant_type=authorization_code", srv.client.appId, srv.client.appSecret, code)

		handleAuthorize(c, url, state)
	})

	srv.Get(srv.urlPrefix+"/refresh-token", func(c *mel.Context) {
		refreshToken := c.Query("refresh_token")

		url := fmt.Sprintf("https://api.weixin.qq.com/sns/oauth2/refresh_token?appid=%s&grant_type=refresh_token&refresh_token=%s", srv.client.appId, refreshToken)

		handleAuthorize(c, url, "")
	})

	srv.Get(srv.urlPrefix+"/signature", func(c *mel.Context) {
		timestamp := c.Query("timestamp")
		noncestr := c.Query("noncestr")
		url := c.Query("url")
		refresh := c.Query("refresh")

		var ticket string
		var err error
		if refresh != "" && (refresh == "true" || refresh == "True" || refresh == "1") {
			ticket, err = srv.client.RefreshTicket("")
		} else {
			ticket, err = srv.client.Ticket()
		}
		if err != nil {
			c.AbortWithError(http.StatusInternalServerError, err)
			return
		}

		strs := sort.StringSlice{
			"timestamp=" + timestamp,
			"noncestr=" + noncestr,
			"url=" + url,
			"jsapi_ticket=" + ticket,
		}
		strs.Sort()
		h := sha1.New()
		buf := bufio.NewWriterSize(h, 1024)
		for i, s := range strs {
			buf.WriteString(s)
			if i < len(strs)-1 {
				buf.WriteByte('&')
			}
		}
		buf.Flush()
		sign := hex.EncodeToString(h.Sum(nil))
		c.JSON(http.StatusOK, map[string]string{
			"signature": sign,
		})
		srv.logger.Infow("signature", "strs", strs, "sign", sign)
	})

	return srv
}

func computeSign(elements ...string) string {
	strs := sort.StringSlice(elements)
	strs.Sort()

	h := sha1.New()

	buf := bufio.NewWriterSize(h, 1024)
	for _, s := range strs {
		buf.WriteString(s)
	}
	buf.Flush()

	return hex.EncodeToString(h.Sum(nil))
}
<commit_msg>Update files<commit_after>package mp

import (
	"bufio"
	"bytes"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base64"
	"encoding/hex"
	"encoding/xml"
	"github.com/ridewindx/mel"
	"github.com/ridewindx/mel/binding"
	"sort"
	"sync"
	"sync/atomic"
	"unsafe"
	"go.uber.org/zap"
	"github.com/jiudaoyun/wechat"
	"net/http"
	"fmt"
	"encoding/json"
	"strings"
	"github.com/ridewindx/melware"
)

type Server struct {
	*mel.Mel
	urlPrefix string

	appID string // App ID
	ID    string // Wechat ID

	tokenMutex sync.Mutex
	token      unsafe.Pointer

	aesKeyMutex sync.Mutex
	aesKey      unsafe.Pointer

	client            *Client
	middlewares       []Handler
	messageHandlerMap map[string]Handler
	eventHandlerMap   map[string]Handler

	logger *zap.SugaredLogger
}

func (srv *Server) setURLPrefix(urlPrefix string) {
	if !strings.HasPrefix(urlPrefix, "/") {
		urlPrefix = "/" + urlPrefix
	}
	urlPrefix = strings.TrimRight(urlPrefix, "/")
	srv.urlPrefix = urlPrefix
}

func (srv *Server) SetID(id string) {
	srv.ID = id
}

func (srv *Server) SetAppID(appID string) {
	srv.appID = appID
}

func (srv *Server) SetClient(client *Client) {
	srv.client = client
}

type Token struct {
	current string
	last    string
}

type AESKey struct {
	current string
	last    string
}

func (srv *Server) GetToken() (string, string) {
	p := (*Token)(atomic.LoadPointer(&srv.token))
	if p != nil {
		return p.current, p.last
	}
	return "", ""
}

func (srv *Server) SetToken(token string) {
	if token == "" {
		return
	}

	srv.tokenMutex.Lock()
	defer srv.tokenMutex.Unlock()

	current, _ := srv.GetToken()
	if token == current {
		return
	}

	t := Token{
		current: token,
		last:    current,
	}
	atomic.StorePointer(&srv.token, unsafe.Pointer(&t))
}

func (srv *Server) deleteLastToken() {
	srv.tokenMutex.Lock()
	defer srv.tokenMutex.Unlock()

	current, last := srv.GetToken()
	if last == "" {
		return
	}

	t := Token{
		current: current,
	}
	atomic.StorePointer(&srv.token, unsafe.Pointer(&t))
}

func (srv *Server) GetAESKey() (string, string) {
	p := (*AESKey)(atomic.LoadPointer(&srv.aesKey))
	if p != nil {
		return p.current, p.last
	}
	return "", ""
}

func (srv *Server) SetAESKey(base64AESKey string) {
	if len(base64AESKey) != 43 {
		return
	}
	aesKey, err := base64.StdEncoding.DecodeString(base64AESKey + "=")
	if err != nil {
		return
	}

	srv.aesKeyMutex.Lock()
	defer srv.aesKeyMutex.Unlock()

	current, _ := srv.GetAESKey()
	if bytes.Equal(aesKey, []byte(current)) {
		return
	}

	k := AESKey{
		current: string(aesKey),
		last:    current,
	}
	atomic.StorePointer(&srv.aesKey, unsafe.Pointer(&k))
}

func (srv *Server) deleteLastAESKey() {
	srv.aesKeyMutex.Lock()
	defer srv.aesKeyMutex.Unlock()

	current, last := srv.GetAESKey()
	if last == "" {
		return
	}

	k := AESKey{
		current: current,
	}
	atomic.StorePointer(&srv.aesKey, unsafe.Pointer(&k))
}

func (srv *Server) Use(middlewares ...Handler) {
	srv.middlewares = append(srv.middlewares, middlewares...)
	if len(srv.middlewares)+1 > int(abortIndex) {
		panic("too many middlewares")
	}
}

func (srv *Server) HandleMessage(msgType string, handler Handler) {
	srv.messageHandlerMap[msgType] = handler
}

func (srv *Server) HandleEvent(eventType string, handler Handler) {
	srv.eventHandlerMap[eventType] = handler
}

func (srv *Server) GetVerifyFile(filename string, content []byte) {
	srv.Get(srv.urlPrefix+filename, func(c *mel.Context) {
		c.Data(200, "text/plain", content)
	})
}

func NewServer(token, aesKey string, urlPrefix ...string) *Server {
	srv := &Server{
		Mel:               mel.New(),
		messageHandlerMap: make(map[string]Handler),
		eventHandlerMap:   make(map[string]Handler),
		logger:            wechat.Sugar,
	}

	srv.SetToken(token)
	srv.SetAESKey(aesKey)

	if len(urlPrefix) > 0 {
		srv.setURLPrefix(urlPrefix[0])
	}

	equal := func(a, b string) bool {
		return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
	}

	verifySignReturnToken := func(signature, timestamp, nonce string) string {
		currentToken, lastToken := srv.GetToken()
		token := currentToken

		isValid := func() bool {
			computedSignature := computeSign(token, timestamp, nonce)
			return equal(signature, computedSignature)
		}

		if isValid() {
			srv.deleteLastToken()
			return token
		}

		if lastToken != "" {
			token = lastToken
			if isValid() {
				return token
			}
		}

		return ""
	}

	verifySign := func(c *mel.Context) bool {
		signature := c.Query("signature")
		timestamp := c.Query("timestamp")
		nonce := c.Query("nonce")

		return verifySignReturnToken(signature, timestamp, nonce) != ""
	}

	type EncryptMsg struct {
		ToUserName string `xml:"ToUserName"`
		Encrypt    string `xml:"Encrypt"`
	}

	cors := melware.CorsAllowAll()
	cors.AllowCredentials = false
	srv.Mel.Use(cors.Middleware())

	srv.Head("/", func(c *mel.Context) { // health check
		c.Status(200)
	})

	srv.Get(srv.urlPrefix+"/", func(c *mel.Context) {
		if verifySign(c) {
			echostr := c.Query("echostr")
			c.Text(200, echostr)
		}
	})

	handleMessage := func(event *Event) interface{} {
		var handler Handler
		var ok bool
		if event.Type == MessageEvent {
			handler, ok = srv.eventHandlerMap[event.Event]
		} else {
			handler, ok = srv.messageHandlerMap[event.Type]
		}
		if !ok {
			return nil // no registered handler, just respond with empty string
		}

		ctx := &Context{
			Client:   srv.client,
			index:    preStartIndex,
			handlers: append(srv.middlewares, handler),
			Event:    event,
		}

		ctx.Next()

		return ctx.response
	}

	srv.Post(srv.urlPrefix+"/", func(c *mel.Context) {
		encryptType := c.Query("encrypt_type")
		signature := c.Query("signature")
		timestamp := c.Query("timestamp")
		nonce := c.Query("nonce")

		switch encryptType {
		case "aes":
			token := verifySignReturnToken(signature, timestamp, nonce)
			if token == "" {
				srv.logger.Error("Verify sign empty token")
				return
			}

			msgSign := c.Query("msg_signature")

			var obj EncryptMsg
			err := c.BindWith(&obj, binding.XML)
			if err != nil {
				srv.logger.Errorw("Bind with XML failed", "error", err)
				return
			}

			if srv.ID != "" && !equal(obj.ToUserName, srv.ID) {
				srv.logger.Errorw("Wechat ID inconsistent", "id", srv.ID, "ToUserName", obj.ToUserName)
				return
			}

			computedSign := computeSign(token, timestamp, nonce, obj.Encrypt)
			if !equal(computedSign, msgSign) {
				srv.logger.Errorw("Signature inconsistent")
				return
			}

			encryptedMsg, err := base64.StdEncoding.DecodeString(obj.Encrypt)
			if err != nil {
				srv.logger.Errorw("Decode base64 string failed", "error", err)
				return
			}

			current, last := srv.GetAESKey()
			aesKey := current
			random, msg, appId, err := decryptMsg(encryptedMsg, []byte(aesKey))
			if err != nil {
				if last == "" {
					srv.logger.Errorw("Decrypt AES msg failed", "error", err)
					return
				}
				aesKey = last
				random, msg, appId, err = decryptMsg(encryptedMsg, []byte(aesKey))
				if err != nil {
					srv.logger.Errorw("Decrypt AES msg failed", "error", err)
					return
				}
			} else {
				srv.deleteLastAESKey()
			}
			if srv.appID != "" && string(appId) != srv.appID {
				srv.logger.Errorw("AppID inconsistent", "AppID", appId)
				return
			}

			var event Event
			if err = xml.Unmarshal(msg, &event); err != nil {
				srv.logger.Errorw("Unmarshal msg failed", "error", err)
				return
			}

			repBytes, err := xml.Marshal(handleMessage(&event))
			if err != nil {
				srv.logger.Errorw("Marshal msg failed", "error", err)
				return
			}

			encryptedRepBytes := encryptMsg(random, repBytes, appId, []byte(aesKey))
			encryptedRepStr := base64.StdEncoding.EncodeToString(encryptedRepBytes)
			repSignature := computeSign(token, timestamp, nonce, encryptedRepStr)

			type EncryptRepMsg struct {
				Encrypt      string
				MsgSignature string
				TimeStamp    string
				Nonce        string
			}

			err = c.XML(200, &EncryptRepMsg{encryptedRepStr, repSignature, timestamp, nonce})
			if err != nil {
				srv.logger.Errorw("Reply msg failed", "error", err)
			}

		case "", "raw":
			if !verifySign(c) {
				return
			}

			var event Event
			err := c.BindWith(&event, binding.XML)
			if err != nil {
				return
			}

			c.XML(200, handleMessage(&event))

		default:
			return
		}
	})

	handleAuthorize := func(c *mel.Context, url string, state string) {
		rep, err := srv.client.Client.Get(url)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
		defer rep.Body.Close()

		if rep.StatusCode != http.StatusOK {
			c.AbortWithError(http.StatusUnauthorized, fmt.Errorf("http.Status: %s", rep.Status))
			return
		}

		type Result struct {
			AccessToken  string `json:"access_token"`
			ExpiresIn    string `json:"expires_in"`
			RefreshToken string `json:"refresh_token"`
			OpenID       string `json:"openid"`
			Scope        string `json:"scope"`
			State        string `json:"state,omitempty"`
		}

		type ResultWithErr struct {
			Result
			Err
		}

		var result ResultWithErr
		err = json.NewDecoder(rep.Body).Decode(&result)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}

		if result.Code() != OK {
			c.AbortWithError(http.StatusUnauthorized, &result)
			return
		}

		result.State = state
		c.JSON(http.StatusOK, &result.Result)
		srv.logger.Infof("/token", "result", result.Result)
	}

	srv.Get(srv.urlPrefix+"/token", func(c *mel.Context) {
		code := c.Query("code")
		state := c.Query("state")

		url := fmt.Sprintf("https://api.weixin.qq.com/sns/oauth2/access_token?appid=%s&secret=%s&code=%s&grant_type=authorization_code", srv.client.appId, srv.client.appSecret, code)

		handleAuthorize(c, url, state)
	})

	srv.Get(srv.urlPrefix+"/refresh-token", func(c *mel.Context) {
		refreshToken := c.Query("refresh_token")

		url := fmt.Sprintf("https://api.weixin.qq.com/sns/oauth2/refresh_token?appid=%s&grant_type=refresh_token&refresh_token=%s", srv.client.appId, refreshToken)

		handleAuthorize(c, url, "")
	})

	srv.Get(srv.urlPrefix+"/signature", func(c *mel.Context) {
		timestamp := c.Query("timestamp")
		noncestr := c.Query("noncestr")
		url := c.Query("url")
		refresh := c.Query("refresh")

		var ticket string
		var err error
		if refresh != "" && (refresh == "true" || refresh == "True" || refresh == "1") {
			ticket, err = srv.client.RefreshTicket("")
		} else {
			ticket, err = srv.client.Ticket()
		}
		if err != nil {
			c.AbortWithError(http.StatusInternalServerError, err)
			return
		}

		strs := sort.StringSlice{
			"timestamp=" + timestamp,
			"noncestr=" + noncestr,
			"url=" + url,
			"jsapi_ticket=" + ticket,
		}
		strs.Sort()
		h := sha1.New()
		buf := bufio.NewWriterSize(h, 1024)
		for i, s := range strs {
			buf.WriteString(s)
			if i < len(strs)-1 {
				buf.WriteByte('&')
			}
		}
		buf.Flush()
		sign := hex.EncodeToString(h.Sum(nil))
		c.JSON(http.StatusOK, map[string]string{
			"signature": sign,
		})
		srv.logger.Infow("signature", "strs", strs, "sign", sign)
	})

	return srv
}

func computeSign(elements ...string) string {
	strs := sort.StringSlice(elements)
	strs.Sort()

	h := sha1.New()

	buf := bufio.NewWriterSize(h, 1024)
	for _, s := range strs {
		buf.WriteString(s)
	}
	buf.Flush()

	return hex.EncodeToString(h.Sum(nil))
}
<|endoftext|>