<commit_before>// Copyright 2013 Joe Walnes and the websocketd team.
// All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package libwebsocketd

import (
	"errors"
	"fmt"
	"golang.org/x/net/websocket"
	"net"
	"net/http"
	"net/http/cgi"
	"net/textproto"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
)

var ForkNotAllowedError = errors.New("too many forks active")

// WebsocketdServer presents http.Handler interface for requests libwebsocketd is handling.
type WebsocketdServer struct {
	Config *Config
	Log    *LogScope
	forks  chan byte
}

// NewWebsocketdServer creates WebsocketdServer struct with pre-determined config, logscope and maxforks limit
func NewWebsocketdServer(config *Config, log *LogScope, maxforks int) *WebsocketdServer {
	mux := &WebsocketdServer{
		Config: config,
		Log:    log,
	}
	if maxforks > 0 {
		mux.forks = make(chan byte, maxforks)
	}
	return mux
}

// wshandshake returns closure to verify websocket origin header according to configured rules
func (h *WebsocketdServer) wshandshake(log *LogScope) func(*websocket.Config, *http.Request) error {
	return func(wsconf *websocket.Config, req *http.Request) error {
		if len(h.Config.Headers)+len(h.Config.HeadersWs) > 0 {
			if wsconf.Header == nil {
				wsconf.Header = http.Header(make(map[string][]string))
			}
			pushHeaders(wsconf.Header, h.Config.Headers)
			pushHeaders(wsconf.Header, h.Config.HeadersWs)
		}
		return checkOrigin(wsconf, req, h.Config, log)
	}
}

func splitMimeHeader(s string) (string, string) {
	p := strings.IndexByte(s, ':')
	if p < 0 {
		return s, ""
	}
	key := textproto.CanonicalMIMEHeaderKey(s[:p])

	for p = p + 1; p < len(s); p++ {
		if s[p] != ' ' {
			break
		}
	}
	return key, s[p:]
}

func pushHeaders(h http.Header, hdrs []string) {
	for _, hstr := range hdrs {
		h.Add(splitMimeHeader(hstr))
	}
}

// ServeHTTP muxes between WebSocket handler, CGI handler, DevConsole, Static HTML or 404.
func (h *WebsocketdServer) ServeHTTP(w http.ResponseWriter, req *http.Request) {

	log := h.Log.NewLevel(h.Log.LogFunc)
	log.Associate("url", h.TellURL("http", req.Host, req.RequestURI))

	pushHeaders(w.Header(), h.Config.Headers)
	if h.Config.CommandName != "" || h.Config.UsingScriptDir {
		hdrs := req.Header
		upgradeRe := regexp.MustCompile("(?i)(^|[,\\s])Upgrade($|[,\\s])")
		// WebSocket, limited to size of h.forks
		if strings.ToLower(hdrs.Get("Upgrade")) == "websocket" && upgradeRe.MatchString(hdrs.Get("Connection")) {
			if h.noteForkCreated() == nil {
				defer h.noteForkCompled()

				handler, err := NewWebsocketdHandler(h, req, log)
				if err != nil {
					if err == ScriptNotFoundError {
						log.Access("session", "NOT FOUND: %s", err)
						http.Error(w, "404 Not Found", 404)
					} else {
						log.Access("session", "INTERNAL ERROR: %s", err)
						http.Error(w, "500 Internal Server Error", 500)
					}
					return
				}

				// Now we are ready for connection upgrade dance...
				wsServer := &websocket.Server{
					Handshake: h.wshandshake(log),
					Handler:   handler.wshandler(log),
				}
				wsServer.ServeHTTP(w, req)
			} else {
				log.Error("http", "Max of possible forks already active, upgrade rejected")
				http.Error(w, "429 Too Many Requests", 429)
			}
			return
		}
	}

	pushHeaders(w.Header(), h.Config.HeadersHTTP)

	// Dev console (if enabled)
	if h.Config.DevConsole {
		log.Access("http", "DEVCONSOLE")
		content := ConsoleContent
		content = strings.Replace(content, "{{license}}", License, -1)
		content = strings.Replace(content, "{{addr}}", h.TellURL("ws", req.Host, req.RequestURI), -1)
		http.ServeContent(w, req, ".html", h.Config.StartupTime, strings.NewReader(content))
		return
	}

	// CGI scripts, limited to size of h.forks
	if h.Config.CgiDir != "" {
		filePath := path.Join(h.Config.CgiDir, fmt.Sprintf(".%s", filepath.FromSlash(req.URL.Path)))
		if fi, err := os.Stat(filePath); err == nil && !fi.IsDir() {

			log.Associate("cgiscript", filePath)
			if h.noteForkCreated() == nil {
				defer h.noteForkCompled()

				// Make variables to supplement cgi... Environ it uses will show empty list.
				envlen := len(h.Config.ParentEnv)
				cgienv := make([]string, envlen+1)
				if envlen > 0 {
					copy(cgienv, h.Config.ParentEnv)
				}
				cgienv[envlen] = "SERVER_SOFTWARE=" + h.Config.ServerSoftware
				cgiHandler := &cgi.Handler{
					Path: filePath,
					Env: []string{
						"SERVER_SOFTWARE=" + h.Config.ServerSoftware,
					},
				}
				log.Access("http", "CGI")
				cgiHandler.ServeHTTP(w, req)
			} else {
				log.Error("http", "Fork not allowed since maxforks amount has been reached. CGI was not run.")
				http.Error(w, "429 Too Many Requests", 429)
			}
			return
		}
	}

	// Static files
	if h.Config.StaticDir != "" {
		handler := http.FileServer(http.Dir(h.Config.StaticDir))
		log.Access("http", "STATIC")
		handler.ServeHTTP(w, req)
		return
	}

	// 404
	log.Access("http", "NOT FOUND")
	http.NotFound(w, req)
}

var canonicalHostname string

// TellURL is a helper function that changes http to https or ws to wss in case if SSL is used
func (h *WebsocketdServer) TellURL(scheme, host, path string) string {
	if host[0] == ':' {
		if canonicalHostname == "" {
			var err error
			canonicalHostname, err = os.Hostname()
			if err != nil {
				canonicalHostname = "UNKNOWN"
			}
		}
		host = canonicalHostname + host
	}
	if h.Config.Ssl {
		return scheme + "s://" + host + path
	}
	return scheme + "://" + host + path
}

func (h *WebsocketdServer) noteForkCreated() error {
	// note that forks can be nil since the construct could've been created by
	// someone who is not using NewWebsocketdServer
	if h.forks != nil {
		select {
		case h.forks <- 1:
			return nil
		default:
			return ForkNotAllowedError
		}
	} else {
		return nil
	}
}

func (h *WebsocketdServer) noteForkCompled() {
	if h.forks != nil { // see comment in noteForkCreated
		select {
		case <-h.forks:
			return
		default:
			// This could only happen if the completion handler called more times than creation handler above
			// Code should be audited to not allow this to happen, it's desired to have test that would
			// make sure this is impossible but it is not exist yet.
			panic("Cannot deplet number of allowed forks, something is not right in code!")
		}
	}
	return
}

func checkOrigin(wsconf *websocket.Config, req *http.Request, config *Config, log *LogScope) (err error) {
	// check for origin to be correct in future
	// handshaker triggers answering with 403 if error was returned
	// We keep behavior of original handshaker that populates this field
	origin := req.Header.Get("Origin")
	if origin == "" || (origin == "null" && config.AllowOrigins == nil) {
		// we don't want to trust string "null" if there is any
		// enforcements are active
		req.Header.Set("Origin", "file:")
	}

	wsconf.Origin, err = websocket.Origin(wsconf, req)
	if err == nil && wsconf.Origin == nil {
		log.Access("session", "rejected null origin")
		return fmt.Errorf("null origin not allowed")
	}
	if err != nil {
		log.Access("session", "Origin parsing error: %s", err)
		return err
	}
	log.Associate("origin", wsconf.Origin.String())

	// If some origin restrictions are present:
	if config.SameOrigin || config.AllowOrigins != nil {
		originServer, originPort, err := tellHostPort(wsconf.Origin.Host, wsconf.Origin.Scheme == "https")
		if err != nil {
			log.Access("session", "Origin hostname parsing error: %s", err)
			return err
		}
		if config.SameOrigin {
			localServer, localPort, err := tellHostPort(req.Host, req.TLS != nil)
			if err != nil {
				log.Access("session", "Request hostname parsing error: %s", err)
				return err
			}
			if originServer != localServer || originPort != localPort {
				log.Access("session", "Same origin policy mismatch")
				return fmt.Errorf("same origin policy violated")
			}
		}
		if config.AllowOrigins != nil {
			matchFound := false
			for _, allowed := range config.AllowOrigins {
				if pos := strings.Index(allowed, "://"); pos > 0 {
					// allowed schema has to match
					allowedURL, err := url.Parse(allowed)
					if err != nil {
						continue // pass bad URLs in origin list
					}
					if allowedURL.Scheme != wsconf.Origin.Scheme {
						continue // mismatch
					}
					allowed = allowed[pos+3:]
				}
				allowServer, allowPort, err := tellHostPort(allowed, false)
				if err != nil {
					continue // unparseable
				}
				if allowPort == "80" && allowed[len(allowed)-3:] != ":80" {
					// any port is allowed, host names need to match
					matchFound = allowServer == originServer
				} else {
					// exact match of host names and ports
					matchFound = allowServer == originServer && allowPort == originPort
				}
				if matchFound {
					break
				}
			}
			if !matchFound {
				log.Access("session", "Origin is not listed in allowed list")
				return fmt.Errorf("origin list matches were not found")
			}
		}
	}
	return nil
}

func tellHostPort(host string, ssl bool) (server, port string, err error) {
	server, port, err = net.SplitHostPort(host)
	if err != nil {
		if addrerr, ok := err.(*net.AddrError); ok && strings.Contains(addrerr.Err, "missing port") {
			server = host
			if ssl {
				port = "443"
			} else {
				port = "80"
			}
			err = nil
		}
	}
	return server, port, err
}
<commit_msg>Fixes #163, panic in TellURL<commit_after>// Copyright 2013 Joe Walnes and the websocketd team.
// All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package libwebsocketd

import (
	"errors"
	"fmt"
	"golang.org/x/net/websocket"
	"net"
	"net/http"
	"net/http/cgi"
	"net/textproto"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
)

var ForkNotAllowedError = errors.New("too many forks active")

// WebsocketdServer presents http.Handler interface for requests libwebsocketd is handling.
type WebsocketdServer struct {
	Config *Config
	Log    *LogScope
	forks  chan byte
}

// NewWebsocketdServer creates WebsocketdServer struct with pre-determined config, logscope and maxforks limit
func NewWebsocketdServer(config *Config, log *LogScope, maxforks int) *WebsocketdServer {
	mux := &WebsocketdServer{
		Config: config,
		Log:    log,
	}
	if maxforks > 0 {
		mux.forks = make(chan byte, maxforks)
	}
	return mux
}

// wshandshake returns closure to verify websocket origin header according to configured rules
func (h *WebsocketdServer) wshandshake(log *LogScope) func(*websocket.Config, *http.Request) error {
	return func(wsconf *websocket.Config, req *http.Request) error {
		if len(h.Config.Headers)+len(h.Config.HeadersWs) > 0 {
			if wsconf.Header == nil {
				wsconf.Header = http.Header(make(map[string][]string))
			}
			pushHeaders(wsconf.Header, h.Config.Headers)
			pushHeaders(wsconf.Header, h.Config.HeadersWs)
		}
		return checkOrigin(wsconf, req, h.Config, log)
	}
}

func splitMimeHeader(s string) (string, string) {
	p := strings.IndexByte(s, ':')
	if p < 0 {
		return s, ""
	}
	key := textproto.CanonicalMIMEHeaderKey(s[:p])

	for p = p + 1; p < len(s); p++ {
		if s[p] != ' ' {
			break
		}
	}
	return key, s[p:]
}

func pushHeaders(h http.Header, hdrs []string) {
	for _, hstr := range hdrs {
		h.Add(splitMimeHeader(hstr))
	}
}

// ServeHTTP muxes between WebSocket handler, CGI handler, DevConsole, Static HTML or 404.
func (h *WebsocketdServer) ServeHTTP(w http.ResponseWriter, req *http.Request) {

	log := h.Log.NewLevel(h.Log.LogFunc)
	log.Associate("url", h.TellURL("http", req.Host, req.RequestURI))

	pushHeaders(w.Header(), h.Config.Headers)
	if h.Config.CommandName != "" || h.Config.UsingScriptDir {
		hdrs := req.Header
		upgradeRe := regexp.MustCompile("(?i)(^|[,\\s])Upgrade($|[,\\s])")
		// WebSocket, limited to size of h.forks
		if strings.ToLower(hdrs.Get("Upgrade")) == "websocket" && upgradeRe.MatchString(hdrs.Get("Connection")) {
			if h.noteForkCreated() == nil {
				defer h.noteForkCompled()

				handler, err := NewWebsocketdHandler(h, req, log)
				if err != nil {
					if err == ScriptNotFoundError {
						log.Access("session", "NOT FOUND: %s", err)
						http.Error(w, "404 Not Found", 404)
					} else {
						log.Access("session", "INTERNAL ERROR: %s", err)
						http.Error(w, "500 Internal Server Error", 500)
					}
					return
				}

				// Now we are ready for connection upgrade dance...
				wsServer := &websocket.Server{
					Handshake: h.wshandshake(log),
					Handler:   handler.wshandler(log),
				}
				wsServer.ServeHTTP(w, req)
			} else {
				log.Error("http", "Max of possible forks already active, upgrade rejected")
				http.Error(w, "429 Too Many Requests", 429)
			}
			return
		}
	}

	pushHeaders(w.Header(), h.Config.HeadersHTTP)

	// Dev console (if enabled)
	if h.Config.DevConsole {
		log.Access("http", "DEVCONSOLE")
		content := ConsoleContent
		content = strings.Replace(content, "{{license}}", License, -1)
		content = strings.Replace(content, "{{addr}}", h.TellURL("ws", req.Host, req.RequestURI), -1)
		http.ServeContent(w, req, ".html", h.Config.StartupTime, strings.NewReader(content))
		return
	}

	// CGI scripts, limited to size of h.forks
	if h.Config.CgiDir != "" {
		filePath := path.Join(h.Config.CgiDir, fmt.Sprintf(".%s", filepath.FromSlash(req.URL.Path)))
		if fi, err := os.Stat(filePath); err == nil && !fi.IsDir() {

			log.Associate("cgiscript", filePath)
			if h.noteForkCreated() == nil {
				defer h.noteForkCompled()

				// Make variables to supplement cgi... Environ it uses will show empty list.
				envlen := len(h.Config.ParentEnv)
				cgienv := make([]string, envlen+1)
				if envlen > 0 {
					copy(cgienv, h.Config.ParentEnv)
				}
				cgienv[envlen] = "SERVER_SOFTWARE=" + h.Config.ServerSoftware
				cgiHandler := &cgi.Handler{
					Path: filePath,
					Env: []string{
						"SERVER_SOFTWARE=" + h.Config.ServerSoftware,
					},
				}
				log.Access("http", "CGI")
				cgiHandler.ServeHTTP(w, req)
			} else {
				log.Error("http", "Fork not allowed since maxforks amount has been reached. CGI was not run.")
				http.Error(w, "429 Too Many Requests", 429)
			}
			return
		}
	}

	// Static files
	if h.Config.StaticDir != "" {
		handler := http.FileServer(http.Dir(h.Config.StaticDir))
		log.Access("http", "STATIC")
		handler.ServeHTTP(w, req)
		return
	}

	// 404
	log.Access("http", "NOT FOUND")
	http.NotFound(w, req)
}

var canonicalHostname string

// TellURL is a helper function that changes http to https or ws to wss in case if SSL is used
func (h *WebsocketdServer) TellURL(scheme, host, path string) string {
	if len(host) > 0 && host[0] == ':' {
		if canonicalHostname == "" {
			var err error
			canonicalHostname, err = os.Hostname()
			if err != nil {
				canonicalHostname = "UNKNOWN"
			}
		}
		host = canonicalHostname + host
	}
	if h.Config.Ssl {
		return scheme + "s://" + host + path
	}
	return scheme + "://" + host + path
}

func (h *WebsocketdServer) noteForkCreated() error {
	// note that forks can be nil since the construct could've been created by
	// someone who is not using NewWebsocketdServer
	if h.forks != nil {
		select {
		case h.forks <- 1:
			return nil
		default:
			return ForkNotAllowedError
		}
	} else {
		return nil
	}
}

func (h *WebsocketdServer) noteForkCompled() {
	if h.forks != nil { // see comment in noteForkCreated
		select {
		case <-h.forks:
			return
		default:
			// This could only happen if the completion handler called more times than creation handler above
			// Code should be audited to not allow this to happen, it's desired to have test that would
			// make sure this is impossible but it is not exist yet.
			panic("Cannot deplet number of allowed forks, something is not right in code!")
		}
	}
	return
}

func checkOrigin(wsconf *websocket.Config, req *http.Request, config *Config, log *LogScope) (err error) {
	// check for origin to be correct in future
	// handshaker triggers answering with 403 if error was returned
	// We keep behavior of original handshaker that populates this field
	origin := req.Header.Get("Origin")
	if origin == "" || (origin == "null" && config.AllowOrigins == nil) {
		// we don't want to trust string "null" if there is any
		// enforcements are active
		req.Header.Set("Origin", "file:")
	}

	wsconf.Origin, err = websocket.Origin(wsconf, req)
	if err == nil && wsconf.Origin == nil {
		log.Access("session", "rejected null origin")
		return fmt.Errorf("null origin not allowed")
	}
	if err != nil {
		log.Access("session", "Origin parsing error: %s", err)
		return err
	}
	log.Associate("origin", wsconf.Origin.String())

	// If some origin restrictions are present:
	if config.SameOrigin || config.AllowOrigins != nil {
		originServer, originPort, err := tellHostPort(wsconf.Origin.Host, wsconf.Origin.Scheme == "https")
		if err != nil {
			log.Access("session", "Origin hostname parsing error: %s", err)
			return err
		}
		if config.SameOrigin {
			localServer, localPort, err := tellHostPort(req.Host, req.TLS != nil)
			if err != nil {
				log.Access("session", "Request hostname parsing error: %s", err)
				return err
			}
			if originServer != localServer || originPort != localPort {
				log.Access("session", "Same origin policy mismatch")
				return fmt.Errorf("same origin policy violated")
			}
		}
		if config.AllowOrigins != nil {
			matchFound := false
			for _, allowed := range config.AllowOrigins {
				if pos := strings.Index(allowed, "://"); pos > 0 {
					// allowed schema has to match
					allowedURL, err := url.Parse(allowed)
					if err != nil {
						continue // pass bad URLs in origin list
					}
					if allowedURL.Scheme != wsconf.Origin.Scheme {
						continue // mismatch
					}
					allowed = allowed[pos+3:]
				}
				allowServer, allowPort, err := tellHostPort(allowed, false)
				if err != nil {
					continue // unparseable
				}
				if allowPort == "80" && allowed[len(allowed)-3:] != ":80" {
					// any port is allowed, host names need to match
					matchFound = allowServer == originServer
				} else {
					// exact match of host names and ports
					matchFound = allowServer == originServer && allowPort == originPort
				}
				if matchFound {
					break
				}
			}
			if !matchFound {
				log.Access("session", "Origin is not listed in allowed list")
				return fmt.Errorf("origin list matches were not found")
			}
		}
	}
	return nil
}

func tellHostPort(host string, ssl bool) (server, port string, err error) {
	server, port, err = net.SplitHostPort(host)
	if err != nil {
		if addrerr, ok := err.(*net.AddrError); ok && strings.Contains(addrerr.Err, "missing port") {
			server = host
			if ssl {
				port = "443"
			} else {
				port = "80"
			}
			err = nil
		}
	}
	return server, port, err
}
<|endoftext|>


<commit_before>// The MIT License (MIT)

// Copyright (c) 2016 Maciej Borzecki

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

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

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

import (
	"bytes"
	"encoding/xml"
	"github.com/bboozzoo/q3stats/util"
	"github.com/pkg/errors"
	"io"
	"io/ioutil"
	"log"
	"os"
)

type RawStat struct {
	Name  string `xml:"name,attr"`
	Value int    `xml:"value,attr"`
}

type RawItem struct {
	Name    string `xml:"name,attr"`
	Pickups uint    `xml:"pickups,attr"`
	Time    uint    `xml:"time,attr"`
}

type RawWeapon struct {
	Name  string `xml:"name,attr"`
	Hits  uint   `xml:"hits,attr"`
	Shots uint   `xml:"shots,attr"`
	Kills uint   `xml:"kills,attr"`
}

type RawPlayer struct {
	Name     string      `xml:"name,attr"`
	Stats    []RawStat   `xml:"stat"`
	Items    []RawItem   `xml:"items>item"`
	Weapons  []RawWeapon `xml:"weapons>weapon"`
	Powerups []RawItem   `xml:"powerups>item"`
}

type RawMatch struct {
	Datetime string      `xml:"datetime,attr"`
	Map      string      `xml:"map,attr"`
	Type     string      `xml:"type,attr"`
	Duration uint        `xml:"duration,attr"`
	Players  []RawPlayer `xml:"player"`
	DataHash string
}

func LoadMatch(src io.Reader) (*RawMatch, error) {
	raw, err := ioutil.ReadAll(src)
	if err != nil {
		return nil, errors.Wrap(err, "failed to load match data")
	}

	var match RawMatch
	err = xml.Unmarshal(raw, &match)
	if err != nil {
		return nil, errors.Wrap(err, "unmarshal failed")
	}

	match.DataHash = util.DataHash(raw)

	log.Printf("match: %+v", match)
	return &match, nil
}

func LoadMatchFile(path string) (*RawMatch, error) {

	in, err := os.Open(path)
	if err != nil {
		return nil, errors.Wrap(err, "failed to open match file")
	}
	defer in.Close()

	return LoadMatch(in)
}

func LoadMatchData(data []byte) (*RawMatch, error) {
	buf := bytes.NewBuffer(data)
	return LoadMatch(buf)
}
<commit_msg>loader/match: document item Time field<commit_after>// The MIT License (MIT)

// Copyright (c) 2016 Maciej Borzecki

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

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

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

import (
	"bytes"
	"encoding/xml"
	"github.com/bboozzoo/q3stats/util"
	"github.com/pkg/errors"
	"io"
	"io/ioutil"
	"log"
	"os"
)

type RawStat struct {
	Name  string `xml:"name,attr"`
	Value int    `xml:"value,attr"`
}

type RawItem struct {
	Name    string `xml:"name,attr"`
	Pickups uint   `xml:"pickups,attr"`
	// item held time in ms
	Time uint `xml:"time,attr"`
}

type RawWeapon struct {
	Name  string `xml:"name,attr"`
	Hits  uint   `xml:"hits,attr"`
	Shots uint   `xml:"shots,attr"`
	Kills uint   `xml:"kills,attr"`
}

type RawPlayer struct {
	Name     string      `xml:"name,attr"`
	Stats    []RawStat   `xml:"stat"`
	Items    []RawItem   `xml:"items>item"`
	Weapons  []RawWeapon `xml:"weapons>weapon"`
	Powerups []RawItem   `xml:"powerups>item"`
}

type RawMatch struct {
	Datetime string      `xml:"datetime,attr"`
	Map      string      `xml:"map,attr"`
	Type     string      `xml:"type,attr"`
	Duration uint        `xml:"duration,attr"`
	Players  []RawPlayer `xml:"player"`
	DataHash string
}

func LoadMatch(src io.Reader) (*RawMatch, error) {
	raw, err := ioutil.ReadAll(src)
	if err != nil {
		return nil, errors.Wrap(err, "failed to load match data")
	}

	var match RawMatch
	err = xml.Unmarshal(raw, &match)
	if err != nil {
		return nil, errors.Wrap(err, "unmarshal failed")
	}

	match.DataHash = util.DataHash(raw)

	log.Printf("match: %+v", match)
	return &match, nil
}

func LoadMatchFile(path string) (*RawMatch, error) {

	in, err := os.Open(path)
	if err != nil {
		return nil, errors.Wrap(err, "failed to open match file")
	}
	defer in.Close()

	return LoadMatch(in)
}

func LoadMatchData(data []byte) (*RawMatch, error) {
	buf := bytes.NewBuffer(data)
	return LoadMatch(buf)
}
<|endoftext|>


<commit_before>package rtmptee

import (
	"os"
	"os/signal"
	"runtime/debug"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/deckarep/golang-set"
	"github.com/pwaller/barrier"
	"github.com/rcrowley/go-metrics"
	"github.com/vrischmann/go-metrics-influxdb"
)

type App struct {
	Config *Config
	Flows  map[string][]*Flow
	quit   barrier.Barrier
}

func NewApp(config *Config) *App {
	return &App{
		Config: config,
		Flows:  make(map[string][]*Flow),
	}
}

func (app *App) Run() error {
	go app.handleSigint()
	go app.handleSigusr1()

	if app.Config.Times.IdleTime > 0 {
		go ShowIdleness(app.Config.Times.IdleTime)
	}

	if app.Config.Metrics.Influx != nil {
		go influxdb.InfluxDB(
			metrics.DefaultRegistry,
			10*time.Second,
			app.Config.Metrics.Influx.Host,
			app.Config.Metrics.Influx.Database,
			app.Config.Metrics.Influx.Username,
			app.Config.Metrics.Influx.Password,
		)
	}

	prevStreams := mapset.NewSet()
	nginxRtmp := NewNginxRtmp(app.Config)

	ticker := TickNow(app.Config.Times.NginxRtmpPollInterval)
	resetTimer := NewRestartableTimer(app.Config.Times.NginxRtmpServerTimeout)

	numStreamsMetric := metrics.GetOrRegister("streams", metrics.NewGauge()).(metrics.Gauge)
	for {
		select {
		case <-resetTimer.C:
			log.Warn("No reply from server, assuming all streams gone")
			resetTimer.Stop()
			for stream := range prevStreams.Iter() {
				app.removeStream(stream.(string))
			}
			prevStreams = mapset.NewSet()

		case <-ticker:
			curStreams, err := nginxRtmp.GetActiveStreams()
			if err != nil {
				Catch(err)
				continue
			}
			resetTimer.Restart()

			numStreamsMetric.Update(int64(curStreams.Cardinality()))

			for stream := range curStreams.Difference(prevStreams).Iter() {
				app.addStream(stream.(string))
			}
			for stream := range prevStreams.Difference(curStreams).Iter() {
				app.removeStream(stream.(string))
			}

			// All streams gone? => Good time for a GC run
			if prevStreams.Cardinality() > 0 && curStreams.Cardinality() == 0 {
				debug.FreeOSMemory()
			}

			prevStreams = curStreams

		case <-app.quit.Barrier():
			resetTimer.Stop()
			for _, flows := range app.Flows {
				for _, flow := range flows {
					flow.Stop()
				}
			}
			return nil
		}
	}
}

func (app *App) addStream(stream string) {
	logged := false

	for flowName, flowConfig := range app.Config.Flows {
		if flowConfig.Regexp.MatchString(stream) {
			if !logged {
				logged = true
				log.WithFields(log.Fields{
					"stream": stream,
					"match":  true,
				}).Warn("New stream")
			}

			app.addFlow(flowName, stream, flowConfig.Source, flowConfig.Sinks)
		}
	}

	if !logged {
		log.WithFields(log.Fields{
			"stream": stream,
			"match":  false,
		}).Debug("New stream, ignoring")
	}
}

func (app *App) addFlow(name string, stream string, sourceTemplate CmdData, sinkTemplates map[string]CmdData) {
	vars := map[string]string{
		"{stream}": stream,
	}

	source := sourceTemplate.Replace(vars)
	sinks := make(map[string]CmdData, len(sinkTemplates))
	for sinkName, sinkTemplate := range sinkTemplates {
		sinks[sinkName] = sinkTemplate.Replace(vars)
	}

	flow := NewFlow(name, app.Config, source, sinks, log.WithFields(log.Fields{
		"name":   name,
		"stream": stream,
	}))
	flow.Start()

	if _, ok := app.Flows[stream]; !ok {
		app.Flows[stream] = make([]*Flow, 0, 1)
	}
	app.Flows[stream] = append(app.Flows[stream], flow)
}

func (app *App) removeStream(stream string) {
	log.WithFields(log.Fields{
		"stream": stream,
	}).Warn("Stream gone")

	flows, ok := app.Flows[stream]
	if !ok {
		return
	}

	for _, flow := range flows {
		flow.log.Info("Stopping flow")
		flow.Stop()
	}

	delete(app.Flows, stream)
}

func (app *App) handleSigint() {
	channel := make(chan os.Signal, 1)
	signal.Notify(channel, os.Interrupt)
	<-channel
	if app.Config.Debug {
		panic("Interrupted")
	} else {
		log.Info("interrupted, shutting down...")
		app.quit.Fall()
	}
}

func (app *App) handleSigusr1() {
	channel := make(chan os.Signal, 1)
	signal.Notify(channel, syscall.SIGUSR1)
	for range channel {
		panic("Debug")
	}
}
<commit_msg>Make SIGUSR1 write the stacktrace to a file<commit_after>package rtmptee

import (
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"runtime/debug"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/deckarep/golang-set"
	"github.com/pwaller/barrier"
	"github.com/rcrowley/go-metrics"
	"github.com/vrischmann/go-metrics-influxdb"
)

type App struct {
	Config *Config
	Flows  map[string][]*Flow
	quit   barrier.Barrier
}

func NewApp(config *Config) *App {
	return &App{
		Config: config,
		Flows:  make(map[string][]*Flow),
	}
}

func (app *App) Run() error {
	go app.handleSigint()
	go app.handleSigusr1()

	if app.Config.Times.IdleTime > 0 {
		go ShowIdleness(app.Config.Times.IdleTime)
	}

	if app.Config.Metrics.Influx != nil {
		go influxdb.InfluxDB(
			metrics.DefaultRegistry,
			10*time.Second,
			app.Config.Metrics.Influx.Host,
			app.Config.Metrics.Influx.Database,
			app.Config.Metrics.Influx.Username,
			app.Config.Metrics.Influx.Password,
		)
	}

	prevStreams := mapset.NewSet()
	nginxRtmp := NewNginxRtmp(app.Config)

	ticker := TickNow(app.Config.Times.NginxRtmpPollInterval)
	resetTimer := NewRestartableTimer(app.Config.Times.NginxRtmpServerTimeout)

	numStreamsMetric := metrics.GetOrRegister("streams", metrics.NewGauge()).(metrics.Gauge)
	for {
		select {
		case <-resetTimer.C:
			log.Warn("No reply from server, assuming all streams gone")
			resetTimer.Stop()
			for stream := range prevStreams.Iter() {
				app.removeStream(stream.(string))
			}
			prevStreams = mapset.NewSet()

		case <-ticker:
			curStreams, err := nginxRtmp.GetActiveStreams()
			if err != nil {
				Catch(err)
				continue
			}
			resetTimer.Restart()

			numStreamsMetric.Update(int64(curStreams.Cardinality()))

			for stream := range curStreams.Difference(prevStreams).Iter() {
				app.addStream(stream.(string))
			}
			for stream := range prevStreams.Difference(curStreams).Iter() {
				app.removeStream(stream.(string))
			}

			// All streams gone? => Good time for a GC run
			if prevStreams.Cardinality() > 0 && curStreams.Cardinality() == 0 {
				debug.FreeOSMemory()
			}

			prevStreams = curStreams

		case <-app.quit.Barrier():
			resetTimer.Stop()
			for _, flows := range app.Flows {
				for _, flow := range flows {
					flow.Stop()
				}
			}
			return nil
		}
	}
}

func (app *App) addStream(stream string) {
	logged := false

	for flowName, flowConfig := range app.Config.Flows {
		if flowConfig.Regexp.MatchString(stream) {
			if !logged {
				logged = true
				log.WithFields(log.Fields{
					"stream": stream,
					"match":  true,
				}).Warn("New stream")
			}

			app.addFlow(flowName, stream, flowConfig.Source, flowConfig.Sinks)
		}
	}

	if !logged {
		log.WithFields(log.Fields{
			"stream": stream,
			"match":  false,
		}).Debug("New stream, ignoring")
	}
}

func (app *App) addFlow(name string, stream string, sourceTemplate CmdData, sinkTemplates map[string]CmdData) {
	vars := map[string]string{
		"{stream}": stream,
	}

	source := sourceTemplate.Replace(vars)
	sinks := make(map[string]CmdData, len(sinkTemplates))
	for sinkName, sinkTemplate := range sinkTemplates {
		sinks[sinkName] = sinkTemplate.Replace(vars)
	}

	flow := NewFlow(name, app.Config, source, sinks, log.WithFields(log.Fields{
		"name":   name,
		"stream": stream,
	}))
	flow.Start()

	if _, ok := app.Flows[stream]; !ok {
		app.Flows[stream] = make([]*Flow, 0, 1)
	}
	app.Flows[stream] = append(app.Flows[stream], flow)
}

func (app *App) removeStream(stream string) {
	log.WithFields(log.Fields{
		"stream": stream,
	}).Warn("Stream gone")

	flows, ok := app.Flows[stream]
	if !ok {
		return
	}

	for _, flow := range flows {
		flow.log.Info("Stopping flow")
		flow.Stop()
	}

	delete(app.Flows, stream)
}

func (app *App) handleSigint() {
	channel := make(chan os.Signal, 1)
	signal.Notify(channel, os.Interrupt)
	<-channel
	if app.Config.Debug {
		panic("Interrupted")
	} else {
		log.Info("interrupted, shutting down...")
		app.quit.Fall()
	}
}

func (app *App) handleSigusr1() {
	channel := make(chan os.Signal, 1)
	signal.Notify(channel, syscall.SIGUSR1)
	for range channel {
		stack := make([]byte, 524288)
		length := runtime.Stack(stack, true)
		name := fmt.Sprintf("/tmp/rtmptee.%d.stack", os.Getpid())

		f, err := os.Create(name)
		if err != nil {
			continue
		}
		f.Write(stack[:length])
		f.Close()

		log.Infof("Wrote stacktrace to %s", name)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015-present Oursky Ltd.
//
// 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.

// +build zmq

package zmq

import (
	"testing"
	"time"

	"github.com/zeromq/goczmq"

	. "github.com/smartystreets/goconvey/convey"
)

func workerSock(t *testing.T, id string, addr string) *goczmq.Sock {
	sock := goczmq.NewSock(goczmq.Dealer)
	sock.SetIdentity(id)
	if err := sock.Connect(addr); err != nil {
		t.Fatalf("Failed to create worker to addr = %s", addr)
	}
	return sock
}

func clientSock(t *testing.T, id string, addr string) *goczmq.Sock {
	sock := goczmq.NewSock(goczmq.Dealer)
	sock.SetIdentity(id)
	if err := sock.Connect(addr); err != nil {
		t.Fatalf("Failed to create client to addr = %s", addr)
	}
	return sock
}

func bytesArray(ss ...string) (bs [][]byte) {
	for _, s := range ss {
		bs = append(bs, []byte(s))
	}
	return
}

func TestWorker(t *testing.T) {
	Convey("Test workerQueue", t, func() {
		address1 := []byte("address1")
		address2 := []byte("address2")
		address3 := []byte("address3")

		Convey("Add and pick a worker", func() {
			q := newWorkerQueue()
			q.Add(newWorker(address1))
			So(q.Len(), ShouldEqual, 1)
			addr := q.Next()
			So(addr, ShouldResemble, address1)
			So(q.Len(), ShouldEqual, 0)
		})

		Convey("Add duplicated worker", func() {
			q := newWorkerQueue()
			q.Add(newWorker(address1))
			So(q.Len(), ShouldEqual, 1)
			q.Add(newWorker(address1))
			So(q.Len(), ShouldEqual, 1)

			q.Add(newWorker(address2))
			q.Add(newWorker(address3))
			So(q.Len(), ShouldEqual, 3)
		})

		Convey("Return the last Add worker first", func() {
			q := newWorkerQueue()
			q.Add(newWorker(address1))
			q.Add(newWorker(address2))
			q.Add(newWorker(address3))
			So(q.Len(), ShouldEqual, 3)

			So(q.Next(), ShouldResemble, address3)
			So(q.Next(), ShouldResemble, address2)
			So(q.Next(), ShouldResemble, address1)
		})

		Convey("Tick a non exist worker", func() {
			q := newWorkerQueue()
			q.Tick(newWorker(address1))
			So(q.Len(), ShouldEqual, 0)
			q.Add(newWorker(address2))
			So(q.Len(), ShouldEqual, 1)
			q.Tick(newWorker(address1))
			So(q.Len(), ShouldEqual, 1)
		})

		Convey("Pruge multiple expired workers", func() {
			q := newWorkerQueue()
			q.Add(newWorker(address1))
			q.Add(newWorker(address2))
			So(q.Len(), ShouldEqual, 2)

			// Wait the worker to time out
			time.Sleep((HeartbeatLiveness + 1) * HeartbeatInterval)
			q.Add(newWorker(address3))
			So(q.Len(), ShouldEqual, 3)
			q.Purge()
			So(q.Len(), ShouldEqual, 1)
		})
	})
}
func TestBrokerWorker(t *testing.T) {
	const (
		workerAddr = "inproc://plugin.test"
	)
	broker, err := NewBroker("", workerAddr)
	if err != nil {
		t.Fatalf("Failed to init broker: %v", err)
	}
	go broker.Run()
	go broker.Channler()

	Convey("Test Broker with worker", t, func() {
		Convey("receive Ready signal will register the worker", func() {
			w := workerSock(t, "ready", workerAddr)
			defer w.Destroy()
			w.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers

			So(broker.workers.Len(), ShouldEqual, 1)
			w.SendMessage(bytesArray(Shutdown))
		})

		Convey("receive multiple Ready signal will register all workers", func() {
			w1 := workerSock(t, "ready1", workerAddr)
			defer w1.Destroy()
			w1.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers
			w2 := workerSock(t, "ready2", workerAddr)
			defer w2.Destroy()
			w2.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers

			So(broker.workers.Len(), ShouldEqual, 2)
			w1.SendMessage(bytesArray(Shutdown))
			w2.SendMessage(bytesArray(Shutdown))
		})

		Convey("reveice Heartbeat without Reay will not register the worker", func() {
			w := workerSock(t, "heartbeat", workerAddr)
			defer w.Destroy()
			w.SendMessage(bytesArray(Heartbeat))
			// Wait the poller to get the message
			time.Sleep(HeartbeatInterval)
			So(broker.workers.Len(), ShouldEqual, 0)
		})

		Convey("broker RPC will timeout", func() {
			w := workerSock(t, "worker", workerAddr)
			w.SetRcvtimeo(heartbeatIntervalMS)
			defer w.Destroy()
			w.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers

			So(broker.workers.Len(), ShouldEqual, 1)

			reqChan := make(chan chan []byte)
			broker.RPC(reqChan, []byte(("from server")))
			respChan := <-reqChan
			msg := <-respChan
			So(msg, ShouldResemble, []byte{0})
		})

		Convey("broker RPC recive worker reply", func() {
			w := workerSock(t, "worker", workerAddr)
			w.SetRcvtimeo(heartbeatIntervalMS)
			defer w.Destroy()
			w.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers
			So(broker.workers.Len(), ShouldEqual, 1)

			reqChan := make(chan chan []byte)
			broker.RPC(reqChan, []byte(("from server")))
			respChan := <-reqChan

			msg, _ := w.RecvMessage()
			So(len(msg), ShouldEqual, 3)
			So(msg[2], ShouldResemble, []byte("from server"))
			msg[2] = []byte("from worker")
			w.SendMessage(msg)

			resp := <-respChan
			So(resp, ShouldResemble, []byte("from worker"))
		})

		Convey("send message from server to multple plugin", func() {
			w := workerSock(t, "worker", workerAddr)
			w.SetRcvtimeo(heartbeatIntervalMS)
			defer w.Destroy()
			w.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers
			w2 := workerSock(t, "worker2", workerAddr)
			w2.SetRcvtimeo(heartbeatIntervalMS)
			defer w2.Destroy()
			w2.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers

			So(broker.workers.Len(), ShouldEqual, 2)

			reqChan := make(chan chan []byte)
			broker.RPC(reqChan, []byte(("from server")))
			respChan := <-reqChan

			msg, _ := w2.RecvMessage()
			So(len(msg), ShouldEqual, 3)
			So(msg[2], ShouldResemble, []byte("from server"))
			msg[2] = []byte("from worker2")
			w2.SendMessage(msg)

			resp := <-respChan
			So(resp, ShouldResemble, []byte("from worker2"))
		})
	})
}
<commit_msg>Add multiple message and plugin test<commit_after>// Copyright 2015-present Oursky Ltd.
//
// 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.

// +build zmq

package zmq

import (
	"sync"
	"testing"
	"time"

	"github.com/zeromq/goczmq"

	. "github.com/smartystreets/goconvey/convey"
)

func workerSock(t *testing.T, id string, addr string) *goczmq.Sock {
	sock := goczmq.NewSock(goczmq.Dealer)
	sock.SetIdentity(id)
	if err := sock.Connect(addr); err != nil {
		t.Fatalf("Failed to create worker to addr = %s", addr)
	}
	return sock
}

func clientSock(t *testing.T, id string, addr string) *goczmq.Sock {
	sock := goczmq.NewSock(goczmq.Dealer)
	sock.SetIdentity(id)
	if err := sock.Connect(addr); err != nil {
		t.Fatalf("Failed to create client to addr = %s", addr)
	}
	return sock
}

func bytesArray(ss ...string) (bs [][]byte) {
	for _, s := range ss {
		bs = append(bs, []byte(s))
	}
	return
}

func TestWorker(t *testing.T) {
	Convey("Test workerQueue", t, func() {
		address1 := []byte("address1")
		address2 := []byte("address2")
		address3 := []byte("address3")

		Convey("Add and pick a worker", func() {
			q := newWorkerQueue()
			q.Add(newWorker(address1))
			So(q.Len(), ShouldEqual, 1)
			addr := q.Next()
			So(addr, ShouldResemble, address1)
			So(q.Len(), ShouldEqual, 0)
		})

		Convey("Add duplicated worker", func() {
			q := newWorkerQueue()
			q.Add(newWorker(address1))
			So(q.Len(), ShouldEqual, 1)
			q.Add(newWorker(address1))
			So(q.Len(), ShouldEqual, 1)

			q.Add(newWorker(address2))
			q.Add(newWorker(address3))
			So(q.Len(), ShouldEqual, 3)
		})

		Convey("Return the last Add worker first", func() {
			q := newWorkerQueue()
			q.Add(newWorker(address1))
			q.Add(newWorker(address2))
			q.Add(newWorker(address3))
			So(q.Len(), ShouldEqual, 3)

			So(q.Next(), ShouldResemble, address3)
			So(q.Next(), ShouldResemble, address2)
			So(q.Next(), ShouldResemble, address1)
		})

		Convey("Tick a non exist worker", func() {
			q := newWorkerQueue()
			q.Tick(newWorker(address1))
			So(q.Len(), ShouldEqual, 0)
			q.Add(newWorker(address2))
			So(q.Len(), ShouldEqual, 1)
			q.Tick(newWorker(address1))
			So(q.Len(), ShouldEqual, 1)
		})

		Convey("Pruge multiple expired workers", func() {
			q := newWorkerQueue()
			q.Add(newWorker(address1))
			q.Add(newWorker(address2))
			So(q.Len(), ShouldEqual, 2)

			// Wait the worker to time out
			time.Sleep((HeartbeatLiveness + 1) * HeartbeatInterval)
			q.Add(newWorker(address3))
			So(q.Len(), ShouldEqual, 3)
			q.Purge()
			So(q.Len(), ShouldEqual, 1)
		})
	})
}
func TestBrokerWorker(t *testing.T) {
	const (
		workerAddr = "inproc://plugin.test"
	)
	broker, err := NewBroker("", workerAddr)
	if err != nil {
		t.Fatalf("Failed to init broker: %v", err)
	}
	go broker.Run()
	go broker.Channler()

	Convey("Test Broker with worker", t, func() {
		Convey("receive Ready signal will register the worker", func() {
			w := workerSock(t, "ready", workerAddr)
			defer w.Destroy()
			w.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers

			So(broker.workers.Len(), ShouldEqual, 1)
			w.SendMessage(bytesArray(Shutdown))
		})

		Convey("receive multiple Ready signal will register all workers", func() {
			w1 := workerSock(t, "ready1", workerAddr)
			defer w1.Destroy()
			w1.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers
			w2 := workerSock(t, "ready2", workerAddr)
			defer w2.Destroy()
			w2.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers

			So(broker.workers.Len(), ShouldEqual, 2)
			w1.SendMessage(bytesArray(Shutdown))
			w2.SendMessage(bytesArray(Shutdown))
		})

		Convey("reveice Heartbeat without Reay will not register the worker", func() {
			w := workerSock(t, "heartbeat", workerAddr)
			defer w.Destroy()
			w.SendMessage(bytesArray(Heartbeat))
			// Wait the poller to get the message
			time.Sleep(HeartbeatInterval)
			So(broker.workers.Len(), ShouldEqual, 0)
		})

		Convey("broker RPC will timeout", func() {
			w := workerSock(t, "worker", workerAddr)
			w.SetRcvtimeo(heartbeatIntervalMS)
			defer w.Destroy()
			w.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers

			So(broker.workers.Len(), ShouldEqual, 1)

			reqChan := make(chan chan []byte)
			broker.RPC(reqChan, []byte(("from server")))
			respChan := <-reqChan
			msg := <-respChan
			So(msg, ShouldResemble, []byte{0})
		})

		Convey("broker RPC recive worker reply", func() {
			w := workerSock(t, "worker", workerAddr)
			w.SetRcvtimeo(heartbeatIntervalMS)
			defer w.Destroy()
			w.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers
			So(broker.workers.Len(), ShouldEqual, 1)

			reqChan := make(chan chan []byte)
			broker.RPC(reqChan, []byte(("from server")))
			respChan := <-reqChan

			msg, _ := w.RecvMessage()
			So(len(msg), ShouldEqual, 3)
			So(msg[2], ShouldResemble, []byte("from server"))
			msg[2] = []byte("from worker")
			w.SendMessage(msg)

			resp := <-respChan
			So(resp, ShouldResemble, []byte("from worker"))
		})

		Convey("send message from server to multple plugin", func() {
			w := workerSock(t, "worker", workerAddr)
			w.SetRcvtimeo(heartbeatIntervalMS)
			defer func() {
				w.SendMessage(bytesArray(Shutdown))
				w.Destroy()
			}()
			w.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers
			w2 := workerSock(t, "worker2", workerAddr)
			w2.SetRcvtimeo(heartbeatIntervalMS)
			defer func() {
				w2.SendMessage(bytesArray(Shutdown))
				w2.Destroy()
			}()
			w2.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers

			So(broker.workers.Len(), ShouldEqual, 2)

			reqChan := make(chan chan []byte)
			broker.RPC(reqChan, []byte(("from server")))
			respChan := <-reqChan

			msg, _ := w2.RecvMessage()
			So(len(msg), ShouldEqual, 3)
			So(msg[2], ShouldResemble, []byte("from server"))
			msg[2] = []byte("from worker2")
			w2.SendMessage(msg)

			resp := <-respChan
			So(resp, ShouldResemble, []byte("from worker2"))
		})

		Convey("send multiple message from server to multple plugin", func(c C) {
			w := workerSock(t, "mworker", workerAddr)
			w.SetRcvtimeo(heartbeatIntervalMS)
			defer func() {
				w.SendMessage(bytesArray(Shutdown))
				w.Destroy()
			}()
			w.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers
			w2 := workerSock(t, "mworker2", workerAddr)
			w2.SetRcvtimeo(heartbeatIntervalMS)
			defer func() {
				w2.SendMessage(bytesArray(Shutdown))
				w2.Destroy()
			}()
			w2.SendMessage(bytesArray(Ready))
			<-broker.freshWorkers

			So(broker.workers.Len(), ShouldEqual, 2)

			reqChan := make(chan chan []byte)
			broker.RPC(reqChan, []byte(("from server")))
			respChan := <-reqChan

			req2Chan := make(chan chan []byte)
			broker.RPC(req2Chan, []byte(("from server")))
			resp2Chan := <-req2Chan

			go func() {
				msg, _ := w.RecvMessage()
				c.So(len(msg), ShouldEqual, 3)
				c.So(msg[2], ShouldResemble, []byte("from server"))
				msg[2] = []byte("from worker1")
				w.SendMessage(msg)
			}()

			go func() {
				msg, _ := w2.RecvMessage()
				c.So(len(msg), ShouldEqual, 3)
				c.So(msg[2], ShouldResemble, []byte("from server"))
				msg[2] = []byte("from worker2")
				w2.SendMessage(msg)
			}()

			var wg sync.WaitGroup
			wg.Add(2)
			go func() {
				defer wg.Done()
				resp := <-respChan
				c.So(resp, ShouldResemble, []byte("from worker2"))
			}()
			go func() {
				defer wg.Done()
				resp := <-resp2Chan
				c.So(resp, ShouldResemble, []byte("from worker1"))
			}()
			wg.Wait()

		})
	})
}
<|endoftext|>


<commit_before>package link

import (
	"debug/elf"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sync"

	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/internal"
)

var (
	uprobeEventsPath = filepath.Join(tracefsPath, "uprobe_events")

	// rgxUprobeSymbol is used to strip invalid characters from the uprobe symbol
	// as they are not allowed to be used as the EVENT token in tracefs.
	rgxUprobeSymbol = regexp.MustCompile("[^a-zA-Z0-9]+")

	uprobeRetprobeBit = struct {
		once  sync.Once
		value uint64
		err   error
	}{}

	uprobeRefCtrOffsetPMUPath = "/sys/bus/event_source/devices/uprobe/format/ref_ctr_offset"
	// elixir.bootlin.com/linux/v5.15-rc7/source/kernel/events/core.c#L9799
	uprobeRefCtrOffsetShift = 32
	haveRefCtrOffsetPMU     = internal.FeatureTest("RefCtrOffsetPMU", "4.20", func() error {
		_, err := os.Stat(uprobeRefCtrOffsetPMUPath)
		if err != nil {
			return internal.ErrNotSupported
		}
		return nil
	})

	// ErrNoSymbol indicates that the given symbol was not found
	// in the ELF symbols table.
	ErrNoSymbol = errors.New("not found")
)

// Executable defines an executable program on the filesystem.
type Executable struct {
	// Path of the executable on the filesystem.
	path string
	// Parsed ELF symbols and dynamic symbols offsets.
	offsets map[string]uint64
}

// UprobeOptions defines additional parameters that will be used
// when loading Uprobes.
type UprobeOptions struct {
	// Symbol offset. Must be provided in case of external symbols (shared libs).
	// If set, overrides the offset eventually parsed from the executable.
	Offset uint64
	// Only set the uprobe on the given process ID. Useful when tracing
	// shared library calls or programs that have many running instances.
	PID int
	// Automatically manage SDT reference counts (semaphores).
	//
	// If this field is set, the Kernel will increment/decrement the
	// semaphore located in the process memory at the provided address on
	// probe attach/detach.
	//
	// See also:
	// sourceware.org/systemtap/wiki/UserSpaceProbeImplementation (Semaphore Handling)
	// github.com/torvalds/linux/commit/1cc33161a83d
	// github.com/torvalds/linux/commit/a6ca88b241d5
	RefCtrOffset uint64
}

// To open a new Executable, use:
//
//  OpenExecutable("/bin/bash")
//
// The returned value can then be used to open Uprobe(s).
func OpenExecutable(path string) (*Executable, error) {
	if path == "" {
		return nil, fmt.Errorf("path cannot be empty")
	}

	f, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open file '%s': %w", path, err)
	}
	defer f.Close()

	se, err := internal.NewSafeELFFile(f)
	if err != nil {
		return nil, fmt.Errorf("parse ELF file: %w", err)
	}

	if se.Type != elf.ET_EXEC && se.Type != elf.ET_DYN {
		// ELF is not an executable or a shared object.
		return nil, errors.New("the given file is not an executable or a shared object")
	}

	ex := Executable{
		path:    path,
		offsets: make(map[string]uint64),
	}

	if err := ex.load(se); err != nil {
		return nil, err
	}

	return &ex, nil
}

func (ex *Executable) load(f *internal.SafeELFFile) error {
	syms, err := f.Symbols()
	if err != nil && !errors.Is(err, elf.ErrNoSymbols) {
		return err
	}

	dynsyms, err := f.DynamicSymbols()
	if err != nil && !errors.Is(err, elf.ErrNoSymbols) {
		return err
	}

	syms = append(syms, dynsyms...)

	for _, s := range syms {
		if elf.ST_TYPE(s.Info) != elf.STT_FUNC {
			// Symbol not associated with a function or other executable code.
			continue
		}

		off := s.Value

		// Loop over ELF segments.
		for _, prog := range f.Progs {
			// Skip uninteresting segments.
			if prog.Type != elf.PT_LOAD || (prog.Flags&elf.PF_X) == 0 {
				continue
			}

			if prog.Vaddr <= s.Value && s.Value < (prog.Vaddr+prog.Memsz) {
				// If the symbol value is contained in the segment, calculate
				// the symbol offset.
				//
				// fn symbol offset = fn symbol VA - .text VA + .text offset
				//
				// stackoverflow.com/a/40249502
				off = s.Value - prog.Vaddr + prog.Off
				break
			}
		}

		ex.offsets[s.Name] = off
	}

	return nil
}

func (ex *Executable) offset(symbol string) (uint64, error) {
	if off, ok := ex.offsets[symbol]; ok {
		// Symbols with location 0 from section undef are shared library calls and
		// are relocated before the binary is executed. Dynamic linking is not
		// implemented by the library, so mark this as unsupported for now.
		//
		// Since only offset values are stored and not elf.Symbol, if the value is 0,
		// assume it's an external symbol.
		if off == 0 {
			return 0, fmt.Errorf("cannot resolve %s library call '%s', "+
				"consider providing the offset via options: %w", ex.path, symbol, ErrNotSupported)
		}
		return off, nil
	}
	return 0, fmt.Errorf("symbol %s: %w", symbol, ErrNoSymbol)
}

// Uprobe attaches the given eBPF program to a perf event that fires when the
// given symbol starts executing in the given Executable.
// For example, /bin/bash::main():
//
//  ex, _ = OpenExecutable("/bin/bash")
//  ex.Uprobe("main", prog, nil)
//
// When using symbols which belongs to shared libraries,
// an offset must be provided via options:
//
//  up, err := ex.Uprobe("main", prog, &UprobeOptions{Offset: 0x123})
//
// Losing the reference to the resulting Link (up) will close the Uprobe
// and prevent further execution of prog. The Link must be Closed during
// program shutdown to avoid leaking system resources.
//
// Functions provided by shared libraries can currently not be traced and
// will result in an ErrNotSupported.
func (ex *Executable) Uprobe(symbol string, prog *ebpf.Program, opts *UprobeOptions) (Link, error) {
	u, err := ex.uprobe(symbol, prog, opts, false)
	if err != nil {
		return nil, err
	}

	err = u.attach(prog)
	if err != nil {
		u.Close()
		return nil, err
	}

	return u, nil
}

// Uretprobe attaches the given eBPF program to a perf event that fires right
// before the given symbol exits. For example, /bin/bash::main():
//
//  ex, _ = OpenExecutable("/bin/bash")
//  ex.Uretprobe("main", prog, nil)
//
// When using symbols which belongs to shared libraries,
// an offset must be provided via options:
//
//  up, err := ex.Uretprobe("main", prog, &UprobeOptions{Offset: 0x123})
//
// Losing the reference to the resulting Link (up) will close the Uprobe
// and prevent further execution of prog. The Link must be Closed during
// program shutdown to avoid leaking system resources.
//
// Functions provided by shared libraries can currently not be traced and
// will result in an ErrNotSupported.
func (ex *Executable) Uretprobe(symbol string, prog *ebpf.Program, opts *UprobeOptions) (Link, error) {
	u, err := ex.uprobe(symbol, prog, opts, true)
	if err != nil {
		return nil, err
	}

	err = u.attach(prog)
	if err != nil {
		u.Close()
		return nil, err
	}

	return u, nil
}

// uprobe opens a perf event for the given binary/symbol and attaches prog to it.
// If ret is true, create a uretprobe.
func (ex *Executable) uprobe(symbol string, prog *ebpf.Program, opts *UprobeOptions, ret bool) (*perfEvent, error) {
	if prog == nil {
		return nil, fmt.Errorf("prog cannot be nil: %w", errInvalidInput)
	}
	if prog.Type() != ebpf.Kprobe {
		return nil, fmt.Errorf("eBPF program type %s is not Kprobe: %w", prog.Type(), errInvalidInput)
	}
	if opts == nil {
		opts = &UprobeOptions{}
	}

	offset := opts.Offset
	if offset == 0 {
		off, err := ex.offset(symbol)
		if err != nil {
			return nil, err
		}
		offset = off
	}

	pid := opts.PID
	if pid == 0 {
		pid = perfAllThreads
	}

	if opts.RefCtrOffset != 0 {
		if err := haveRefCtrOffsetPMU(); err != nil {
			return nil, fmt.Errorf("uprobe ref_ctr_offset: %w", err)
		}
	}

	args := probeArgs{
		symbol:       symbol,
		path:         ex.path,
		offset:       offset,
		pid:          pid,
		refCtrOffset: opts.RefCtrOffset,
		ret:          ret,
	}

	// Use uprobe PMU if the kernel has it available.
	tp, err := pmuUprobe(args)
	if err == nil {
		return tp, nil
	}
	if err != nil && !errors.Is(err, ErrNotSupported) {
		return nil, fmt.Errorf("creating perf_uprobe PMU: %w", err)
	}

	// Use tracefs if uprobe PMU is missing.
	args.symbol = uprobeSanitizedSymbol(symbol)
	tp, err = tracefsUprobe(args)
	if err != nil {
		return nil, fmt.Errorf("creating trace event '%s:%s' in tracefs: %w", ex.path, symbol, err)
	}

	return tp, nil
}

// pmuUprobe opens a perf event based on the uprobe PMU.
func pmuUprobe(args probeArgs) (*perfEvent, error) {
	return pmuProbe(uprobeType, args)
}

// tracefsUprobe creates a Uprobe tracefs entry.
func tracefsUprobe(args probeArgs) (*perfEvent, error) {
	return tracefsProbe(uprobeType, args)
}

// uprobeSanitizedSymbol replaces every invalid characted for the tracefs api with an underscore.
func uprobeSanitizedSymbol(symbol string) string {
	return rgxUprobeSymbol.ReplaceAllString(symbol, "_")
}

// uprobeToken creates the PATH:OFFSET(REF_CTR_OFFSET) token for the tracefs api.
func uprobeToken(args probeArgs) string {
	po := fmt.Sprintf("%s:%#x", args.path, args.offset)

	if args.refCtrOffset != 0 {
		// This is not documented in Documentation/trace/uprobetracer.txt.
		// elixir.bootlin.com/linux/v5.15-rc7/source/kernel/trace/trace.c#L5564
		po += fmt.Sprintf("(%#x)", args.refCtrOffset)
	}

	return po
}

func uretprobeBit() (uint64, error) {
	uprobeRetprobeBit.once.Do(func() {
		uprobeRetprobeBit.value, uprobeRetprobeBit.err = determineRetprobeBit(uprobeType)
	})
	return uprobeRetprobeBit.value, uprobeRetprobeBit.err
}
<commit_msg>link/uprobe: allow specifying bpf cookie via options<commit_after>package link

import (
	"debug/elf"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sync"

	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/internal"
)

var (
	uprobeEventsPath = filepath.Join(tracefsPath, "uprobe_events")

	// rgxUprobeSymbol is used to strip invalid characters from the uprobe symbol
	// as they are not allowed to be used as the EVENT token in tracefs.
	rgxUprobeSymbol = regexp.MustCompile("[^a-zA-Z0-9]+")

	uprobeRetprobeBit = struct {
		once  sync.Once
		value uint64
		err   error
	}{}

	uprobeRefCtrOffsetPMUPath = "/sys/bus/event_source/devices/uprobe/format/ref_ctr_offset"
	// elixir.bootlin.com/linux/v5.15-rc7/source/kernel/events/core.c#L9799
	uprobeRefCtrOffsetShift = 32
	haveRefCtrOffsetPMU     = internal.FeatureTest("RefCtrOffsetPMU", "4.20", func() error {
		_, err := os.Stat(uprobeRefCtrOffsetPMUPath)
		if err != nil {
			return internal.ErrNotSupported
		}
		return nil
	})

	// ErrNoSymbol indicates that the given symbol was not found
	// in the ELF symbols table.
	ErrNoSymbol = errors.New("not found")
)

// Executable defines an executable program on the filesystem.
type Executable struct {
	// Path of the executable on the filesystem.
	path string
	// Parsed ELF symbols and dynamic symbols offsets.
	offsets map[string]uint64
}

// UprobeOptions defines additional parameters that will be used
// when loading Uprobes.
type UprobeOptions struct {
	// Symbol offset. Must be provided in case of external symbols (shared libs).
	// If set, overrides the offset eventually parsed from the executable.
	Offset uint64
	// Only set the uprobe on the given process ID. Useful when tracing
	// shared library calls or programs that have many running instances.
	PID int
	// Automatically manage SDT reference counts (semaphores).
	//
	// If this field is set, the Kernel will increment/decrement the
	// semaphore located in the process memory at the provided address on
	// probe attach/detach.
	//
	// See also:
	// sourceware.org/systemtap/wiki/UserSpaceProbeImplementation (Semaphore Handling)
	// github.com/torvalds/linux/commit/1cc33161a83d
	// github.com/torvalds/linux/commit/a6ca88b241d5
	RefCtrOffset uint64
	// Arbitrary value that can be fetched from an eBPF program
	// via `bpf_get_attach_cookie()`.
	//
	// Needs kernel 5.15+.
	Cookie uint64
}

// To open a new Executable, use:
//
//  OpenExecutable("/bin/bash")
//
// The returned value can then be used to open Uprobe(s).
func OpenExecutable(path string) (*Executable, error) {
	if path == "" {
		return nil, fmt.Errorf("path cannot be empty")
	}

	f, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open file '%s': %w", path, err)
	}
	defer f.Close()

	se, err := internal.NewSafeELFFile(f)
	if err != nil {
		return nil, fmt.Errorf("parse ELF file: %w", err)
	}

	if se.Type != elf.ET_EXEC && se.Type != elf.ET_DYN {
		// ELF is not an executable or a shared object.
		return nil, errors.New("the given file is not an executable or a shared object")
	}

	ex := Executable{
		path:    path,
		offsets: make(map[string]uint64),
	}

	if err := ex.load(se); err != nil {
		return nil, err
	}

	return &ex, nil
}

func (ex *Executable) load(f *internal.SafeELFFile) error {
	syms, err := f.Symbols()
	if err != nil && !errors.Is(err, elf.ErrNoSymbols) {
		return err
	}

	dynsyms, err := f.DynamicSymbols()
	if err != nil && !errors.Is(err, elf.ErrNoSymbols) {
		return err
	}

	syms = append(syms, dynsyms...)

	for _, s := range syms {
		if elf.ST_TYPE(s.Info) != elf.STT_FUNC {
			// Symbol not associated with a function or other executable code.
			continue
		}

		off := s.Value

		// Loop over ELF segments.
		for _, prog := range f.Progs {
			// Skip uninteresting segments.
			if prog.Type != elf.PT_LOAD || (prog.Flags&elf.PF_X) == 0 {
				continue
			}

			if prog.Vaddr <= s.Value && s.Value < (prog.Vaddr+prog.Memsz) {
				// If the symbol value is contained in the segment, calculate
				// the symbol offset.
				//
				// fn symbol offset = fn symbol VA - .text VA + .text offset
				//
				// stackoverflow.com/a/40249502
				off = s.Value - prog.Vaddr + prog.Off
				break
			}
		}

		ex.offsets[s.Name] = off
	}

	return nil
}

func (ex *Executable) offset(symbol string) (uint64, error) {
	if off, ok := ex.offsets[symbol]; ok {
		// Symbols with location 0 from section undef are shared library calls and
		// are relocated before the binary is executed. Dynamic linking is not
		// implemented by the library, so mark this as unsupported for now.
		//
		// Since only offset values are stored and not elf.Symbol, if the value is 0,
		// assume it's an external symbol.
		if off == 0 {
			return 0, fmt.Errorf("cannot resolve %s library call '%s', "+
				"consider providing the offset via options: %w", ex.path, symbol, ErrNotSupported)
		}
		return off, nil
	}
	return 0, fmt.Errorf("symbol %s: %w", symbol, ErrNoSymbol)
}

// Uprobe attaches the given eBPF program to a perf event that fires when the
// given symbol starts executing in the given Executable.
// For example, /bin/bash::main():
//
//  ex, _ = OpenExecutable("/bin/bash")
//  ex.Uprobe("main", prog, nil)
//
// When using symbols which belongs to shared libraries,
// an offset must be provided via options:
//
//  up, err := ex.Uprobe("main", prog, &UprobeOptions{Offset: 0x123})
//
// Losing the reference to the resulting Link (up) will close the Uprobe
// and prevent further execution of prog. The Link must be Closed during
// program shutdown to avoid leaking system resources.
//
// Functions provided by shared libraries can currently not be traced and
// will result in an ErrNotSupported.
func (ex *Executable) Uprobe(symbol string, prog *ebpf.Program, opts *UprobeOptions) (Link, error) {
	u, err := ex.uprobe(symbol, prog, opts, false)
	if err != nil {
		return nil, err
	}

	err = u.attach(prog)
	if err != nil {
		u.Close()
		return nil, err
	}

	return u, nil
}

// Uretprobe attaches the given eBPF program to a perf event that fires right
// before the given symbol exits. For example, /bin/bash::main():
//
//  ex, _ = OpenExecutable("/bin/bash")
//  ex.Uretprobe("main", prog, nil)
//
// When using symbols which belongs to shared libraries,
// an offset must be provided via options:
//
//  up, err := ex.Uretprobe("main", prog, &UprobeOptions{Offset: 0x123})
//
// Losing the reference to the resulting Link (up) will close the Uprobe
// and prevent further execution of prog. The Link must be Closed during
// program shutdown to avoid leaking system resources.
//
// Functions provided by shared libraries can currently not be traced and
// will result in an ErrNotSupported.
func (ex *Executable) Uretprobe(symbol string, prog *ebpf.Program, opts *UprobeOptions) (Link, error) {
	u, err := ex.uprobe(symbol, prog, opts, true)
	if err != nil {
		return nil, err
	}

	err = u.attach(prog)
	if err != nil {
		u.Close()
		return nil, err
	}

	return u, nil
}

// uprobe opens a perf event for the given binary/symbol and attaches prog to it.
// If ret is true, create a uretprobe.
func (ex *Executable) uprobe(symbol string, prog *ebpf.Program, opts *UprobeOptions, ret bool) (*perfEvent, error) {
	if prog == nil {
		return nil, fmt.Errorf("prog cannot be nil: %w", errInvalidInput)
	}
	if prog.Type() != ebpf.Kprobe {
		return nil, fmt.Errorf("eBPF program type %s is not Kprobe: %w", prog.Type(), errInvalidInput)
	}
	if opts == nil {
		opts = &UprobeOptions{}
	}

	offset := opts.Offset
	if offset == 0 {
		off, err := ex.offset(symbol)
		if err != nil {
			return nil, err
		}
		offset = off
	}

	pid := opts.PID
	if pid == 0 {
		pid = perfAllThreads
	}

	if opts.RefCtrOffset != 0 {
		if err := haveRefCtrOffsetPMU(); err != nil {
			return nil, fmt.Errorf("uprobe ref_ctr_offset: %w", err)
		}
	}

	args := probeArgs{
		symbol:       symbol,
		path:         ex.path,
		offset:       offset,
		pid:          pid,
		refCtrOffset: opts.RefCtrOffset,
		ret:          ret,
		cookie:       opts.Cookie,
	}

	// Use uprobe PMU if the kernel has it available.
	tp, err := pmuUprobe(args)
	if err == nil {
		return tp, nil
	}
	if err != nil && !errors.Is(err, ErrNotSupported) {
		return nil, fmt.Errorf("creating perf_uprobe PMU: %w", err)
	}

	// Use tracefs if uprobe PMU is missing.
	args.symbol = uprobeSanitizedSymbol(symbol)
	tp, err = tracefsUprobe(args)
	if err != nil {
		return nil, fmt.Errorf("creating trace event '%s:%s' in tracefs: %w", ex.path, symbol, err)
	}

	return tp, nil
}

// pmuUprobe opens a perf event based on the uprobe PMU.
func pmuUprobe(args probeArgs) (*perfEvent, error) {
	return pmuProbe(uprobeType, args)
}

// tracefsUprobe creates a Uprobe tracefs entry.
func tracefsUprobe(args probeArgs) (*perfEvent, error) {
	return tracefsProbe(uprobeType, args)
}

// uprobeSanitizedSymbol replaces every invalid characted for the tracefs api with an underscore.
func uprobeSanitizedSymbol(symbol string) string {
	return rgxUprobeSymbol.ReplaceAllString(symbol, "_")
}

// uprobeToken creates the PATH:OFFSET(REF_CTR_OFFSET) token for the tracefs api.
func uprobeToken(args probeArgs) string {
	po := fmt.Sprintf("%s:%#x", args.path, args.offset)

	if args.refCtrOffset != 0 {
		// This is not documented in Documentation/trace/uprobetracer.txt.
		// elixir.bootlin.com/linux/v5.15-rc7/source/kernel/trace/trace.c#L5564
		po += fmt.Sprintf("(%#x)", args.refCtrOffset)
	}

	return po
}

func uretprobeBit() (uint64, error) {
	uprobeRetprobeBit.once.Do(func() {
		uprobeRetprobeBit.value, uprobeRetprobeBit.err = determineRetprobeBit(uprobeType)
	})
	return uprobeRetprobeBit.value, uprobeRetprobeBit.err
}
<|endoftext|>


<commit_before>package matasano

import (
	"encoding/hex"
	"encoding/base64"
	"bytes"
	"errors"
	"crypto/aes"
)

func HexToBase64(input string) (output string, err error) {
	bytes, decodeError := hex.DecodeString(input)
	if decodeError != nil {
		return "", decodeError
	}
	return base64.StdEncoding.EncodeToString(bytes), nil
}

func FixedXor(input string, mask string) (output string, err error) {
	var decodeErr error
	var inputBytes, maskBytes []byte
	inputBytes, decodeErr = hex.DecodeString(input)
	if decodeErr != nil {
		return "", decodeErr
	}
	maskBytes, decodeErr = hex.DecodeString(mask)
	if decodeErr != nil {
		return "", decodeErr
	}

	resultBytes := bytes.Map(makeRepeatingXorClojure(maskBytes), inputBytes)
	return hex.EncodeToString(resultBytes), nil
}

func FixedXorWithSingleByteMask(input []byte, mask byte) []byte {
	return bytes.Map(makeRepeatingXorClojure([]byte{mask}), input)
}

func RepeatingXor(input []byte, mask []byte) []byte {
	return bytes.Map(makeRepeatingXorClojure(mask), input)
}

func makeRepeatingXorClojure(mask [] byte) func(r rune) rune {
	i := 0
	maskLength := len(mask)
	return func(r rune) rune {
		result := rune(byte(r) ^ mask[i])
		i = (i+1)%maskLength
		return result
	}
}

func HemingDistance(a []byte, b []byte) (uint, error) {
	if len(a) != len(b) {
		return 0, errors.New("Both arrays should be the same length.")
	}
	diffBits := uint(0)
	for i := 0 ; i < len(a) ; i++ {
		diff := a[i] ^ b[i]
		for j := uint(0) ; j < 8 ; j++ {
			diffBits += uint((diff >> j) & 1)
		}
	}
	return diffBits, nil
}

func DecryptAES128ECB(encrypted []byte, key []byte) ([]byte, error) {
	cipher, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	decrypted := new(bytes.Buffer)
	input := bytes.NewBuffer(encrypted)
	bufIn := make([]byte, cipher.BlockSize())
	bufOut := make([]byte, cipher.BlockSize())
	for {
		_, err := input.Read(bufIn)
		if err != nil {
			// EOF
			break
		}
		cipher.Decrypt(bufOut, bufIn)
		decrypted.Write(bufOut)
	}
	return decrypted.Bytes(), nil
}

/*
	Counts the amount of matching blocks in a byte array given a certain block size (in bytes).
	Every match will be single-counted, so if block a matches block b, then block b matching block a will not be
	counted.
 */
func CountMatchingBlocks(data []byte, blockSize int) int {
	blocks := len(data) / blockSize
	matchingBlocks := 0
	for i := 0 ; i < blocks ; i++ {
		offset := i * blockSize
		this := data[offset:offset+blockSize]
		for j := i + 1 ; j < blocks ; j++ {
			offset := j * blockSize
			that := data[offset:offset+blockSize]
			if bytes.Equal(this, that) {
				matchingBlocks++
			}
		}
	}
	return matchingBlocks
}

func PadBufferUsingPKCS7(input []byte, blocksize int) []byte {
	if len(input) >= blocksize {
		return input
	}
	padding := byte(blocksize - len(input))
	var output bytes.Buffer
	output.Write(input)
	for i := byte(0) ; i < padding ; i++ {
		output.WriteByte(padding)
	}
	return output.Bytes()
}
<commit_msg>Padding now also works for buffers of mulitples of blocksize.<commit_after>package matasano

import (
	"encoding/hex"
	"encoding/base64"
	"bytes"
	"errors"
	"crypto/aes"
)

func HexToBase64(input string) (output string, err error) {
	bytes, decodeError := hex.DecodeString(input)
	if decodeError != nil {
		return "", decodeError
	}
	return base64.StdEncoding.EncodeToString(bytes), nil
}

func FixedXor(input string, mask string) (output string, err error) {
	var decodeErr error
	var inputBytes, maskBytes []byte
	inputBytes, decodeErr = hex.DecodeString(input)
	if decodeErr != nil {
		return "", decodeErr
	}
	maskBytes, decodeErr = hex.DecodeString(mask)
	if decodeErr != nil {
		return "", decodeErr
	}

	resultBytes := bytes.Map(makeRepeatingXorClojure(maskBytes), inputBytes)
	return hex.EncodeToString(resultBytes), nil
}

func FixedXorWithSingleByteMask(input []byte, mask byte) []byte {
	return bytes.Map(makeRepeatingXorClojure([]byte{mask}), input)
}

func RepeatingXor(input []byte, mask []byte) []byte {
	return bytes.Map(makeRepeatingXorClojure(mask), input)
}

func makeRepeatingXorClojure(mask [] byte) func(r rune) rune {
	i := 0
	maskLength := len(mask)
	return func(r rune) rune {
		result := rune(byte(r) ^ mask[i])
		i = (i+1)%maskLength
		return result
	}
}

func HemingDistance(a []byte, b []byte) (uint, error) {
	if len(a) != len(b) {
		return 0, errors.New("Both arrays should be the same length.")
	}
	diffBits := uint(0)
	for i := 0 ; i < len(a) ; i++ {
		diff := a[i] ^ b[i]
		for j := uint(0) ; j < 8 ; j++ {
			diffBits += uint((diff >> j) & 1)
		}
	}
	return diffBits, nil
}

func DecryptAES128ECB(encrypted []byte, key []byte) ([]byte, error) {
	cipher, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	decrypted := new(bytes.Buffer)
	input := bytes.NewBuffer(encrypted)
	bufIn := make([]byte, cipher.BlockSize())
	bufOut := make([]byte, cipher.BlockSize())
	for {
		_, err := input.Read(bufIn)
		if err != nil {
			// EOF
			break
		}
		cipher.Decrypt(bufOut, bufIn)
		decrypted.Write(bufOut)
	}
	return decrypted.Bytes(), nil
}

/*
	Counts the amount of matching blocks in a byte array given a certain block size (in bytes).
	Every match will be single-counted, so if block a matches block b, then block b matching block a will not be
	counted.
 */
func CountMatchingBlocks(data []byte, blockSize int) int {
	blocks := len(data) / blockSize
	matchingBlocks := 0
	for i := 0 ; i < blocks ; i++ {
		offset := i * blockSize
		this := data[offset:offset+blockSize]
		for j := i + 1 ; j < blocks ; j++ {
			offset := j * blockSize
			that := data[offset:offset+blockSize]
			if bytes.Equal(this, that) {
				matchingBlocks++
			}
		}
	}
	return matchingBlocks
}

func PadBufferUsingPKCS7(input []byte, blocksize int) []byte {
	remainder := len(input) % blocksize
	if remainder == 0 {
		return input
	}
	padding := byte(blocksize - remainder)
	var output bytes.Buffer
	output.Write(input)
	for i := byte(0) ; i < padding ; i++ {
		output.WriteByte(padding)
	}
	return output.Bytes()
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"io"
	"log"
	"net/http"

	"github.com/terranodo/tegola/cache"
)

//	TileCacheHandler implements a request cache for tiles on requests when the URLs
//	have a /:z/:x/:y scheme suffix (i.e. /osm/1/3/4.pbf)
func TileCacheHandler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error

		//	check if a cache backend exists
		if Cache == nil {
			//	nope. move on
			next.ServeHTTP(w, r)
			return
		}

		//	parse our URI into a cache key structure
		key, err := cache.ParseKey(r.URL.Path)
		if err != nil {
			log.Println("cache middleware: ParseKey err: %v", err)
			next.ServeHTTP(w, r)
			return
		}

		//	use the URL path as the key
		cachedTile, hit, err := Cache.Get(key)
		if err != nil {
			log.Printf("cache middleware: error reading from cache: %v", err)
			next.ServeHTTP(w, r)
			return
		}
		//	cache miss
		if !hit {
			//	buffer which will hold a copy of the response for writing to the cache
			var buff bytes.Buffer

			//	ovewrite our current responseWriter with a tileCacheResponseWriter
			w = newTileCacheResponseWriter(w, &buff)

			next.ServeHTTP(w, r)

			if err := Cache.Set(key, buff.Bytes()); err != nil {
				log.Println("cache response writer err: %v", err)
			}
			return
		}

		//	TODO: how configurable do we want the CORS policy to be?
		//	set CORS header
		w.Header().Add("Access-Control-Allow-Origin", "*")

		//	mimetype for protocol buffers
		w.Header().Add("Content-Type", "application/x-protobuf")

		//	communicate the cache is being used
		w.Header().Add("Tegola-Cache", "HIT")

		w.Write(cachedTile)
		return
	})
}

func newTileCacheResponseWriter(resp http.ResponseWriter, w io.Writer) http.ResponseWriter {
	return &tileCacheResponseWriter{
		resp:  resp,
		multi: io.MultiWriter(w, resp),
	}
}

//	tileCacheResponsWriter wraps http.ResponseWriter (https://golang.org/pkg/net/http/#ResponseWriter)
//	to additionally write the response to a cache when there is a cache MISS
type tileCacheResponseWriter struct {
	resp  http.ResponseWriter
	multi io.Writer
}

func (w *tileCacheResponseWriter) Header() http.Header {
	//	communicate the tegola cache is being used
	w.resp.Header().Set("Tegola-Cache", "MISS")

	return w.resp.Header()
}

func (w *tileCacheResponseWriter) Write(b []byte) (int, error) {
	//	write to our multi writer
	return w.multi.Write(b)
}

func (w *tileCacheResponseWriter) WriteHeader(i int) {
	w.resp.WriteHeader(i)
}
<commit_msg>cehck for canceled requests before writing to cache in cache middleware. #96<commit_after>package server

import (
	"bytes"
	"io"
	"log"
	"net/http"

	"github.com/terranodo/tegola/cache"
)

//	TileCacheHandler implements a request cache for tiles on requests when the URLs
//	have a /:z/:x/:y scheme suffix (i.e. /osm/1/3/4.pbf)
func TileCacheHandler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error

		//	check if a cache backend exists
		if Cache == nil {
			//	nope. move on
			next.ServeHTTP(w, r)
			return
		}

		//	parse our URI into a cache key structure
		key, err := cache.ParseKey(r.URL.Path)
		if err != nil {
			log.Println("cache middleware: ParseKey err: %v", err)
			next.ServeHTTP(w, r)
			return
		}

		//	use the URL path as the key
		cachedTile, hit, err := Cache.Get(key)
		if err != nil {
			log.Printf("cache middleware: error reading from cache: %v", err)
			next.ServeHTTP(w, r)
			return
		}
		//	cache miss
		if !hit {
			//	buffer which will hold a copy of the response for writing to the cache
			var buff bytes.Buffer

			//	ovewrite our current responseWriter with a tileCacheResponseWriter
			w = newTileCacheResponseWriter(w, &buff)

			next.ServeHTTP(w, r)

			//	check if our request context has been canceled
			if r.Context().Err() != nil {
				return
			}

			if err := Cache.Set(key, buff.Bytes()); err != nil {
				log.Println("cache response writer err: %v", err)
			}
			return
		}

		//	TODO: how configurable do we want the CORS policy to be?
		//	set CORS header
		w.Header().Add("Access-Control-Allow-Origin", "*")

		//	mimetype for protocol buffers
		w.Header().Add("Content-Type", "application/x-protobuf")

		//	communicate the cache is being used
		w.Header().Add("Tegola-Cache", "HIT")

		w.Write(cachedTile)
		return
	})
}

func newTileCacheResponseWriter(resp http.ResponseWriter, w io.Writer) http.ResponseWriter {
	return &tileCacheResponseWriter{
		resp:  resp,
		multi: io.MultiWriter(w, resp),
	}
}

//	tileCacheResponsWriter wraps http.ResponseWriter (https://golang.org/pkg/net/http/#ResponseWriter)
//	to additionally write the response to a cache when there is a cache MISS
type tileCacheResponseWriter struct {
	resp  http.ResponseWriter
	multi io.Writer
}

func (w *tileCacheResponseWriter) Header() http.Header {
	//	communicate the tegola cache is being used
	w.resp.Header().Set("Tegola-Cache", "MISS")

	return w.resp.Header()
}

func (w *tileCacheResponseWriter) Write(b []byte) (int, error) {
	//	write to our multi writer
	return w.multi.Write(b)
}

func (w *tileCacheResponseWriter) WriteHeader(i int) {
	w.resp.WriteHeader(i)
}
<|endoftext|>


<commit_before>package mysqldb

import (
	"fmt"
	l "libs/log"
	. "libs/parser"
	"strings"
)

var (

	//ID表
	idtbl = TblField{
		FieldList: []Field{
			Field{"id", "BIGINT(20) UNSIGNED", false, true},
			Field{"stub", "VARCHAR(10)", false, false},
		},
		IndexInfo: map[string][]string{
			"PRIMARY":             []string{"id"},
			"UNIQUE INDEX `stub`": []string{"stub"},
		},
	}
	//角色表
	roletbl = TblField{
		FieldList: []Field{
			Field{"uid", "BIGINT(20) UNSIGNED", false, false},
			Field{"account", "CHAR(100)", false, false},
			Field{"rolename", "CHAR(64)", false, false},
			Field{"createtime", "DATETIME", false, false},
			Field{"lastlogintime", "DATETIME", false, false},
			Field{"locktime", "DATETIME", false, false},
			Field{"roleindex", "TINYINT(1) UNSIGNED", false, false},
			Field{"roleinfo", "TEXT", true, false},
			Field{"entity", "CHAR(100)", false, false},
			Field{"deleted", "TINYINT(1) UNSIGNED", true, false},
			Field{"locked", "TINYINT(1) UNSIGNED", true, false},
			Field{"status", "TINYINT(1) UNSIGNED", false, false},
			Field{"serverid", "CHAR(64)", true, false},
			Field{"scene", "CHAR(64)", true, false},
			Field{"scene_x", "FLOAT", true, false},
			Field{"scene_y", "FLOAT", true, false},
			Field{"scene_z", "FLOAT", true, false},
			Field{"scene_dir", "FLOAT", true, false},
			Field{"landtimes", "INT(10) UNSIGNED", false, false},
		},
		IndexInfo: map[string][]string{
			"PRIMARY":      []string{"uid"},
			"INDEX `role`": []string{"account", "rolename"},
		},
	}
)

type Sync struct {
	conn       *MySql
	db         string
	datasrouce string
}

func (s *Sync) syncTable(tbl_name string, tblinfo TblField, incNum int) error {
	if oldfield, empty, _ := s.conn.QueryTblInfo(s.db, tbl_name); empty {
		err := s.conn.CreateTable(tbl_name, tblinfo, incNum)
		if err != nil {
			panic(err)
		}
		l.TraceInfo("sync", "create table ", tbl_name, " ok")
	} else {
		alert := make([]Field, 0, 10)
		add := make([]Field, 0, 10)
		for _, newf := range tblinfo.FieldList {
			find := false
			for _, oldf := range oldfield.FieldList {
				if strings.ToLower(newf.Name) == strings.ToLower(oldf.Name) {
					if strings.ToLower(oldf.Type) != strings.ToLower(newf.Type) {
						alert = append(alert, newf)
					}
					find = true
					break
				}
			}
			if !find {
				add = append(add, newf)
			}
		}

		del := make([]Field, 0, 10)
		for _, oldf := range oldfield.FieldList {
			find := false
			for _, newf := range tblinfo.FieldList {
				if strings.ToLower(newf.Name) == strings.ToLower(oldf.Name) {
					find = true
					break
				}
			}
			if !find {
				del = append(del, oldf)
			}
		}

		for _, f := range alert {
			var nulltype string
			if !f.IsNull {
				nulltype = "NOT NULL"
			} else {
				nulltype = "DEFAULT NULL"
			}
			autoinc := ""
			if f.AutoInc {
				autoinc = " AUTO_INCREMENT"
			}

			sql := fmt.Sprintf("ALTER TABLE `%s` CHANGE `%s` `%s` %s %s %s;", tbl_name, f.Name, f.Name, f.Type, nulltype, autoinc)
			_, err := s.conn.Exec(sql)
			if err != nil {
				panic(err)
			}

			l.TraceInfo("sync", sql, " ok")
		}

		for _, f := range add {
			var nulltype string
			if !f.IsNull {
				nulltype = "NOT NULL"
			} else {
				nulltype = "DEFAULT NULL"
			}
			autoinc := ""
			if f.AutoInc {
				autoinc = " AUTO_INCREMENT"
			}

			/*if len(del) > 0 {
				sql := fmt.Sprintf("ALTER TABLE `%s` CHANGE `%s` `%s` %s %s %s;", tbl_name, del[0].Name, f.Name, f.Type, nulltype, autoinc)
				_, err := s.conn.Exec(sql)
				if err != nil {
					panic(err)
				}

				l.TraceInfo("sync", sql, " ok")
				del = del[1:]
				break
			} */
			sql := fmt.Sprintf("ALTER TABLE `%s` ADD `%s` %s %s %s;", tbl_name, f.Name, f.Type, nulltype, autoinc)
			_, err := s.conn.Exec(sql)
			if err != nil {
				panic(err)
			}

			l.TraceInfo("sync", sql, "ok")

		}

		for _, f := range del {
			sql := fmt.Sprintf("ALTER TABLE `%s`  DROP `%s` ;", tbl_name, f.Name)
			_, err := s.conn.Exec(sql)
			if err != nil {
				panic(err)
			}

			l.TraceInfo("sync", sql, " ok")
		}
	}

	return nil
}

func (s *Sync) syncObject(obj *Object) {
	tbl_name := "tbl_" + strings.ToLower(obj.Name)
	tblinfo := obj.CreateTable()
	s.syncTable(tbl_name, tblinfo, 0)

	records := obj.CreateRecordTable()
	for k := range records {
		tbl_name = strings.ToLower(fmt.Sprintf("tbl_%s_%s", obj.Name, k))

		s.syncTable(tbl_name, records[k], 0)
	}
}

func (s *Sync) syncSysDb() {
	//取用户名
	row, _ := s.conn.Query("SELECT CURRENT_USER()")
	row.Next()
	var user string
	row.Scan(&user)
	row.Close()
	s.syncTable("sys_uid", idtbl, 1000)
	s.syncTable("role_info", roletbl, 0)
}

func (s *Sync) SyncDB(path string, drop bool, role string) {

	defer func() {
		if e := recover(); e != nil {
			fmt.Println(e)
		}
		s.conn.Close()
	}()
	if drop {
		s.conn.DropDB(s.db)
		l.TraceInfo("sync", "drop db ", s.db)
	}

	s.conn.Exec(fmt.Sprintf("CREATE DATABASE IF NOT EXISTS `%s`", s.db))
	s.conn.UseDB(s.db)
	s.syncSysDb()

	for k, t := range userTables {
		s.syncTable(k, t, 0)
	}

	LoadAllDef(path)

	for _, obj := range Defs {
		if obj.Name == role || obj.Persistent == "true" {
			s.syncObject(obj)
			l.TraceInfo("sync", "process ", obj.Name, " complete")
		}

	}
	l.TraceInfo("sync", "all complete")
}

func CreateSyncDB(db string, datasource string) *Sync {
	s := &Sync{}
	c, err := NewConn("mysql", datasource)
	if err != nil {
		l.TraceInfo("dbmgr", err)
		panic(err)
	}

	var conn *MySql
	var ok bool
	if conn, ok = c.(*MySql); !ok {
		panic("convert failed")
	}
	s.conn = conn
	s.db = db
	s.datasrouce = datasource
	return s
}
<commit_msg>修正数据库同步的问题<commit_after>package mysqldb

import (
	"fmt"
	l "libs/log"
	. "libs/parser"
	"strings"
)

var (

	//ID表
	idtbl = TblField{
		FieldList: []Field{
			Field{"id", "BIGINT(20) UNSIGNED", false, true},
			Field{"stub", "VARCHAR(10)", false, false},
		},
		IndexInfo: map[string][]string{
			"PRIMARY":             []string{"id"},
			"UNIQUE INDEX `stub`": []string{"stub"},
		},
	}
	//角色表
	roletbl = TblField{
		FieldList: []Field{
			Field{"uid", "BIGINT(20) UNSIGNED", false, false},
			Field{"account", "CHAR(100)", false, false},
			Field{"rolename", "CHAR(64)", false, false},
			Field{"createtime", "DATETIME", false, false},
			Field{"lastlogintime", "DATETIME", false, false},
			Field{"locktime", "DATETIME", false, false},
			Field{"roleindex", "TINYINT(1) UNSIGNED", false, false},
			Field{"roleinfo", "TEXT", true, false},
			Field{"entity", "CHAR(100)", false, false},
			Field{"deleted", "TINYINT(1) UNSIGNED", true, false},
			Field{"locked", "TINYINT(1) UNSIGNED", true, false},
			Field{"status", "TINYINT(1) UNSIGNED", false, false},
			Field{"serverid", "CHAR(64)", true, false},
			Field{"scene", "CHAR(64)", true, false},
			Field{"scene_x", "FLOAT", true, false},
			Field{"scene_y", "FLOAT", true, false},
			Field{"scene_z", "FLOAT", true, false},
			Field{"scene_dir", "FLOAT", true, false},
			Field{"landtimes", "INT(10) UNSIGNED", false, false},
		},
		IndexInfo: map[string][]string{
			"PRIMARY":      []string{"uid"},
			"INDEX `role`": []string{"account", "rolename"},
		},
	}
)

type Sync struct {
	conn       *MySql
	db         string
	datasrouce string
}

func (s *Sync) syncTable(tbl_name string, tblinfo TblField, incNum int) error {
	if oldfield, empty, _ := s.conn.QueryTblInfo(s.db, tbl_name); empty {
		err := s.conn.CreateTable(tbl_name, tblinfo, incNum)
		if err != nil {
			panic(err)
		}
		l.TraceInfo("sync", "create table ", tbl_name, " ok")
	} else {
		alert := make([]Field, 0, 10)
		add := make([]Field, 0, 10)
		for _, newf := range tblinfo.FieldList {
			find := false
			for _, oldf := range oldfield.FieldList {
				if strings.ToLower(newf.Name) == strings.ToLower(oldf.Name) {
					if strings.ToLower(oldf.Type) != strings.ToLower(newf.Type) {
						alert = append(alert, newf)
					}
					find = true
					break
				}
			}
			if !find {
				add = append(add, newf)
			}
		}

		del := make([]Field, 0, 10)
		for _, oldf := range oldfield.FieldList {
			find := false
			for _, newf := range tblinfo.FieldList {
				if strings.ToLower(newf.Name) == strings.ToLower(oldf.Name) {
					find = true
					break
				}
			}
			if !find {
				del = append(del, oldf)
			}
		}

		for _, f := range alert {
			var nulltype string
			if !f.IsNull {
				nulltype = "NOT NULL"
			} else {
				nulltype = "DEFAULT NULL"
			}
			autoinc := ""
			if f.AutoInc {
				autoinc = " AUTO_INCREMENT"
			}

			sql := fmt.Sprintf("ALTER TABLE `%s` CHANGE `%s` `%s` %s %s %s;", tbl_name, f.Name, f.Name, f.Type, nulltype, autoinc)
			_, err := s.conn.Exec(sql)
			if err != nil {
				panic(err)
			}

			l.TraceInfo("sync", sql, " ok")
		}

		for _, f := range add {
			var nulltype string
			if !f.IsNull {
				nulltype = "NOT NULL"
			} else {
				nulltype = "DEFAULT NULL"
			}
			autoinc := ""
			if f.AutoInc {
				autoinc = " AUTO_INCREMENT"
			}

			/*if len(del) > 0 {
				sql := fmt.Sprintf("ALTER TABLE `%s` CHANGE `%s` `%s` %s %s %s;", tbl_name, del[0].Name, f.Name, f.Type, nulltype, autoinc)
				_, err := s.conn.Exec(sql)
				if err != nil {
					panic(err)
				}

				l.TraceInfo("sync", sql, " ok")
				del = del[1:]
				break
			} */
			sql := fmt.Sprintf("ALTER TABLE `%s` ADD `%s` %s %s %s;", tbl_name, f.Name, f.Type, nulltype, autoinc)
			_, err := s.conn.Exec(sql)
			if err != nil {
				panic(err)
			}

			l.TraceInfo("sync", sql, "ok")

		}

		for _, f := range del {
			sql := fmt.Sprintf("ALTER TABLE `%s`  DROP `%s` ;", tbl_name, f.Name)
			_, err := s.conn.Exec(sql)
			if err != nil {
				panic(err)
			}

			l.TraceInfo("sync", sql, " ok")
		}
	}

	return nil
}

func (s *Sync) syncObject(obj *Object) {
	tbl_name := "tbl_" + strings.ToLower(obj.Name)
	tblinfo := obj.CreateTable()
	s.syncTable(tbl_name, tblinfo, 0)

	records := obj.CreateRecordTable()
	for k := range records {
		tbl_name = strings.ToLower(fmt.Sprintf("tbl_%s_%s", obj.Name, k))

		s.syncTable(tbl_name, records[k], 0)
	}
}

func (s *Sync) syncSysDb() {
	//取用户名
	row, _ := s.conn.Query("SELECT CURRENT_USER()")
	row.Next()
	var user string
	row.Scan(&user)
	row.Close()
	s.syncTable("sys_uid", idtbl, 1000)
	s.syncTable("role_info", roletbl, 0)
}

func (s *Sync) SyncDB(path string, drop bool, role string) {

	defer func() {
		if e := recover(); e != nil {
			fmt.Println(e)
		}
		s.conn.Close()
	}()
	if drop {
		s.conn.DropDB(s.db)
		l.TraceInfo("sync", "drop db ", s.db)
	}

	s.conn.Exec(fmt.Sprintf("CREATE DATABASE IF NOT EXISTS `%s`", s.db))
	s.conn.UseDB(s.db)
	s.syncSysDb()

	for k, t := range userTables {
		s.syncTable(k, t, 0)
	}

	LoadAllDef(path)

	for _, obj := range Defs {
		if obj.Name == role || obj.Persistent == "true" {
			s.syncObject(obj)
			l.TraceInfo("sync", "process ", obj.Name, " complete")
		}

	}
	l.TraceInfo("sync", "all complete")
}

func CreateSyncDB(db string, datasource string) *Sync {
	s := &Sync{}
	c, err := NewMysqlConn(datasource)
	if err != nil {
		l.TraceInfo("dbmgr", err)
		panic(err)
	}

	s.conn = conn
	s.db = db
	s.datasrouce = datasource
	return s
}
<|endoftext|>


<commit_before>package enroll

import (
	"time"

	"github.com/satori/go.uuid"
)

type Payload struct {
	PayloadType         string      `json:"type" db:"type"`
	PayloadVersion      int         `json:"version" db:"version"`
	PayloadIdentifier   string      `json:"identifier" db:"identifier"`
	PayloadUUID         string      `json:"uuid" db:"uuid"`
	PayloadDisplayName  string      `json:"displayname" db:"displayname"`
	PayloadDescription  string      `json:"description,omitempty" db:"description"`
	PayloadOrganization string      `json:"organization,omitempty" db:"organization"`
	PayloadScope        string      `json:"scope" db:"scope" plist:",omitempty"`
	PayloadContent      interface{} `json:"content,omitempty" plist:"PayloadContent,omitempty"`
}

type Profile struct {
	PayloadContent           []interface{}     `json:"content,omitempty" db:"content"`
	PayloadDescription       string            `json:"description,omitempty" db:"description"`
	PayloadDisplayName       string            `json:"displayname,omitempty" db:"displayname"`
	PayloadExpirationDate    *time.Time        `json:"expiration_date,omitempty" db:"expiration_date" plist:",omitempty"`
	PayloadIdentifier        string            `json:"identifier" db:"identifier"`
	PayloadOrganization      string            `json:"organization,omitempty" db:"organization"`
	PayloadUUID              string            `json:"uuid" db:"uuid"`
	PayloadRemovalDisallowed bool              `json:"removal_disallowed" db:"removal_disallowed" plist:",omitempty"`
	PayloadType              string            `json:"type" db:"type"`
	PayloadVersion           int               `json:"version" db:"version"`
	PayloadScope             string            `json:"scope" db:"scope" plist:",omitempty"`
	RemovalDate              *time.Time        `json:"removal_date" db:"removal_date" plist:"-" plist:",omitempty"`
	DurationUntilRemoval     float32           `json:"duration_until_removal" db:"duration_until_removal" plist:",omitempty"`
	ConsentText              map[string]string `json:"consent_text" db:"consent_text" plist:",omitempty"`
}

func NewProfile() *Profile {
	payloadUuid := uuid.NewV4()

	return &Profile{
		PayloadVersion: 1,
		PayloadType:    "Configuration",
		PayloadUUID:    payloadUuid.String(),
	}
}

func NewPayload(payloadType string) *Payload {
	payloadUuid := uuid.NewV4()

	return &Payload{
		PayloadVersion: 1,
		PayloadType:    payloadType,
		PayloadUUID:    payloadUuid.String(),
	}
}

type SCEPPayloadContent struct {
	CAFingerprint []byte `plist:"CAFingerprint,omitempty"` // NSData
	Challenge     string `plist:"Challenge,omitempty"`
	Keysize       int
	KeyType       string `plist:"KeyType"`
	KeyUsage      int    `plist:"KeyUsage"`
	Name          string
	Subject       [][][]string `plist:"Subject,omitempty"`
	URL           string
}

// AccessRights define the management rights of the MDM server over the device.
// May not be zero. If 2 is specified, 1 must also be specified. If 128 is specified, 64 must also be specified.
type AccessRights int

const (
	// Allow inspection of installed configuration profiles.
	ProfileInspection AccessRights = 1 << iota

	// Allow installation and removal of configuration profiles.
	ProfileInstallAndRemoval

	// Allow device lock and passcode removal.
	DeviceLock

	// Allow device erase.
	DeviceErase

	// Allow query of Device Information (device capacity, serial number).
	DeviceInformationQuery

	// 	Allow query of Network Information (phone/SIM numbers, MAC addresses).
	NetworkInformationQuery

	// Allow inspection of installed provisioning profiles.
	ProvisioningProfileInspection

	//  Allow installation and removal of provisioning profiles.
	ProvisioningProfileInstallAndRemoval

	// Allow inspection of installed applications.
	ApplicationInspection

	// Allow restriction-related queries.
	RestrictionQuery

	// Allow security-related queries.
	SecurityQuery

	// Allow manipulation of settings.
	// Availability: Available in iOS 5.0 and later. Available in macOS 10.9 for certain commands.
	SettingsManipulation

	// Allow app management.
	// Availability: Available in iOS 5.0 and later. Available in macOS 10.9 for certain commands.
	AppManagement
)

func allRights() AccessRights {
	return ProfileInspection |
		ProfileInstallAndRemoval |
		DeviceLock |
		DeviceErase |
		DeviceInformationQuery |
		NetworkInformationQuery |
		ProvisioningProfileInspection |
		ProvisioningProfileInstallAndRemoval |
		ApplicationInspection |
		RestrictionQuery |
		SecurityQuery |
		SettingsManipulation |
		AppManagement
}

// TODO: Actually this is one of those non-nested payloads that doesnt respect the PayloadContent key.
type MDMPayloadContent struct {
	Payload
	AccessRights            AccessRights
	CheckInURL              string
	CheckOutWhenRemoved     bool
	IdentityCertificateUUID string
	ServerCapabilities      []string `plist:"ServerCapabilities,omitempty"`
	SignMessage             bool     `plist:"SignMessage,omitempty"`
	ServerURL               string
	Topic                   string
}

type ProfileServicePayload struct {
	URL              string
	Challenge        string `plist:",omitempty"`
	DeviceAttributes []string
}
<commit_msg>Revert "Correct SCEP profile payload (#492)" (#652)<commit_after>package enroll

import (
	"time"

	"github.com/satori/go.uuid"
)

type Payload struct {
	PayloadType         string      `json:"type" db:"type"`
	PayloadVersion      int         `json:"version" db:"version"`
	PayloadIdentifier   string      `json:"identifier" db:"identifier"`
	PayloadUUID         string      `json:"uuid" db:"uuid"`
	PayloadDisplayName  string      `json:"displayname" db:"displayname"`
	PayloadDescription  string      `json:"description,omitempty" db:"description"`
	PayloadOrganization string      `json:"organization,omitempty" db:"organization"`
	PayloadScope        string      `json:"scope" db:"scope" plist:",omitempty"`
	PayloadContent      interface{} `json:"content,omitempty" plist:"PayloadContent,omitempty"`
}

type Profile struct {
	PayloadContent           []interface{}     `json:"content,omitempty" db:"content"`
	PayloadDescription       string            `json:"description,omitempty" db:"description"`
	PayloadDisplayName       string            `json:"displayname,omitempty" db:"displayname"`
	PayloadExpirationDate    *time.Time        `json:"expiration_date,omitempty" db:"expiration_date" plist:",omitempty"`
	PayloadIdentifier        string            `json:"identifier" db:"identifier"`
	PayloadOrganization      string            `json:"organization,omitempty" db:"organization"`
	PayloadUUID              string            `json:"uuid" db:"uuid"`
	PayloadRemovalDisallowed bool              `json:"removal_disallowed" db:"removal_disallowed" plist:",omitempty"`
	PayloadType              string            `json:"type" db:"type"`
	PayloadVersion           int               `json:"version" db:"version"`
	PayloadScope             string            `json:"scope" db:"scope" plist:",omitempty"`
	RemovalDate              *time.Time        `json:"removal_date" db:"removal_date" plist:"-" plist:",omitempty"`
	DurationUntilRemoval     float32           `json:"duration_until_removal" db:"duration_until_removal" plist:",omitempty"`
	ConsentText              map[string]string `json:"consent_text" db:"consent_text" plist:",omitempty"`
}

func NewProfile() *Profile {
	payloadUuid := uuid.NewV4()

	return &Profile{
		PayloadVersion: 1,
		PayloadType:    "Configuration",
		PayloadUUID:    payloadUuid.String(),
	}
}

func NewPayload(payloadType string) *Payload {
	payloadUuid := uuid.NewV4()

	return &Payload{
		PayloadVersion: 1,
		PayloadType:    payloadType,
		PayloadUUID:    payloadUuid.String(),
	}
}

type SCEPPayloadContent struct {
	CAFingerprint []byte `plist:"CAFingerprint,omitempty"` // NSData
	Challenge     string `plist:"Challenge,omitempty"`
	Keysize       int
	KeyType       string `plist:"Key Type"`
	KeyUsage      int    `plist:"Key Usage"`
	Name          string
	Subject       [][][]string `plist:"Subject,omitempty"`
	URL           string
}

// AccessRights define the management rights of the MDM server over the device.
// May not be zero. If 2 is specified, 1 must also be specified. If 128 is specified, 64 must also be specified.
type AccessRights int

const (
	// Allow inspection of installed configuration profiles.
	ProfileInspection AccessRights = 1 << iota

	// Allow installation and removal of configuration profiles.
	ProfileInstallAndRemoval

	// Allow device lock and passcode removal.
	DeviceLock

	// Allow device erase.
	DeviceErase

	// Allow query of Device Information (device capacity, serial number).
	DeviceInformationQuery

	// 	Allow query of Network Information (phone/SIM numbers, MAC addresses).
	NetworkInformationQuery

	// Allow inspection of installed provisioning profiles.
	ProvisioningProfileInspection

	//  Allow installation and removal of provisioning profiles.
	ProvisioningProfileInstallAndRemoval

	// Allow inspection of installed applications.
	ApplicationInspection

	// Allow restriction-related queries.
	RestrictionQuery

	// Allow security-related queries.
	SecurityQuery

	// Allow manipulation of settings.
	// Availability: Available in iOS 5.0 and later. Available in macOS 10.9 for certain commands.
	SettingsManipulation

	// Allow app management.
	// Availability: Available in iOS 5.0 and later. Available in macOS 10.9 for certain commands.
	AppManagement
)

func allRights() AccessRights {
	return ProfileInspection |
		ProfileInstallAndRemoval |
		DeviceLock |
		DeviceErase |
		DeviceInformationQuery |
		NetworkInformationQuery |
		ProvisioningProfileInspection |
		ProvisioningProfileInstallAndRemoval |
		ApplicationInspection |
		RestrictionQuery |
		SecurityQuery |
		SettingsManipulation |
		AppManagement
}

// TODO: Actually this is one of those non-nested payloads that doesnt respect the PayloadContent key.
type MDMPayloadContent struct {
	Payload
	AccessRights            AccessRights
	CheckInURL              string
	CheckOutWhenRemoved     bool
	IdentityCertificateUUID string
	ServerCapabilities      []string `plist:"ServerCapabilities,omitempty"`
	SignMessage             bool     `plist:"SignMessage,omitempty"`
	ServerURL               string
	Topic                   string
}

type ProfileServicePayload struct {
	URL              string
	Challenge        string `plist:",omitempty"`
	DeviceAttributes []string
}
<|endoftext|>