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

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"sync"
	"sync/atomic"

	"github.com/01org/cc-oci-runtime/proxy/api"

	"github.com/golang/glog"
)

// Main struct holding the proxy state
type proxy struct {
	// Protect concurrent accesses from separate client goroutines to this
	// structure fields
	sync.Mutex

	// proxy socket
	listener net.Listener

	// vms are hashed by their containerID
	vms map[string]*vm
}

// Represents a client, either a cc-oci-runtime or cc-shim process having
// opened a socket to the proxy
type client struct {
	id    uint64
	proxy *proxy
	vm    *vm

	conn net.Conn
}

func (c *client) info(lvl glog.Level, msg string) {
	if !glog.V(lvl) {
		return
	}
	glog.Infof("[client #%d] %s", c.id, msg)
}

func (c *client) infof(lvl glog.Level, fmt string, a ...interface{}) {
	if !glog.V(lvl) {
		return
	}
	a = append(a, 0)
	copy(a[1:], a[0:])
	a[0] = c.id
	glog.Infof("[client #%d] "+fmt, a...)
}

// "hello"
func helloHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	hello := api.Hello{}

	if err := json.Unmarshal(data, &hello); err != nil {
		response.SetError(err)
		return
	}

	if hello.ContainerID == "" || hello.CtlSerial == "" || hello.IoSerial == "" {
		response.SetErrorMsg("malformed hello command")
	}

	proxy := client.proxy
	proxy.Lock()
	if _, ok := proxy.vms[hello.ContainerID]; ok {

		proxy.Unlock()
		response.SetErrorf("%s: container already registered",
			hello.ContainerID)
		return
	}

	client.infof(1, "hello(containerId=%s,ctlSerial=%s,ioSerial=%s", hello.ContainerID,
		hello.CtlSerial, hello.IoSerial)

	vm := newVM(hello.ContainerID, hello.CtlSerial, hello.IoSerial)
	proxy.vms[hello.ContainerID] = vm
	proxy.Unlock()

	if err := vm.Connect(); err != nil {
		proxy.Lock()
		delete(proxy.vms, hello.ContainerID)
		proxy.Unlock()
		response.SetError(err)
		return
	}

	client.vm = vm
}

// "attach"
func attachHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	proxy := client.proxy

	attach := api.Attach{}
	if err := json.Unmarshal(data, &attach); err != nil {
		response.SetError(err)
		return
	}

	proxy.Lock()
	vm := proxy.vms[attach.ContainerID]
	proxy.Unlock()

	if vm == nil {
		response.SetErrorf("unknown containerID: %s", attach.ContainerID)
		return
	}

	client.infof(1, "attach(containerId=%s)", attach.ContainerID)

	client.vm = vm
}

// "bye"
func byeHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	proxy := client.proxy

	bye := api.Bye{}
	if err := json.Unmarshal(data, &bye); err != nil {
		response.SetError(err)
		return
	}

	proxy.Lock()
	vm := proxy.vms[bye.ContainerID]
	proxy.Unlock()

	if vm == nil {
		response.SetErrorf("unknown containerID: %s", bye.ContainerID)
		return
	}

	client.info(1, "bye()")

	proxy.Lock()
	delete(proxy.vms, vm.containerID)
	proxy.Unlock()

	client.vm = nil
	vm.Close()
}

// "allocateIO"
func allocateIoHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	vm := client.vm

	allocateIo := api.AllocateIo{}
	if err := json.Unmarshal(data, &allocateIo); err != nil {
		response.SetError(err)
		return
	}

	if allocateIo.NStreams < 1 || allocateIo.NStreams > 2 {
		response.SetErrorf("asking for unexpected number of streams (%d)",
			allocateIo.NStreams)
	}

	if vm == nil {
		response.SetErrorMsg("client not attached to a vm")
		return
	}

	client.infof(1, "allocateIo(nStreams=%d)", allocateIo.NStreams)

	// We'll send c0 to the client, keep c1
	c0, c1, err := Socketpair()
	if err != nil {
		response.SetError(err)
		return
	}

	f0, err := c0.File()
	if err != nil {
		response.SetError(err)
		return
	}

	ioBase := vm.AllocateIo(allocateIo.NStreams, client.id, c1)

	client.infof(1, "-> %d streams allocated, ioBase=%d", allocateIo.NStreams, ioBase)

	response.AddResult("ioBase", ioBase)
	response.SetFile(f0)

	// File() dups the underlying fd, so it's safe to close c0 here (will
	// keep the c0 <-> c1 connection alive).
	c0.Close()
}

// "hyper"
func hyperHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	hyper := api.Hyper{}
	vm := client.vm

	if err := json.Unmarshal(data, &hyper); err != nil {
		response.SetError(err)
		return
	}

	if vm == nil {
		response.SetErrorMsg("client not attached to a vm")
		return
	}

	client.infof(1, "cmd=%s, data=%s)", hyper.HyperName, hyper.Data)

	err := vm.SendMessage(hyper.HyperName, hyper.Data)
	response.SetError(err)
}

func newProxy() *proxy {
	return &proxy{
		vms: make(map[string]*vm),
	}
}

// This variable is populated at link time with the value of:
//   ${locatestatedir}/run/cc-oci-runtime/proxy
var socketPath string

func (proxy *proxy) init() error {
	var l net.Listener
	var err error

	fds := listenFds()

	if len(fds) > 1 {
		return fmt.Errorf("too many activated sockets (%d)", len(fds))
	} else if len(fds) == 1 {
		fd := fds[0]
		l, err = net.FileListener(fd)
		if err != nil {
			return fmt.Errorf("couldn't listen on socket: %v", err)
		}
	} else {
		if err = os.Remove(socketPath); err != nil && !os.IsNotExist(err) {
			return fmt.Errorf("couldn't remove exiting socket: %v", err)
		}
		l, err = net.ListenUnix("unix", &net.UnixAddr{Name: socketPath, Net: "unix"})
		if err != nil {
			return fmt.Errorf("couldn't create AF_UNIX socket: %v", err)
		}
		if err = os.Chmod(socketPath, 0666|os.ModeSocket); err != nil {
			return fmt.Errorf("couldn't set mode on socket: %v", err)
		}
	}

	proxy.listener = l

	return nil
}

var nextClientID = uint64(1)

func (proxy *proxy) serveNewClient(proto *protocol, newConn net.Conn) {
	newClient := &client{
		id:    nextClientID,
		proxy: proxy,
		conn:  newConn,
	}

	atomic.AddUint64(&nextClientID, 1)

	// Unfortunately it's hard to find out information on the peer
	// at the other end of a unix socket. We use a per-client ID to
	// identify connections.
	newClient.info(1, "client connected")

	if err := proto.Serve(newConn, newClient); err != nil && err != io.EOF {
		newClient.infof(1, "error serving client: %v", err)
	}

	newConn.Close()
	newClient.info(1, "connection closed")
}

func (proxy *proxy) serve() {

	// Define the client (runtime/shim) <-> proxy protocol
	proto := newProtocol()
	proto.Handle("hello", helloHandler)
	proto.Handle("attach", attachHandler)
	proto.Handle("bye", byeHandler)
	proto.Handle("allocateIO", allocateIoHandler)
	proto.Handle("hyper", hyperHandler)

	glog.V(1).Info("proxy started")

	for {
		conn, err := proxy.listener.Accept()
		if err != nil {
			fmt.Fprintln(os.Stderr, "couldn't accept connection:", err)
			continue
		}

		go proxy.serveNewClient(proto, conn)
	}
}

func proxyMain() {
	proxy := newProxy()
	if err := proxy.init(); err != nil {
		fmt.Fprintln(os.Stderr, "init:", err.Error())
		os.Exit(1)
	}
	proxy.serve()
}

func initLogging() {
	// We print logs on stderr by default.
	flag.Set("logtostderr", "true")

	// It can be practical to use an environment variable to trigger a verbose output
	level := os.Getenv("CC_PROXY_LOG_LEVEL")
	if level != "" {
		flag.Set("v", level)
	}
}

func main() {
	initLogging()

	flag.Parse()
	defer glog.Flush()

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

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"sync"
	"sync/atomic"

	"github.com/01org/cc-oci-runtime/proxy/api"

	"github.com/golang/glog"
)

// Main struct holding the proxy state
type proxy struct {
	// Protect concurrent accesses from separate client goroutines to this
	// structure fields
	sync.Mutex

	// proxy socket
	listener net.Listener

	// vms are hashed by their containerID
	vms map[string]*vm
}

// Represents a client, either a cc-oci-runtime or cc-shim process having
// opened a socket to the proxy
type client struct {
	id    uint64
	proxy *proxy
	vm    *vm

	conn net.Conn
}

func (c *client) info(lvl glog.Level, msg string) {
	if !glog.V(lvl) {
		return
	}
	glog.Infof("[client #%d] %s", c.id, msg)
}

func (c *client) infof(lvl glog.Level, fmt string, a ...interface{}) {
	if !glog.V(lvl) {
		return
	}
	a = append(a, 0)
	copy(a[1:], a[0:])
	a[0] = c.id
	glog.Infof("[client #%d] "+fmt, a...)
}

// "hello"
func helloHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	hello := api.Hello{}

	if err := json.Unmarshal(data, &hello); err != nil {
		response.SetError(err)
		return
	}

	if hello.ContainerID == "" || hello.CtlSerial == "" || hello.IoSerial == "" {
		response.SetErrorMsg("malformed hello command")
	}

	proxy := client.proxy
	proxy.Lock()
	if _, ok := proxy.vms[hello.ContainerID]; ok {

		proxy.Unlock()
		response.SetErrorf("%s: container already registered",
			hello.ContainerID)
		return
	}

	client.infof(1, "hello(containerId=%s,ctlSerial=%s,ioSerial=%s", hello.ContainerID,
		hello.CtlSerial, hello.IoSerial)

	vm := newVM(hello.ContainerID, hello.CtlSerial, hello.IoSerial)
	proxy.vms[hello.ContainerID] = vm
	proxy.Unlock()

	if err := vm.Connect(); err != nil {
		proxy.Lock()
		delete(proxy.vms, hello.ContainerID)
		proxy.Unlock()
		response.SetError(err)
		return
	}

	client.vm = vm
}

// "attach"
func attachHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	proxy := client.proxy

	attach := api.Attach{}
	if err := json.Unmarshal(data, &attach); err != nil {
		response.SetError(err)
		return
	}

	proxy.Lock()
	vm := proxy.vms[attach.ContainerID]
	proxy.Unlock()

	if vm == nil {
		response.SetErrorf("unknown containerID: %s", attach.ContainerID)
		return
	}

	client.infof(1, "attach(containerId=%s)", attach.ContainerID)

	client.vm = vm
}

// "bye"
func byeHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	proxy := client.proxy

	bye := api.Bye{}
	if err := json.Unmarshal(data, &bye); err != nil {
		response.SetError(err)
		return
	}

	proxy.Lock()
	vm := proxy.vms[bye.ContainerID]
	proxy.Unlock()

	if vm == nil {
		response.SetErrorf("unknown containerID: %s", bye.ContainerID)
		return
	}

	client.info(1, "bye()")

	proxy.Lock()
	delete(proxy.vms, vm.containerID)
	proxy.Unlock()

	client.vm = nil
	vm.Close()
}

// "allocateIO"
func allocateIoHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	vm := client.vm

	allocateIo := api.AllocateIo{}
	if err := json.Unmarshal(data, &allocateIo); err != nil {
		response.SetError(err)
		return
	}

	if allocateIo.NStreams < 1 || allocateIo.NStreams > 2 {
		response.SetErrorf("asking for unexpected number of streams (%d)",
			allocateIo.NStreams)
	}

	if vm == nil {
		response.SetErrorMsg("client not attached to a vm")
		return
	}

	client.infof(1, "allocateIo(nStreams=%d)", allocateIo.NStreams)

	// We'll send c0 to the client, keep c1
	c0, c1, err := Socketpair()
	if err != nil {
		response.SetError(err)
		return
	}

	f0, err := c0.File()
	if err != nil {
		response.SetError(err)
		return
	}

	ioBase := vm.AllocateIo(allocateIo.NStreams, client.id, c1)

	client.infof(1, "-> %d streams allocated, ioBase=%d", allocateIo.NStreams, ioBase)

	response.AddResult("ioBase", ioBase)
	response.SetFile(f0)

	// File() dups the underlying fd, so it's safe to close c0 here (will
	// keep the c0 <-> c1 connection alive).
	c0.Close()
}

// "hyper"
func hyperHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	hyper := api.Hyper{}
	vm := client.vm

	if err := json.Unmarshal(data, &hyper); err != nil {
		response.SetError(err)
		return
	}

	if vm == nil {
		response.SetErrorMsg("client not attached to a vm")
		return
	}

	client.infof(1, "cmd=%s, data=%s)", hyper.HyperName, hyper.Data)

	err := vm.SendMessage(hyper.HyperName, hyper.Data)
	response.SetError(err)
}

func newProxy() *proxy {
	return &proxy{
		vms: make(map[string]*vm),
	}
}

// This variable is populated at link time with the value of:
//   ${locatestatedir}/run/cc-oci-runtime/proxy
var socketPath string

func (proxy *proxy) init() error {
	var l net.Listener
	var err error

	fds := listenFds()

	if len(fds) > 1 {
		return fmt.Errorf("too many activated sockets (%d)", len(fds))
	} else if len(fds) == 1 {
		fd := fds[0]
		l, err = net.FileListener(fd)
		if err != nil {
			return fmt.Errorf("couldn't listen on socket: %v", err)
		}
	} else {
		if err = os.Remove(socketPath); err != nil && !os.IsNotExist(err) {
			return fmt.Errorf("couldn't remove exiting socket: %v", err)
		}
		l, err = net.ListenUnix("unix", &net.UnixAddr{Name: socketPath, Net: "unix"})
		if err != nil {
			return fmt.Errorf("couldn't create AF_UNIX socket: %v", err)
		}
		if err = os.Chmod(socketPath, 0660|os.ModeSocket); err != nil {
			return fmt.Errorf("couldn't set mode on socket: %v", err)
		}
	}

	proxy.listener = l

	return nil
}

var nextClientID = uint64(1)

func (proxy *proxy) serveNewClient(proto *protocol, newConn net.Conn) {
	newClient := &client{
		id:    nextClientID,
		proxy: proxy,
		conn:  newConn,
	}

	atomic.AddUint64(&nextClientID, 1)

	// Unfortunately it's hard to find out information on the peer
	// at the other end of a unix socket. We use a per-client ID to
	// identify connections.
	newClient.info(1, "client connected")

	if err := proto.Serve(newConn, newClient); err != nil && err != io.EOF {
		newClient.infof(1, "error serving client: %v", err)
	}

	newConn.Close()
	newClient.info(1, "connection closed")
}

func (proxy *proxy) serve() {

	// Define the client (runtime/shim) <-> proxy protocol
	proto := newProtocol()
	proto.Handle("hello", helloHandler)
	proto.Handle("attach", attachHandler)
	proto.Handle("bye", byeHandler)
	proto.Handle("allocateIO", allocateIoHandler)
	proto.Handle("hyper", hyperHandler)

	glog.V(1).Info("proxy started")

	for {
		conn, err := proxy.listener.Accept()
		if err != nil {
			fmt.Fprintln(os.Stderr, "couldn't accept connection:", err)
			continue
		}

		go proxy.serveNewClient(proto, conn)
	}
}

func proxyMain() {
	proxy := newProxy()
	if err := proxy.init(); err != nil {
		fmt.Fprintln(os.Stderr, "init:", err.Error())
		os.Exit(1)
	}
	proxy.serve()
}

func initLogging() {
	// We print logs on stderr by default.
	flag.Set("logtostderr", "true")

	// It can be practical to use an environment variable to trigger a verbose output
	level := os.Getenv("CC_PROXY_LOG_LEVEL")
	if level != "" {
		flag.Set("v", level)
	}
}

func main() {
	initLogging()

	flag.Parse()
	defer glog.Flush()

	proxyMain()
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"math/rand"
	"net/http"
	"net/http/httputil"
	"path"
	"stathat.com/c/consistent"
	"strings"
)

// Sattolo shuffle
// https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#Sattolo.27s_algorithm
func shuffle(list []string) {
	for i := range list {
		j := rand.Intn(i + 1)
		list[i], list[j] = list[j], list[i]
	}
}

type CacheProxy struct {
	backends *consistent.Consistent
}

func New() *CacheProxy {
	c := new(CacheProxy)
	c.backends = consistent.New()
	return c
}

func (c *CacheProxy) Manager() *http.ServeMux {
	manager := http.NewServeMux()

	// RESTful interface for cache members
	// GET /members/
	// 	  Returns a list of current cache members

	// PUT /members/member[:port]
	//    Adds "member" as a cache member

	// DELETE /members/member[:port]
	//    Delets "member" from cache ring

	manager.HandleFunc("/members/", func(w http.ResponseWriter, req *http.Request) {
		switch {
		case req.Method == "PUT":
			c.backends.Add(path.Base(req.URL.Path))
		case req.Method == "DELETE":
			c.backends.Remove(path.Base(req.URL.Path))
		}
		io.WriteString(w, strings.Join(c.backends.Members(), ","))
	})

	return manager
}

// Reverse proxy that selects the backend by nearest match to the request URL
// on the consistent hash ring.
func (c *CacheProxy) Proxy() *httputil.ReverseProxy {
	return &httputil.ReverseProxy{
		Director: func(req *http.Request) {
			req.URL.Scheme = "http"
			backends, _ := c.backends.GetN(req.URL.Path, 3)
			shuffle(backends)
			req.URL.Host = backends[0]
			req.Header.Set("X-Backends", strings.Join(backends, ","))
		},
		Transport: &http.Transport{},
	}
}

func main() {
	proxy := New()

	// Replace with runtime arguments for initialization
	proxy.backends.Add("localhost:9091")
	proxy.backends.Add("localhost:9092")
	proxy.backends.Add("localhost:9093")
	proxy.backends.Add("localhost:9094")
	proxy.backends.Add("localhost:9095")

	// Initialize and run manager server in background via goroutine
	go http.ListenAndServe(":9190", proxy.Manager())

	http.ListenAndServe(":9090", proxy.Proxy())
}
<commit_msg>Adding Select method<commit_after>package main

import (
	"io"
	"math/rand"
	"net/http"
	"net/http/httputil"
	"path"
	"stathat.com/c/consistent"
	"strings"
)

// Sattolo shuffle
// https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#Sattolo.27s_algorithm
func shuffle(list []string) {
	for i := range list {
		j := rand.Intn(i + 1)
		list[i], list[j] = list[j], list[i]
	}
}

type CacheProxy struct {
	backends *consistent.Consistent
}

func New() *CacheProxy {
	c := new(CacheProxy)
	c.backends = consistent.New()
	return c
}

func (c *CacheProxy) Manager() *http.ServeMux {
	manager := http.NewServeMux()

	// RESTful interface for cache members
	// GET /members/
	// 	  Returns a list of current cache members

	// PUT /members/member[:port]
	//    Adds "member" as a cache member

	// DELETE /members/member[:port]
	//    Delets "member" from cache ring

	manager.HandleFunc("/members/", func(w http.ResponseWriter, req *http.Request) {
		switch {
		case req.Method == "PUT":
			c.backends.Add(path.Base(req.URL.Path))
		case req.Method == "DELETE":
			c.backends.Remove(path.Base(req.URL.Path))
		}
		io.WriteString(w, strings.Join(c.backends.Members(), ","))
	})

	return manager
}

func (c *CacheProxy) Select(path string) []string {
	selection, _ := c.backends.GetN(path, 3)
	shuffle(selection)
	return selection
}

// Reverse proxy that selects the backend by nearest match to the request URL
// on the consistent hash ring.
func (c *CacheProxy) Proxy() *httputil.ReverseProxy {
	return &httputil.ReverseProxy{
		Director: func(req *http.Request) {
			req.URL.Scheme = "http"
			backends := c.Select(req.URL.Path)
			req.URL.Host = backends[0]
			req.Header.Set("X-Backends", strings.Join(backends, ","))
		},
		Transport: &http.Transport{},
	}
}

func main() {
	proxy := New()

	// Replace with runtime arguments for initialization
	proxy.backends.Add("localhost:9091")
	proxy.backends.Add("localhost:9092")
	proxy.backends.Add("localhost:9093")
	proxy.backends.Add("localhost:9094")
	proxy.backends.Add("localhost:9095")

	// Initialize and run manager server in background via goroutine
	go http.ListenAndServe(":9190", proxy.Manager())

	http.ListenAndServe(":9090", proxy.Proxy())
}
<|endoftext|>


<commit_before>package influxdbclient

import "fmt"

func buildQuery(fields string, filters *Filters, groupby string, serie string, from string, to string, function string) (query string) {
	if len(function) > 0 {
		query = fmt.Sprintf("SELECT %s(\"%s\") FROM \"%s\"", function, fields, serie)
	} else {
		query = fmt.Sprintf("SELECT \"%s\" FROM \"%s\"", fields, serie)
	}

	var filterQuery FilterQuery

	if len(from) > 0 {
		filterQuery.Append(fmt.Sprintf("time > '%s'", from))
	}

	if len(to) > 0 {
		filterQuery.Append(fmt.Sprintf("time < '%s'", to))
	}

	if len(*filters) > 0 {
		filterQuery.AddFilters(filters)
	}

	if len(groupby) > 0 {
		filterQuery.Append(fmt.Sprintf("GROUP BY %s", groupby))
	}

	if len(filterQuery.Content) > 0 {
		query += " WHERE " + filterQuery.Content
	}
	return
}
<commit_msg>minor fix on group by<commit_after>package influxdbclient

import "fmt"

func buildQuery(fields string, filters *Filters, groupby string, serie string, from string, to string, function string) (query string) {
	if len(function) > 0 {
		query = fmt.Sprintf("SELECT %s(\"%s\") FROM \"%s\"", function, fields, serie)
	} else {
		query = fmt.Sprintf("SELECT \"%s\" FROM \"%s\"", fields, serie)
	}

	var filterQuery FilterQuery

	if len(from) > 0 {
		filterQuery.Append(fmt.Sprintf("time > '%s'", from))
	}

	if len(to) > 0 {
		filterQuery.Append(fmt.Sprintf("time < '%s'", to))
	}

	if len(*filters) > 0 {
		filterQuery.AddFilters(filters)
	}

	if len(filterQuery.Content) > 0 {
		query += " WHERE " + filterQuery.Content
	}

	if len(groupby) > 0 {
		query += fmt.Sprintf(" GROUP BY %s", groupby)
	}
	return
}
<|endoftext|>


<commit_before>package calltracker

import "sort"

// Note: The code in this file is derived from example and source code
// in the container/heap package in the Go programming language standard
// library.

// Queue is an ordered queue of calls. The oldest calls are at the beginning,
// and the newest calls are at the end. Age is determined by comparing the
// number of each call.
type Queue []Call

// Push inserts the call c at the end of the queue without verifying the
// correctness of its ordering.
func (q *Queue) Push(c Call) {
	*q = append(*q, c)
}

// Insert inserts the call c into the queue in its correct position. The
// complexity is O(log(n)) where n = len(q).
func (q *Queue) Insert(c Call) {
	*q = append(*q, c)
	q.up(len(*q) - 1)
}

// Remove removes the element at index i from the queue.
// The complexity is O(log(n)) where n = len(q).
func (q *Queue) Remove(i int) {
	n := len(*q) - 1
	if n != i {
		q.swap(i, n)
		q.down(i, n)
		q.up(i)
	}
	*q = (*q)[0:n]
}

// Find returns the index of the first call with the given number in the queue.
// It will return -1 if such a call is not present in the queue.
func (q Queue) Find(number uint64) int {
	n := len(q)
	i := sort.Search(n, func(i int) bool { return q[i].number <= number })
	if i < n && q[i].number == number {
		return i
	}
	return -1
}

func (q Queue) less(i, j int) bool {
	return q[i].number < q[j].number
}

func (q Queue) swap(i, j int) {
	q[i], q[j] = q[j], q[i]
}

func (q Queue) up(j int) {
	for {
		i := (j - 1) / 2 // parent
		if i == j || !q.less(j, i) {
			break
		}
		q.swap(i, j)
		j = i
	}
}

func (q Queue) down(i, n int) {
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
			break
		}
		j := j1 // left child
		if j2 := j1 + 1; j2 < n && !q.less(j1, j2) {
			j = j2 // = 2*i + 2  // right child
		}
		if !q.less(j, i) {
			break
		}
		q.swap(i, j)
		i = j
	}
}
<commit_msg>Fixed comparator in queue.Find()<commit_after>package calltracker

import "sort"

// Note: The code in this file is derived from example and source code
// in the container/heap package in the Go programming language standard
// library.

// Queue is an ordered queue of calls. The oldest calls are at the beginning,
// and the newest calls are at the end. Age is determined by comparing the
// number of each call.
type Queue []Call

// Push inserts the call c at the end of the queue without verifying the
// correctness of its ordering.
func (q *Queue) Push(c Call) {
	*q = append(*q, c)
}

// Remove removes the element at index i from the queue.
// The complexity is O(log(n)) where n = len(q).
func (q *Queue) Remove(i int) {
	*q = append((*q)[:i], (*q)[i+1:]...)
}

// Find returns the index of the first call with the given number in the queue.
// It will return -1 if such a call is not present in the queue.
func (q Queue) Find(number uint64) int {
	n := len(q)
	i := sort.Search(n, func(i int) bool { return q[i].number >= number })
	if i < n && q[i].number == number {
		return i
	}
	return -1
}
<|endoftext|>


<commit_before>//This queue is FIFO based.
//
//It's linked list based with mutex lock for
//safe read and write for multi-threaded programs.
package queue

import (
	"errors"
	"sync"
)

//Internal node for holding the value
//and the pointer to the next node
type node struct {
	Value    interface{} //The data
	NextNode *node       //Next node
}

//Queue is FIFO based queue with linked-list structure
type Queue struct {
	head *node        //Pointer to the first node
	tail *node        //Pointer to the last node
	size int          //Number of nodes in the queue
	mu   sync.RWMutex //Mutex lock for atomic operations
}

//Creates a new queue
func NewQueue() *Queue {
	return &Queue{}
}

//Push add to the end of the queue
func (q *Queue) Push(i interface{}) {

	//Lock it
	q.mu.Lock()

	//Create a new node a contaier for the value
	newNode := node{i, nil}

	//Check if the queue is empty or not
	if q.head == nil {
		//Empty so set the head pointer to the new node
		q.head = &newNode
	} else {
		//Not empty direct the current last node
		//to the new last node
		q.tail.NextNode = &newNode
	}
	//And allways set the tail to the new last node
	q.tail = &newNode

	//Increase the size before return
	q.size++

	//Unlock it
	q.mu.Unlock()

	return
}

//Pop remove and return the oldest node in the queue.
//If the queue is empty error will contain "Queue is empty"
func (q *Queue) Pop() (i interface{}, err error) {

	//Check if the queue is empty
	if q.IsEmpty() {
		return nil, errors.New("Queue is empty")
	}

	//Lock it
	q.mu.Lock()

	//Take out the headNode
	var headNode = *q.head

	//Check if we only got one item left in the queue
	if q.head == q.tail {
		//Then set the tail to nil
		q.tail = nil
	}

	//Move the head pointer one step forward
	q.head = q.head.NextNode

	//Reduce the size before return
	q.size--

	//Unlock it
	q.mu.Unlock()

	return headNode.Value, nil
}

// The number of elements in the queue
func (q *Queue) Size() int {
	//Lock it
	q.mu.RLock()
	//Get the data
	size := q.size
	//Unlock it
	q.mu.RUnlock()

	return size
}

// Returns true if there is elements in the queue, otherwise false
func (q *Queue) IsEmpty() bool {
	//Lock it
	q.mu.RLock()
	//Get the data
	empty := q.head == nil
	//Unlock it
	q.mu.RUnlock()
	return empty
}
<commit_msg>Updated some comments<commit_after>//Package queue provides a FIFO queue with mutex lock for safe read and write
//
//Queue uses linked list internal structure
package queue

import (
	"errors"
	"sync"
)

//Internal node for holding the value
//and the pointer to the next node
type node struct {
	Value    interface{} //The data
	NextNode *node       //Next node
}

//Queue is FIFO based queue with linked-list structure
type Queue struct {
	head *node        //Pointer to the first node
	tail *node        //Pointer to the last node
	size int          //Number of nodes in the queue
	mu   sync.RWMutex //Mutex lock for atomic operations
}

//Creates a new queue
func NewQueue() *Queue {
	return &Queue{}
}

//Push add to the end of the queue
func (q *Queue) Push(i interface{}) {

	//Lock it
	q.mu.Lock()

	//Create a new node a contaier for the value
	newNode := node{i, nil}

	//Check if the queue is empty or not
	if q.head == nil {
		//Empty so set the head pointer to the new node
		q.head = &newNode
	} else {
		//Not empty direct the current last node
		//to the new last node
		q.tail.NextNode = &newNode
	}
	//And allways set the tail to the new last node
	q.tail = &newNode

	//Increase the size before return
	q.size++

	//Unlock it
	q.mu.Unlock()

	return
}

//Pop remove and return the oldest node in the queue.
//If the queue is empty error will contain "Queue is empty"
func (q *Queue) Pop() (i interface{}, err error) {

	//Check if the queue is empty
	if q.IsEmpty() {
		return nil, errors.New("Queue is empty")
	}

	//Lock it
	q.mu.Lock()

	//Take out the headNode
	var headNode = *q.head

	//Check if we only got one item left in the queue
	if q.head == q.tail {
		//Then set the tail to nil
		q.tail = nil
	}

	//Move the head pointer one step forward
	q.head = q.head.NextNode

	//Reduce the size before return
	q.size--

	//Unlock it
	q.mu.Unlock()

	return headNode.Value, nil
}

// The number of elements in the queue
func (q *Queue) Size() int {
	//Lock it
	q.mu.RLock()
	//Get the data
	size := q.size
	//Unlock it
	q.mu.RUnlock()

	return size
}

// Returns true if there is elements in the queue, otherwise false
func (q *Queue) IsEmpty() bool {
	//Lock it
	q.mu.RLock()
	//Get the data
	empty := q.head == nil
	//Unlock it
	q.mu.RUnlock()
	return empty
}
<|endoftext|>


<commit_before>package main

import (
    "bytes"
    "strconv"
    "strings"
    "time"

    "labix.org/v2/mgo/bson"
)

type RedisMessage struct {
    Ts            time.Time
    NumberOfBulks int64
    Bulks         []string

    Stream_id    uint32
    Tuple        *IpPortTuple
    CmdlineTuple *CmdlineTuple
    Direction    uint8

    IsRequest bool
    Message   string
}

type RedisStream struct {
    tcpStream *TcpStream

    data []byte

    parseOffset   int
    parseState    int
    bytesReceived int

    message *RedisMessage
}

type RedisTransaction struct {
    Type         string
    tuple        TcpTuple
    Src          Endpoint
    Dst          Endpoint
    ResponseTime int32
    Ts           int64
    JsTs         time.Time
    ts           time.Time
    cmdline      *CmdlineTuple

    Redis bson.M

    Request_raw  string
    Response_raw string

    timer *time.Timer
}

// Keep sorted for future command addition
var RedisCommands = map[string]struct{}{
    "APPEND":           struct{}{},
    "AUTH":             struct{}{},
    "BGREWRITEAOF":     struct{}{},
    "BGSAVE":           struct{}{},
    "BITCOUNT":         struct{}{},
    "BITOP":            struct{}{},
    "BITPOS":           struct{}{},
    "BLPOP":            struct{}{},
    "BRPOP":            struct{}{},
    "BRPOPLPUSH":       struct{}{},
    "CLIENT GETNAME":   struct{}{},
    "CLIENT KILL":      struct{}{},
    "CLIENT LIST":      struct{}{},
    "CLIENT PAUSE":     struct{}{},
    "CLIENT SETNAME":   struct{}{},
    "CONFIG GET":       struct{}{},
    "CONFIG RESETSTAT": struct{}{},
    "CONFIG REWRITE":   struct{}{},
    "CONFIG SET":       struct{}{},
    "DBSIZE":           struct{}{},
    "DEBUG OBJECT":     struct{}{},
    "DEBUG SEGFAULT":   struct{}{},
    "DECR":             struct{}{},
    "DECRBY":           struct{}{},
    "DEL":              struct{}{},
    "DISCARD":          struct{}{},
    "DUMP":             struct{}{},
    "ECHO":             struct{}{},
    "EVAL":             struct{}{},
    "EVALSHA":          struct{}{},
    "EXEC":             struct{}{},
    "EXISTS":           struct{}{},
    "EXPIRE":           struct{}{},
    "EXPIREAT":         struct{}{},
    "FLUSHALL":         struct{}{},
    "FLUSHDB":          struct{}{},
    "GET":              struct{}{},
    "GETBIT":           struct{}{},
    "GETRANGE":         struct{}{},
    "GETSET":           struct{}{},
    "HDEL":             struct{}{},
    "HEXISTS":          struct{}{},
    "HGET":             struct{}{},
    "HGETALL":          struct{}{},
    "HINCRBY":          struct{}{},
    "HINCRBYFLOAT":     struct{}{},
    "HKEYS":            struct{}{},
    "HLEN":             struct{}{},
    "HMGET":            struct{}{},
    "HMSET":            struct{}{},
    "HSCAN":            struct{}{},
    "HSET":             struct{}{},
    "HSETINX":          struct{}{},
    "HVALS":            struct{}{},
    "INCR":             struct{}{},
    "INCRBY":           struct{}{},
    "INCRBYFLOAT":      struct{}{},
    "INFO":             struct{}{},
    "KEYS":             struct{}{},
    "LASTSAVE":         struct{}{},
    "LINDEX":           struct{}{},
    "LINSERT":          struct{}{},
    "LLEN":             struct{}{},
    "LPOP":             struct{}{},
    "LPUSH":            struct{}{},
    "LPUSHX":           struct{}{},
    "LRANGE":           struct{}{},
    "LREM":             struct{}{},
    "LSET":             struct{}{},
    "LTRIM":            struct{}{},
    "MGET":             struct{}{},
    "MIGRATE":          struct{}{},
    "MONITOR":          struct{}{},
    "MOVE":             struct{}{},
    "MSET":             struct{}{},
    "MSETNX":           struct{}{},
    "MULTI":            struct{}{},
    "OBJECT":           struct{}{},
    "PERSIST":          struct{}{},
    "PEXPIRE":          struct{}{},
    "PEXPIREAT":        struct{}{},
    "PFADD":            struct{}{},
    "PFCOUNT":          struct{}{},
    "PFMERGE":          struct{}{},
    "PING":             struct{}{},
    "PSETEX":           struct{}{},
    "PSUBSCRIBE":       struct{}{},
    "PTTL":             struct{}{},
    "PUBLISH":          struct{}{},
    "PUBSUB":           struct{}{},
    "PUNSUBSCRIBE":     struct{}{},
    "QUIT":             struct{}{},
    "RANDOMKEY":        struct{}{},
    "RENAME":           struct{}{},
    "RENAMENX":         struct{}{},
    "RESTORE":          struct{}{},
    "RPOP":             struct{}{},
    "RPOPLPUSH":        struct{}{},
    "RPUSH":            struct{}{},
    "RPUSHX":           struct{}{},
    "SADD":             struct{}{},
    "SAVE":             struct{}{},
    "SCAN":             struct{}{},
    "SCARD":            struct{}{},
    "SCRIPT EXISTS":    struct{}{},
    "SCRIPT FLUSH":     struct{}{},
    "SCRIPT KILL":      struct{}{},
    "SCRIPT LOAD":      struct{}{},
    "SDIFF":            struct{}{},
    "SDIFFSTORE":       struct{}{},
    "SELECT":           struct{}{},
    "SET":              struct{}{},
    "SETBIT":           struct{}{},
    "SETEX":            struct{}{},
    "SETNX":            struct{}{},
    "SETRANGE":         struct{}{},
    "SHUTDOWN":         struct{}{},
    "SINTER":           struct{}{},
    "SINTERSTORE":      struct{}{},
    "SISMEMBER":        struct{}{},
    "SLAVEOF":          struct{}{},
    "SLOWLOG":          struct{}{},
    "SMEMBERS":         struct{}{},
    "SMOVE":            struct{}{},
    "SORT":             struct{}{},
    "SPOP":             struct{}{},
    "SRANDMEMBER":      struct{}{},
    "SREM":             struct{}{},
    "SSCAN":            struct{}{},
    "STRLEN":           struct{}{},
    "SUBSCRIBE":        struct{}{},
    "SUNION":           struct{}{},
    "SUNIONSTORE":      struct{}{},
    "SYNC":             struct{}{},
    "TIME":             struct{}{},
    "TTL":              struct{}{},
    "TYPE":             struct{}{},
    "UNSUBSCRIBE":      struct{}{},
    "UNWATCH":          struct{}{},
    "WATCH":            struct{}{},
    "ZADD":             struct{}{},
    "ZCARD":            struct{}{},
    "ZCOUNT":           struct{}{},
    "ZINCRBY":          struct{}{},
    "ZINTERSTORE":      struct{}{},
    "ZRANGE":           struct{}{},
    "ZRANGEBYSCORE":    struct{}{},
    "ZRANK":            struct{}{},
    "ZREM":             struct{}{},
    "ZREMRANGEBYLEX":   struct{}{},
    "ZREMRANGEBYRANK":  struct{}{},
    "ZREMRANGEBYSCORE": struct{}{},
    "ZREVRANGE":        struct{}{},
    "ZREVRANGEBYSCORE": struct{}{},
    "ZREVRANK":         struct{}{},
    "ZSCAN":            struct{}{},
    "ZSCORE":           struct{}{},
    "ZUNIONSTORE":      struct{}{},
}

var redisTransactionsMap = make(map[TcpTuple]*RedisTransaction, TransactionsHashSize)

func (stream *RedisStream) PrepareForNewMessage() {
    stream.message.NumberOfBulks = 0
    stream.message.Bulks = []string{}
    stream.message.IsRequest = false
}

func redisMessageParser(s *RedisStream) (bool, bool) {

    var err error
    var value string
    m := s.message

    for s.parseOffset < len(s.data) {

        if s.data[s.parseOffset] == '*' {
            //Multi Bulk Message

            line, off := readLine(s.data, s.parseOffset)

            if len(line) == 3 && line[1] == '-' && line[2] == '1' {
                //NULL Multi Bulk
                s.parseOffset = off
                value = "nil"
            } else {

                m.NumberOfBulks, err = strconv.ParseInt(line[1:], 10, 64)

                if err != nil {
                    ERR("Failed to read number of bulk messages: %s", err)
                    return false, false
                }
                s.parseOffset = off
                m.Bulks = []string{}

                continue
            }

        } else if s.data[s.parseOffset] == '$' {
            // Bulk Reply

            line, off := readLine(s.data, s.parseOffset)

            if len(line) == 3 && line[1] == '-' && line[2] == '1' {
                // NULL Bulk Reply
                value = "nil"
                s.parseOffset = off
            } else {
                length, err := strconv.ParseInt(line[1:], 10, 64)
                if err != nil {
                    ERR("Failed to read bulk message: %s", err)
                    return false, false
                }

                s.parseOffset = off

                line, off = readLine(s.data, s.parseOffset)

                if int64(len(line)) != length {
                    ERR("Wrong length of data: %d instead of %d", len(line), length)
                    return false, false
                }
                value = line
                s.parseOffset = off
            }

        } else if s.data[s.parseOffset] == ':' {
            // Integer reply
            line, off := readLine(s.data, s.parseOffset)
            n, err := strconv.ParseInt(line[1:], 10, 64)

            if err != nil {
                ERR("Failed to read integer reply: %s", err)
                return false, false
            }
            value = string(n)
            s.parseOffset = off

        } else if s.data[s.parseOffset] == '+' {
            // Status Reply
            line, off := readLine(s.data, s.parseOffset)

            value = line[1:]
            s.parseOffset = off
        } else if s.data[s.parseOffset] == '-' {
            // Error Reply
            line, off := readLine(s.data, s.parseOffset)

            value = line[1:]
            s.parseOffset = off
        } else {
            DEBUG("redis", "Unexpected message starting with %s", s.data[s.parseOffset:])
            return false, false
        }

        // add value
        if m.NumberOfBulks > 0 {
            m.NumberOfBulks = m.NumberOfBulks - 1
            m.Bulks = append(m.Bulks, value)

            if len(m.Bulks) == 1 {
                // check if it's a command
                if isRedisCommand(value) {
                    m.IsRequest = true
                }
            }

            if m.NumberOfBulks == 0 {
                // the last bulk received
                m.Message = strings.Join(m.Bulks, " ")
                return true, true
            }
        } else {
            m.Message = value
            return true, true
        }

    }   //end for

    return true, false
}

func readLine(data []byte, offset int) (string, int) {
    q := bytes.Index(data[offset:], []byte("\r\n"))
    return string(data[offset : offset+q]), offset + q + 2
}

func ParseRedis(pkt *Packet, tcp *TcpStream, dir uint8) {
    defer RECOVER("ParseRedis exception")

    if tcp.redisData[dir] == nil {
        tcp.redisData[dir] = &RedisStream{
            tcpStream: tcp,
            data:      pkt.payload,
            message:   &RedisMessage{Ts: pkt.ts},
        }
    } else {
        // concatenate bytes
        tcp.redisData[dir].data = append(tcp.redisData[dir].data, pkt.payload...)
        if len(tcp.redisData[dir].data) > TCP_MAX_DATA_IN_STREAM {
            DEBUG("redis", "Stream data too large, dropping TCP stream")
            tcp.redisData[dir] = nil
            return
        }
    }

    stream := tcp.redisData[dir]
    if stream.message == nil {
        stream.message = &RedisMessage{Ts: pkt.ts}
    }

    ok, complete := redisMessageParser(tcp.redisData[dir])

    if !ok {
        // drop this tcp stream. Will retry parsing with the next
        // segment in it
        tcp.redisData[dir] = nil
        return
    }

    if complete {

        if stream.message.IsRequest {
            DEBUG("redis", "REDIS request message: %s", stream.message.Message)
        } else {
            DEBUG("redis", "REDIS response message: %s", stream.message.Message)
        }

        // all ok, go to next level
        handleRedis(stream.message, tcp, dir)

        // and reset message
        stream.PrepareForNewMessage()
    }

}

func isRedisCommand(key string) bool {
    _, exists := RedisCommands[key]
    return exists
}

func handleRedis(m *RedisMessage, tcp *TcpStream,
    dir uint8) {

    m.Stream_id = tcp.id
    m.Tuple = tcp.tuple
    m.Direction = dir
    m.CmdlineTuple = procWatcher.FindProcessesTuple(tcp.tuple)

    if m.IsRequest {
        receivedRedisRequest(m)
    } else {
        receivedRedisResponse(m)
    }
}

func receivedRedisRequest(msg *RedisMessage) {
    // Add it to the HT
    tuple := TcpTuple{
        Src_ip:    msg.Tuple.Src_ip,
        Dst_ip:    msg.Tuple.Dst_ip,
        Src_port:  msg.Tuple.Src_port,
        Dst_port:  msg.Tuple.Dst_port,
        stream_id: msg.Stream_id,
    }

    trans := redisTransactionsMap[tuple]
    if trans != nil {
        if len(trans.Redis) != 0 {
            WARN("Two requests without a Response. Dropping old request")
        }
    } else {
        trans = &RedisTransaction{Type: "redis", tuple: tuple}
        redisTransactionsMap[tuple] = trans
    }

    trans.Redis = bson.M{
        "request": msg.Message,
    }
    trans.Request_raw = msg.Message

    trans.cmdline = msg.CmdlineTuple
    trans.ts = msg.Ts
    trans.Ts = int64(trans.ts.UnixNano() / 1000) // transactions have microseconds resolution
    trans.JsTs = msg.Ts
    trans.Src = Endpoint{
        Ip:   Ipv4_Ntoa(tuple.Src_ip),
        Port: tuple.Src_port,
        Proc: string(msg.CmdlineTuple.Src),
    }
    trans.Dst = Endpoint{
        Ip:   Ipv4_Ntoa(tuple.Dst_ip),
        Port: tuple.Dst_port,
        Proc: string(msg.CmdlineTuple.Dst),
    }

    if trans.timer != nil {
        trans.timer.Stop()
    }
    trans.timer = time.AfterFunc(TransactionTimeout, func() { trans.Expire() })

}

func (trans *RedisTransaction) Expire() {

    // remove from map
    delete(redisTransactionsMap, trans.tuple)
}

func receivedRedisResponse(msg *RedisMessage) {

    tuple := TcpTuple{
        Src_ip:    msg.Tuple.Src_ip,
        Dst_ip:    msg.Tuple.Dst_ip,
        Src_port:  msg.Tuple.Src_port,
        Dst_port:  msg.Tuple.Dst_port,
        stream_id: msg.Stream_id,
    }
    trans := redisTransactionsMap[tuple]
    if trans == nil {
        WARN("Response from unknown transaction. Ignoring.")
        return
    }
    // check if the request was received
    if len(trans.Redis) == 0 {
        WARN("Response from unknown transaction. Ignoring.")
        return

    }

    trans.Redis["response"] = msg.Message

    trans.Response_raw = msg.Message

    trans.ResponseTime = int32(msg.Ts.Sub(trans.ts).Nanoseconds() / 1e6) // resp_time in milliseconds

    err := Publisher.PublishRedisTransaction(trans)
    if err != nil {
        WARN("Publish failure: %s", err)
    }

    DEBUG("redis", "Redis transaction completed: %s", trans.Redis)

    // remove from map
    delete(redisTransactionsMap, trans.tuple)
    if trans.timer != nil {
        trans.timer.Stop()
    }

}
<commit_msg>Fixed the Redis parser handling of segmented data.<commit_after>package main

import (
    "bytes"
    "strconv"
    "strings"
    "time"

    "labix.org/v2/mgo/bson"
)

type RedisMessage struct {
    Ts            time.Time
    NumberOfBulks int64
    Bulks         []string

    Stream_id    uint32
    Tuple        *IpPortTuple
    CmdlineTuple *CmdlineTuple
    Direction    uint8

    IsRequest bool
    Message   string
}

type RedisStream struct {
    tcpStream *TcpStream

    data []byte

    parseOffset   int
    bytesReceived int

    message *RedisMessage
}

type RedisTransaction struct {
    Type         string
    tuple        TcpTuple
    Src          Endpoint
    Dst          Endpoint
    ResponseTime int32
    Ts           int64
    JsTs         time.Time
    ts           time.Time
    cmdline      *CmdlineTuple

    Redis bson.M

    Request_raw  string
    Response_raw string

    timer *time.Timer
}

// Keep sorted for future command addition
var RedisCommands = map[string]struct{}{
    "APPEND":           struct{}{},
    "AUTH":             struct{}{},
    "BGREWRITEAOF":     struct{}{},
    "BGSAVE":           struct{}{},
    "BITCOUNT":         struct{}{},
    "BITOP":            struct{}{},
    "BITPOS":           struct{}{},
    "BLPOP":            struct{}{},
    "BRPOP":            struct{}{},
    "BRPOPLPUSH":       struct{}{},
    "CLIENT GETNAME":   struct{}{},
    "CLIENT KILL":      struct{}{},
    "CLIENT LIST":      struct{}{},
    "CLIENT PAUSE":     struct{}{},
    "CLIENT SETNAME":   struct{}{},
    "CONFIG GET":       struct{}{},
    "CONFIG RESETSTAT": struct{}{},
    "CONFIG REWRITE":   struct{}{},
    "CONFIG SET":       struct{}{},
    "DBSIZE":           struct{}{},
    "DEBUG OBJECT":     struct{}{},
    "DEBUG SEGFAULT":   struct{}{},
    "DECR":             struct{}{},
    "DECRBY":           struct{}{},
    "DEL":              struct{}{},
    "DISCARD":          struct{}{},
    "DUMP":             struct{}{},
    "ECHO":             struct{}{},
    "EVAL":             struct{}{},
    "EVALSHA":          struct{}{},
    "EXEC":             struct{}{},
    "EXISTS":           struct{}{},
    "EXPIRE":           struct{}{},
    "EXPIREAT":         struct{}{},
    "FLUSHALL":         struct{}{},
    "FLUSHDB":          struct{}{},
    "GET":              struct{}{},
    "GETBIT":           struct{}{},
    "GETRANGE":         struct{}{},
    "GETSET":           struct{}{},
    "HDEL":             struct{}{},
    "HEXISTS":          struct{}{},
    "HGET":             struct{}{},
    "HGETALL":          struct{}{},
    "HINCRBY":          struct{}{},
    "HINCRBYFLOAT":     struct{}{},
    "HKEYS":            struct{}{},
    "HLEN":             struct{}{},
    "HMGET":            struct{}{},
    "HMSET":            struct{}{},
    "HSCAN":            struct{}{},
    "HSET":             struct{}{},
    "HSETINX":          struct{}{},
    "HVALS":            struct{}{},
    "INCR":             struct{}{},
    "INCRBY":           struct{}{},
    "INCRBYFLOAT":      struct{}{},
    "INFO":             struct{}{},
    "KEYS":             struct{}{},
    "LASTSAVE":         struct{}{},
    "LINDEX":           struct{}{},
    "LINSERT":          struct{}{},
    "LLEN":             struct{}{},
    "LPOP":             struct{}{},
    "LPUSH":            struct{}{},
    "LPUSHX":           struct{}{},
    "LRANGE":           struct{}{},
    "LREM":             struct{}{},
    "LSET":             struct{}{},
    "LTRIM":            struct{}{},
    "MGET":             struct{}{},
    "MIGRATE":          struct{}{},
    "MONITOR":          struct{}{},
    "MOVE":             struct{}{},
    "MSET":             struct{}{},
    "MSETNX":           struct{}{},
    "MULTI":            struct{}{},
    "OBJECT":           struct{}{},
    "PERSIST":          struct{}{},
    "PEXPIRE":          struct{}{},
    "PEXPIREAT":        struct{}{},
    "PFADD":            struct{}{},
    "PFCOUNT":          struct{}{},
    "PFMERGE":          struct{}{},
    "PING":             struct{}{},
    "PSETEX":           struct{}{},
    "PSUBSCRIBE":       struct{}{},
    "PTTL":             struct{}{},
    "PUBLISH":          struct{}{},
    "PUBSUB":           struct{}{},
    "PUNSUBSCRIBE":     struct{}{},
    "QUIT":             struct{}{},
    "RANDOMKEY":        struct{}{},
    "RENAME":           struct{}{},
    "RENAMENX":         struct{}{},
    "RESTORE":          struct{}{},
    "RPOP":             struct{}{},
    "RPOPLPUSH":        struct{}{},
    "RPUSH":            struct{}{},
    "RPUSHX":           struct{}{},
    "SADD":             struct{}{},
    "SAVE":             struct{}{},
    "SCAN":             struct{}{},
    "SCARD":            struct{}{},
    "SCRIPT EXISTS":    struct{}{},
    "SCRIPT FLUSH":     struct{}{},
    "SCRIPT KILL":      struct{}{},
    "SCRIPT LOAD":      struct{}{},
    "SDIFF":            struct{}{},
    "SDIFFSTORE":       struct{}{},
    "SELECT":           struct{}{},
    "SET":              struct{}{},
    "SETBIT":           struct{}{},
    "SETEX":            struct{}{},
    "SETNX":            struct{}{},
    "SETRANGE":         struct{}{},
    "SHUTDOWN":         struct{}{},
    "SINTER":           struct{}{},
    "SINTERSTORE":      struct{}{},
    "SISMEMBER":        struct{}{},
    "SLAVEOF":          struct{}{},
    "SLOWLOG":          struct{}{},
    "SMEMBERS":         struct{}{},
    "SMOVE":            struct{}{},
    "SORT":             struct{}{},
    "SPOP":             struct{}{},
    "SRANDMEMBER":      struct{}{},
    "SREM":             struct{}{},
    "SSCAN":            struct{}{},
    "STRLEN":           struct{}{},
    "SUBSCRIBE":        struct{}{},
    "SUNION":           struct{}{},
    "SUNIONSTORE":      struct{}{},
    "SYNC":             struct{}{},
    "TIME":             struct{}{},
    "TTL":              struct{}{},
    "TYPE":             struct{}{},
    "UNSUBSCRIBE":      struct{}{},
    "UNWATCH":          struct{}{},
    "WATCH":            struct{}{},
    "ZADD":             struct{}{},
    "ZCARD":            struct{}{},
    "ZCOUNT":           struct{}{},
    "ZINCRBY":          struct{}{},
    "ZINTERSTORE":      struct{}{},
    "ZRANGE":           struct{}{},
    "ZRANGEBYSCORE":    struct{}{},
    "ZRANK":            struct{}{},
    "ZREM":             struct{}{},
    "ZREMRANGEBYLEX":   struct{}{},
    "ZREMRANGEBYRANK":  struct{}{},
    "ZREMRANGEBYSCORE": struct{}{},
    "ZREVRANGE":        struct{}{},
    "ZREVRANGEBYSCORE": struct{}{},
    "ZREVRANK":         struct{}{},
    "ZSCAN":            struct{}{},
    "ZSCORE":           struct{}{},
    "ZUNIONSTORE":      struct{}{},
}

var redisTransactionsMap = make(map[TcpTuple]*RedisTransaction, TransactionsHashSize)

func (stream *RedisStream) PrepareForNewMessage() {
    stream.data = stream.data[stream.parseOffset:]
    stream.parseOffset = 0
    stream.message.NumberOfBulks = 0
    stream.message.Bulks = []string{}
    stream.message.IsRequest = false
}

func redisMessageParser(s *RedisStream) (bool, bool) {

    var err error
    var value string
    m := s.message

    for s.parseOffset < len(s.data) {

        if s.data[s.parseOffset] == '*' {
            //Multi Bulk Message

            found, line, off := readLine(s.data, s.parseOffset)
            if !found {
                DEBUG("redis", "End of line not found, waiting for more data")
                return true, false
            }

            if len(line) == 3 && line[1] == '-' && line[2] == '1' {
                //NULL Multi Bulk
                s.parseOffset = off
                value = "nil"
            } else {

                m.NumberOfBulks, err = strconv.ParseInt(line[1:], 10, 64)

                if err != nil {
                    ERR("Failed to read number of bulk messages: %s", err)
                    return false, false
                }
                s.parseOffset = off
                m.Bulks = []string{}

                continue
            }

        } else if s.data[s.parseOffset] == '$' {
            old_offset := s.parseOffset
            // Bulk Reply

            found, line, off := readLine(s.data, s.parseOffset)
            if !found {
                DEBUG("redis", "End of line not found, waiting for more data")
                s.parseOffset = old_offset
                return true, false
            }

            if len(line) == 3 && line[1] == '-' && line[2] == '1' {
                // NULL Bulk Reply
                value = "nil"
                s.parseOffset = off
            } else {
                length, err := strconv.ParseInt(line[1:], 10, 64)
                if err != nil {
                    ERR("Failed to read bulk message: %s", err)
                    return false, false
                }

                s.parseOffset = off

                found, line, off = readLine(s.data, s.parseOffset)
                if !found {
                    DEBUG("redis", "End of line not found, waiting for more data")
                    s.parseOffset = old_offset
                    return true, false
                }

                if int64(len(line)) != length {
                    ERR("Wrong length of data: %d instead of %d", len(line), length)
                    return false, false
                }
                value = line
                s.parseOffset = off
            }

        } else if s.data[s.parseOffset] == ':' {
            // Integer reply
            found, line, off := readLine(s.data, s.parseOffset)
            if !found {
                return true, false
            }
            n, err := strconv.ParseInt(line[1:], 10, 64)

            if err != nil {
                ERR("Failed to read integer reply: %s", err)
                return false, false
            }
            value = strconv.Itoa(int(n))
            s.parseOffset = off

        } else if s.data[s.parseOffset] == '+' {
            // Status Reply
            found, line, off := readLine(s.data, s.parseOffset)
            if !found {
                return true, false
            }

            value = line[1:]
            s.parseOffset = off
        } else if s.data[s.parseOffset] == '-' {
            // Error Reply
            found, line, off := readLine(s.data, s.parseOffset)
            if !found {
                return true, false
            }

            value = line[1:]
            s.parseOffset = off
        } else {
            DEBUG("redis", "Unexpected message starting with %s", s.data[s.parseOffset:])
            return false, false
        }

        // add value
        if m.NumberOfBulks > 0 {
            m.NumberOfBulks = m.NumberOfBulks - 1
            m.Bulks = append(m.Bulks, value)

            if len(m.Bulks) == 1 {
                // check if it's a command
                if isRedisCommand(value) {
                    m.IsRequest = true
                }
            }

            if m.NumberOfBulks == 0 {
                // the last bulk received
                m.Message = strings.Join(m.Bulks, " ")
                return true, true
            }
        } else {
            m.Message = value
            return true, true
        }

    }   //end for

    return true, false
}

func readLine(data []byte, offset int) (bool, string, int) {
    q := bytes.Index(data[offset:], []byte("\r\n"))
    if q == -1 {
        return false, "", 0
    }
    return true, string(data[offset : offset+q]), offset + q + 2
}

func ParseRedis(pkt *Packet, tcp *TcpStream, dir uint8) {
    defer RECOVER("ParseRedis exception")

    if tcp.redisData[dir] == nil {
        tcp.redisData[dir] = &RedisStream{
            tcpStream: tcp,
            data:      pkt.payload,
            message:   &RedisMessage{Ts: pkt.ts},
        }
    } else {
        // concatenate bytes
        tcp.redisData[dir].data = append(tcp.redisData[dir].data, pkt.payload...)
        if len(tcp.redisData[dir].data) > TCP_MAX_DATA_IN_STREAM {
            DEBUG("redis", "Stream data too large, dropping TCP stream")
            tcp.redisData[dir] = nil
            return
        }
    }

    stream := tcp.redisData[dir]
    for len(stream.data) > 0 {
        if stream.message == nil {
            stream.message = &RedisMessage{Ts: pkt.ts}
        }

        ok, complete := redisMessageParser(tcp.redisData[dir])

        if !ok {
            // drop this tcp stream. Will retry parsing with the next
            // segment in it
            tcp.redisData[dir] = nil
            DEBUG("redis", "Ignore Redis message. Drop tcp stream. Try parsing with the next segment")
            return
        }

        if complete {

            if stream.message.IsRequest {
                DEBUG("redis", "REDIS request message: %s", stream.message.Message)
            } else {
                DEBUG("redis", "REDIS response message: %s", stream.message.Message)
            }

            // all ok, go to next level
            handleRedis(stream.message, tcp, dir)

            // and reset message
            stream.PrepareForNewMessage()
        } else {
            // wait for more data
            break
        }
    }

}

func isRedisCommand(key string) bool {
    _, exists := RedisCommands[key]
    return exists
}

func handleRedis(m *RedisMessage, tcp *TcpStream,
    dir uint8) {

    m.Stream_id = tcp.id
    m.Tuple = tcp.tuple
    m.Direction = dir
    m.CmdlineTuple = procWatcher.FindProcessesTuple(tcp.tuple)

    if m.IsRequest {
        receivedRedisRequest(m)
    } else {
        receivedRedisResponse(m)
    }
}

func receivedRedisRequest(msg *RedisMessage) {
    // Add it to the HT
    tuple := TcpTuple{
        Src_ip:    msg.Tuple.Src_ip,
        Dst_ip:    msg.Tuple.Dst_ip,
        Src_port:  msg.Tuple.Src_port,
        Dst_port:  msg.Tuple.Dst_port,
        stream_id: msg.Stream_id,
    }

    trans := redisTransactionsMap[tuple]
    if trans != nil {
        if len(trans.Redis) != 0 {
            WARN("Two requests without a Response. Dropping old request")
        }
    } else {
        trans = &RedisTransaction{Type: "redis", tuple: tuple}
        redisTransactionsMap[tuple] = trans
    }

    trans.Redis = bson.M{
        "request": msg.Message,
    }
    trans.Request_raw = msg.Message

    trans.cmdline = msg.CmdlineTuple
    trans.ts = msg.Ts
    trans.Ts = int64(trans.ts.UnixNano() / 1000) // transactions have microseconds resolution
    trans.JsTs = msg.Ts
    trans.Src = Endpoint{
        Ip:   Ipv4_Ntoa(tuple.Src_ip),
        Port: tuple.Src_port,
        Proc: string(msg.CmdlineTuple.Src),
    }
    trans.Dst = Endpoint{
        Ip:   Ipv4_Ntoa(tuple.Dst_ip),
        Port: tuple.Dst_port,
        Proc: string(msg.CmdlineTuple.Dst),
    }

    if trans.timer != nil {
        trans.timer.Stop()
    }
    trans.timer = time.AfterFunc(TransactionTimeout, func() { trans.Expire() })

}

func (trans *RedisTransaction) Expire() {

    // remove from map
    delete(redisTransactionsMap, trans.tuple)
}

func receivedRedisResponse(msg *RedisMessage) {

    tuple := TcpTuple{
        Src_ip:    msg.Tuple.Src_ip,
        Dst_ip:    msg.Tuple.Dst_ip,
        Src_port:  msg.Tuple.Src_port,
        Dst_port:  msg.Tuple.Dst_port,
        stream_id: msg.Stream_id,
    }
    trans := redisTransactionsMap[tuple]
    if trans == nil {
        WARN("Response from unknown transaction. Ignoring.")
        return
    }
    // check if the request was received
    if len(trans.Redis) == 0 {
        WARN("Response from unknown transaction. Ignoring.")
        return

    }

    trans.Redis["response"] = msg.Message

    trans.Response_raw = msg.Message

    trans.ResponseTime = int32(msg.Ts.Sub(trans.ts).Nanoseconds() / 1e6) // resp_time in milliseconds

    err := Publisher.PublishRedisTransaction(trans)
    if err != nil {
        WARN("Publish failure: %s", err)
    }

    DEBUG("redis", "Redis transaction completed: %s", trans.Redis)

    // remove from map
    delete(redisTransactionsMap, trans.tuple)
    if trans.timer != nil {
        trans.timer.Stop()
    }

}
<|endoftext|>


<commit_before>package par2

import (
	"bufio"
	"bytes"
	"crypto/md5"
	"encoding/binary"
	"errors"
	"fmt"
	"hash"
	"io"
	"math/big"
	"os"
)

type Fileset struct {
	setID     [16]byte
	slicelen  uint64
	complete  bool
	files     map[[16]byte]*File
	checksums map[[16]byte]chksum
}

type File struct {
	Name      string
	md5       [16]byte
	length    uint64
	checksums [][16]byte
}

func (f *File) numBlocks(fset *Fileset) int {
	blockcount := int(f.length / fset.slicelen)
	if f.length%fset.slicelen != 0 {
		blockcount++
	}
	return blockcount
}

type chksum struct {
	*File
	blockno int
}

// NewFileset reads r and returns a Fileset that can be used for verification and recovery of the files.
func NewFileset(r io.Reader) (*Fileset, error) {
	fset := &Fileset{}
	fset.files = make(map[[16]byte]*File)
	fset.checksums = make(map[[16]byte]chksum)
	bufr := bufio.NewReader(r)
	for {
		hdr, err := readHeader(bufr)
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}
		if fset.setID == ([16]byte{}) {
			fset.setID = hdr.setID
		} else if hdr.setID != fset.setID {
			// this is weird and shouldn't happen
			return nil, errors.New("mismatched set ID in one file")
		}
		switch hdr.typ {
		case typeFileDesc:
			f, id, _ := readFileDesc(hdr, bufr)
			if f == nil {
				continue
			}
			fi, ok := fset.files[id]
			if ok {
				if fi.Name == "" {
					// file was discovered through some means
					// that didn't include the file info, usually IFSC
					// pkt. Fill in the information now that we know what file this is.
					f.checksums = fi.checksums
					fset.files[id] = f
				}
			} else {
				fset.files[id] = f
			}
		case typeIFSC:
			chksums, id, _ := readIFSC(hdr, bufr)
			if chksums == nil {
				continue
			}
			fi, ok := fset.files[id]
			if !ok {
				fi = new(File)
				fset.files[id] = fi
			}
			if fi.checksums == nil {
				fi.checksums = chksums
			}
			for i, chk := range fi.checksums {
				fset.checksums[chk] = chksum{
					File:    fi,
					blockno: i,
				}
			}
		case typeMain:
			slicelen, ids, _ := readMain(hdr, bufr)
			for _, id := range ids {
				if _, ok := fset.files[id]; !ok {
					fset.files[id] = new(File)
				}
			}
			fset.slicelen = slicelen
		default:
		}
	}
	if !fset.CanVerify() {
		return fset, nil
	}
	return fset, nil
}

// CanVerify returns whether the current fileset can be
// used for verification.
func (f *Fileset) CanVerify() bool {
	if f.complete {
		return true
	}
	if f.slicelen == 0 {
		return false
	}
	for _, file := range f.files {
		if file.Name == "" || file.checksums == nil {
			return false
		}
	}
	f.complete = true
	return true
}

// Verify verifies the files at paths against the fileset.
// It returns a list of matches and how many blocks are needed in order to repair.
func (f *Fileset) Verify(paths []string) ([]*FileMatch, int) {
	if !f.complete {
		return nil, 0
	}
	files := make(map[[16]byte]*File, len(f.files))
	for _, v := range f.files {
		files[v.md5] = v
	}
	matches := make([]*FileMatch, 0, len(paths))
	blocksNeeded := 0
	for _, s := range paths {
		fm, blocksmissing := f.verifyfile(files, s)
		if fm != nil && fm.File != nil {
			delete(files, fm.File.md5)
			blocksNeeded += blocksmissing
			matches = append(matches, fm)
		}
	}
	for _, fi := range files {
		matches = append(matches, &FileMatch{Err: ErrMissing, File: fi})
		blocksNeeded += fi.numBlocks(f)
	}
	return matches, blocksNeeded
}

var ErrMissing = errors.New("par2: file missing")

func (fset *Fileset) verifyfile(files map[[16]byte]*File, s string) (*FileMatch, int) {
	file, err := os.Open(s)
	if err != nil {
		return &FileMatch{Err: err}, 0
	}
	defer file.Close()

	chk := md5.New()
	_, err = io.Copy(chk, file)
	if err != nil {
		return &FileMatch{Err: err}, 0
	}
	var md5sum [16]byte
	chk.Sum(md5sum[:0])
	fi := files[md5sum]
	if fi != nil {
		return &FileMatch{Path: s, File: fi}, 0
	}

	// ok, this file is either corrupted or not part of the recovery set.
	// Since we're downloading these files via yenc (hopefully), we can always
	// assume that they will have holes where they were corrupted.
	_, err = file.Seek(0, os.SEEK_SET)
	if err != nil {
		return &FileMatch{Err: err}, 0
	}
	partialMatch := &FileMatch{}
	for {
		mdchk := md5.New()
		_, err := io.CopyN(mdchk, file, int64(fset.slicelen))
		mdchk.Sum(md5sum[:0])
		if f, ok := fset.checksums[md5sum]; ok {
			if partialMatch.File == nil {
				// ok we have a match, init the block bitmap
				partialMatch.blocks = &big.Int{}
				partialMatch.File = f.File
				partialMatch.Path = s
			} else if partialMatch.File != f.File {
				// shit
				continue
			}
			partialMatch.blocks.SetBit(partialMatch.blocks, f.blockno, 1)
		}
		if err != nil {
			break
		}
	}
	if partialMatch.File == nil {
		// not part of the recovery set.
		return nil, 0
	}
	blockcount := partialMatch.File.numBlocks(fset)
	blocksmissing := 0
	for i := 0; i < blockcount; i++ {
		if partialMatch.blocks.Bit(i) == 0 {
			blocksmissing++
		}
	}
	return partialMatch, blocksmissing
}

type FileMatch struct {
	Err    error
	Path   string
	File   *File
	blocks *big.Int
}

type hdr struct {
	length      uint64
	hash        [16]byte
	setID       [16]byte
	typ         typ
	partialhash hash.Hash
}

func (h *hdr) String() string {
	return fmt.Sprintf("{\n\t%d\n\t%x\n\t%x\n\t%s\n}", h.length, h.hash, h.setID, h.typ)
}

func readHeader(r *bufio.Reader) (h hdr, err error) {
	err = findHeader(r)
	if err != nil {
		return h, err
	}
	defer func() {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}
	}()
	var buf [56]byte
	b := buf[:]
	_, err = io.ReadFull(r, b)
	if err != nil {
		return h, err
	}
	h.length, b = readint(b)
	// for convenience, turn the length into
	// length after the header has been read.
	h.length -= 64
	h.hash, b = readmd5(b)

	// create a partial hash so that we can create a hash of the
	// entire packet.
	h.partialhash = md5.New()
	h.partialhash.Write(b)

	h.setID, b = readmd5(b)

	// i know, it's not an md5sum
	typbuf, b := readmd5(b)
	switch typbuf {
	case magicMain:
		h.typ = typeMain
	case magicFiledesc:
		h.typ = typeFileDesc
	case magicIFSC:
		h.typ = typeIFSC
	case magicRecvSlic:
		h.typ = typeRecvSlic
	case magicCreator:
		h.typ = typeCreator
	default:
		h.typ = typeUnknown
	}

	return h, nil
}

// These are all the required types.
// Going to leave it at that
var (
	magicMain     = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'M', 'a', 'i', 'n', 0, 0, 0, 0}
	magicFiledesc = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'F', 'i', 'l', 'e', 'D', 'e', 's', 'c'}
	magicIFSC     = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'I', 'F', 'S', 'C', 0, 0, 0, 0}
	magicRecvSlic = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'R', 'e', 'c', 'v', 'S', 'l', 'i', 'c'}
	magicCreator  = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'C', 'r', 'e', 'a', 't', 'o', 'r', 0}
)

// corruption might cause the header
// to become unaligned, so just search for the damn
// header.
func findHeader(r *bufio.Reader) error {
	i := 0
	str := "PAR2\x00PKT"
	for {
		//when completely matched
		if i == len(str) {
			return nil
		}
		c, err := r.ReadByte()
		if err != nil {
			return err
		}
		if str[i] == c {
			i++
			continue
		}
		i = 0
		if str[0] == c {
			i++
		}
	}
}

func readFileDesc(h hdr, r *bufio.Reader) (f *File, id [16]byte, err error) {
	buf, err := readPkt(h, r)
	if err != nil {
		return nil, id, nil
	}
	f = new(File)
	id, buf = readmd5(buf)
	f.md5, buf = readmd5(buf)
	_, buf = readmd5(buf)
	f.length, buf = readint(buf)

	// rest of block is name, trim 0 padding.
	i := bytes.LastIndex(buf, zero)
	if i != -1 {
		buf = buf[:i]
	}
	f.Name = string(buf)
	return f, id, nil
}

var zero = []byte{0}

func readIFSC(h hdr, r *bufio.Reader) (ss [][16]byte, id [16]byte, err error) {
	buf, err := readPkt(h, r)
	if err != nil {
		return nil, id, nil
	}
	id, buf = readmd5(buf)
	ss = make([][16]byte, 0, len(buf)/20)
	for len(buf) > 0 {
		var md5h [16]byte
		md5h, buf = readmd5(buf)
		// don't care about the crc, just that it gets consumed
		_, buf = readcrc(buf)
		ss = append(ss, md5h)
	}
	return ss, id, nil
}

func readMain(h hdr, r *bufio.Reader) (slicesize uint64, ids [][16]byte, err error) {
	buf, err := readPkt(h, r)
	if err != nil {
		return 0, nil, nil
	}
	slicesize, buf = readint(buf)
	numfiles, buf := readcrc(buf)
	ids = make([][16]byte, 0, numfiles)
	for i := uint32(0); i < numfiles; i++ {
		var nid [16]byte
		nid, buf = readmd5(buf)
		ids = append(ids, nid)
	}

	return slicesize, ids, nil
}

func readmd5(b []byte) ([16]byte, []byte) {
	var ret [16]byte
	copy(ret[:], b[:16])
	b = b[16:]
	return ret, b
}

func readint(b []byte) (uint64, []byte) {
	ret := binary.LittleEndian.Uint64(b)
	b = b[8:]
	return ret, b
}

func readcrc(b []byte) (uint32, []byte) {
	ret := binary.LittleEndian.Uint32(b)
	b = b[4:]
	return ret, b
}

// to be used for packets that can be read in a single go. Verifies
// the md5 sum of the packet as well
func readPkt(h hdr, r *bufio.Reader) ([]byte, error) {
	buf := make([]byte, h.length)
	_, err := io.ReadFull(r, buf)
	if err != nil {
		return nil, err
	}
	// check if the packet is valid
	h.partialhash.Write(buf)
	hash := h.partialhash.Sum(nil)
	for i, b := range h.hash {
		if hash[i] != b {
			return nil, errors.New("mismatch packet md5 and packet contents")
		}
	}
	return buf, nil
}
<commit_msg>remove file level md5<commit_after>package par2

import (
	"bufio"
	"bytes"
	"crypto/md5"
	"encoding/binary"
	"errors"
	"fmt"
	"hash"
	"io"
	"math/big"
	"os"
)

type Fileset struct {
	setID     [16]byte
	slicelen  uint64
	complete  bool
	files     map[[16]byte]*File
	checksums map[[16]byte]chksum
}

type File struct {
	Name      string
	length    uint64
	checksums [][16]byte
}

func (f *File) numBlocks(fset *Fileset) int {
	blockcount := int(f.length / fset.slicelen)
	if f.length%fset.slicelen != 0 {
		blockcount++
	}
	return blockcount
}

type chksum struct {
	*File
	blockno int
}

// NewFileset reads r and returns a Fileset that can be used for verification and recovery of the files.
func NewFileset(r io.Reader) (*Fileset, error) {
	fset := &Fileset{}
	fset.files = make(map[[16]byte]*File)
	fset.checksums = make(map[[16]byte]chksum)
	bufr := bufio.NewReader(r)
	for {
		hdr, err := readHeader(bufr)
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}
		if fset.setID == ([16]byte{}) {
			fset.setID = hdr.setID
		} else if hdr.setID != fset.setID {
			// this is weird and shouldn't happen
			return nil, errors.New("mismatched set ID in one file")
		}
		switch hdr.typ {
		case typeFileDesc:
			f, id, _ := readFileDesc(hdr, bufr)
			if f == nil {
				continue
			}
			fi, ok := fset.files[id]
			if ok {
				if fi.Name == "" {
					// file was discovered through some means
					// that didn't include the file info, usually IFSC
					// pkt. Fill in the information now that we know what file this is.
					f.checksums = fi.checksums
					fset.files[id] = f
				}
			} else {
				fset.files[id] = f
			}
		case typeIFSC:
			chksums, id, _ := readIFSC(hdr, bufr)
			if chksums == nil {
				continue
			}
			fi, ok := fset.files[id]
			if !ok {
				fi = new(File)
				fset.files[id] = fi
			}
			if fi.checksums == nil {
				fi.checksums = chksums
			}
			for i, chk := range fi.checksums {
				fset.checksums[chk] = chksum{
					File:    fi,
					blockno: i,
				}
			}
		case typeMain:
			slicelen, ids, _ := readMain(hdr, bufr)
			for _, id := range ids {
				if _, ok := fset.files[id]; !ok {
					fset.files[id] = new(File)
				}
			}
			fset.slicelen = slicelen
		default:
		}
	}
	if !fset.CanVerify() {
		return fset, nil
	}
	return fset, nil
}

// CanVerify returns whether the current fileset can be
// used for verification.
func (f *Fileset) CanVerify() bool {
	if f.complete {
		return true
	}
	if f.slicelen == 0 {
		return false
	}
	for _, file := range f.files {
		if file.Name == "" || file.checksums == nil {
			return false
		}
	}
	f.complete = true
	return true
}

// Verify verifies the files at paths against the fileset.
// It returns a list of matches and how many blocks are needed in order to repair.
func (f *Fileset) Verify(paths []string) ([]*FileMatch, int) {
	if !f.complete {
		return nil, 0
	}
	files := make(map[*File]struct{}, len(f.files))
	for _, v := range f.files {
		files[v] = struct{}{}
	}
	matches := make([]*FileMatch, 0, len(paths))
	blocksNeeded := 0
	for _, s := range paths {
		fm, blocksmissing := f.verifyfile(s)
		if fm != nil && fm.File != nil {
			delete(files, fm.File)
			blocksNeeded += blocksmissing
			matches = append(matches, fm)
		}
	}
	for fi := range files {
		matches = append(matches, &FileMatch{Err: ErrMissing, File: fi})
		blocksNeeded += fi.numBlocks(f)
	}
	return matches, blocksNeeded
}

var ErrMissing = errors.New("par2: file missing")

func (fset *Fileset) verifyfile(s string) (*FileMatch, int) {
	file, err := os.Open(s)
	if err != nil {
		return &FileMatch{Err: err}, 0
	}
	defer file.Close()

	match := &FileMatch{}
	for {
		mdchk := md5.New()
		n, err := io.CopyN(mdchk, file, int64(fset.slicelen))
		if n == 0 {
			break
		}
		if uint64(n) < fset.slicelen {
			// we have a partial block. par2 spec says that we should
			// fill the remainder with 0s
			// Ugh.
			for i := uint64(0); i < fset.slicelen-uint64(n); i++ {
				// can't fail.
				mdchk.Write(zero)
			}
		}
		var md5sum [16]byte
		mdchk.Sum(md5sum[:0])
		if f, ok := fset.checksums[md5sum]; ok {
			if match.File == nil {
				// ok we have a match, init the block bitmap
				match.blocks = &big.Int{}
				match.File = f.File
				match.Path = s
			} else if match.File != f.File {
				// we already decided on one file and now we have
				// another file with the same block.
				//
				// Effort.
				continue
			}
			match.blocks.SetBit(match.blocks, f.blockno, 1)
		}
		if err != nil {
			break
		}
	}
	if match.File == nil {
		// not part of the recovery set.
		return nil, 0
	}
	blockcount := match.File.numBlocks(fset)
	blocksmissing := 0
	for i := 0; i < blockcount; i++ {
		if match.blocks.Bit(i) == 0 {
			blocksmissing++
		}
	}
	return match, blocksmissing
}

type FileMatch struct {
	Err    error
	Path   string
	File   *File
	blocks *big.Int
}

type hdr struct {
	length      uint64
	hash        [16]byte
	setID       [16]byte
	typ         typ
	partialhash hash.Hash
}

func (h *hdr) String() string {
	return fmt.Sprintf("{\n\t%d\n\t%x\n\t%x\n\t%s\n}", h.length, h.hash, h.setID, h.typ)
}

func readHeader(r *bufio.Reader) (h hdr, err error) {
	err = findHeader(r)
	if err != nil {
		return h, err
	}
	defer func() {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}
	}()
	var buf [56]byte
	b := buf[:]
	_, err = io.ReadFull(r, b)
	if err != nil {
		return h, err
	}
	h.length, b = readint(b)
	// for convenience, turn the length into
	// length after the header has been read.
	h.length -= 64
	h.hash, b = readmd5(b)

	// create a partial hash so that we can create a hash of the
	// entire packet.
	h.partialhash = md5.New()
	h.partialhash.Write(b)

	h.setID, b = readmd5(b)

	// i know, it's not an md5sum
	typbuf, b := readmd5(b)
	switch typbuf {
	case magicMain:
		h.typ = typeMain
	case magicFiledesc:
		h.typ = typeFileDesc
	case magicIFSC:
		h.typ = typeIFSC
	case magicRecvSlic:
		h.typ = typeRecvSlic
	case magicCreator:
		h.typ = typeCreator
	default:
		h.typ = typeUnknown
	}

	return h, nil
}

// These are all the required types.
// Going to leave it at that
var (
	magicMain     = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'M', 'a', 'i', 'n', 0, 0, 0, 0}
	magicFiledesc = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'F', 'i', 'l', 'e', 'D', 'e', 's', 'c'}
	magicIFSC     = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'I', 'F', 'S', 'C', 0, 0, 0, 0}
	magicRecvSlic = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'R', 'e', 'c', 'v', 'S', 'l', 'i', 'c'}
	magicCreator  = [16]byte{'P', 'A', 'R', ' ', '2', '.', '0', 0, 'C', 'r', 'e', 'a', 't', 'o', 'r', 0}
)

// corruption might cause the header
// to become unaligned, so just search for the damn
// header.
func findHeader(r *bufio.Reader) error {
	i := 0
	str := "PAR2\x00PKT"
	for {
		//when completely matched
		if i == len(str) {
			return nil
		}
		c, err := r.ReadByte()
		if err != nil {
			return err
		}
		if str[i] == c {
			i++
			continue
		}
		i = 0
		if str[0] == c {
			i++
		}
	}
}

func readFileDesc(h hdr, r *bufio.Reader) (f *File, id [16]byte, err error) {
	buf, err := readPkt(h, r)
	if err != nil {
		return nil, id, nil
	}
	f = new(File)
	id, buf = readmd5(buf)
	_, buf = readmd5(buf)
	_, buf = readmd5(buf)
	f.length, buf = readint(buf)

	// rest of block is name, trim 0 padding.
	i := bytes.LastIndex(buf, zero)
	if i != -1 {
		buf = buf[:i]
	}
	f.Name = string(buf)
	return f, id, nil
}

var zero = []byte{0}

func readIFSC(h hdr, r *bufio.Reader) (ss [][16]byte, id [16]byte, err error) {
	buf, err := readPkt(h, r)
	if err != nil {
		return nil, id, nil
	}
	id, buf = readmd5(buf)
	ss = make([][16]byte, 0, len(buf)/20)
	for len(buf) > 0 {
		var md5h [16]byte
		md5h, buf = readmd5(buf)
		// don't care about the crc, just that it gets consumed
		_, buf = readcrc(buf)
		ss = append(ss, md5h)
	}
	return ss, id, nil
}

func readMain(h hdr, r *bufio.Reader) (slicesize uint64, ids [][16]byte, err error) {
	buf, err := readPkt(h, r)
	if err != nil {
		return 0, nil, nil
	}
	slicesize, buf = readint(buf)
	numfiles, buf := readcrc(buf)
	ids = make([][16]byte, 0, numfiles)
	for i := uint32(0); i < numfiles; i++ {
		var nid [16]byte
		nid, buf = readmd5(buf)
		ids = append(ids, nid)
	}

	return slicesize, ids, nil
}

func readmd5(b []byte) ([16]byte, []byte) {
	var ret [16]byte
	copy(ret[:], b[:16])
	b = b[16:]
	return ret, b
}

func readint(b []byte) (uint64, []byte) {
	ret := binary.LittleEndian.Uint64(b)
	b = b[8:]
	return ret, b
}

func readcrc(b []byte) (uint32, []byte) {
	ret := binary.LittleEndian.Uint32(b)
	b = b[4:]
	return ret, b
}

// to be used for packets that can be read in a single go. Verifies
// the md5 sum of the packet as well
func readPkt(h hdr, r *bufio.Reader) ([]byte, error) {
	buf := make([]byte, h.length)
	_, err := io.ReadFull(r, buf)
	if err != nil {
		return nil, err
	}
	// check if the packet is valid
	h.partialhash.Write(buf)
	hash := h.partialhash.Sum(nil)
	for i, b := range h.hash {
		if hash[i] != b {
			return nil, errors.New("mismatch packet md5 and packet contents")
		}
	}
	return buf, nil
}
<|endoftext|>


<commit_before>// go-rst - A reStructuredText parser for Go
// 2014 (c) The go-rst Authors
// MIT Licensed. See LICENSE for details.
package parse

type itemElements int

const (
	itemEOF itemElements = iota
	itemBodyParagraph
	itemBodyText
	itemBodyInlineEmphasis
	itemBodyInlineStrong
	itemBodyInlineInterpreted
	itemBodyInlineInternalTarget
	itemBodyInlineFootnote
	itemBodyInlineHyperLink
	itemBodyBulletList
	itemBodyEnumeratedList
	itemBodyDefinitionList
	itemBodyFieldList
	itemBodyBibliographicField
	itemBodyOptionList
	itemBodyLiteralBlock
	itemBodyLiteralBlockIndented
	itemBodyQuotedBlock
)

type stateFn func(*lexer) stateFn

type item struct {
	typ itemElements
	pos Pos
	val string
}

type lexer struct {
	name    string
	input   string
	state   stateFn
	pos     Pos
	start   Pos
	width   Pos
	lastPos Pos
	items   chan item
}

func lex(name, input string) *lexer {
	l := &lexer{
		name:  name,
		input: input,
		items: make(chan item),
	}
	go l.run()
	return l
}

func (l *lexer) run() {
	for l.state = lexText; l.state != nil; {
		l.state = l.state(l)
	}
}

func lexText(l *lexer) stateFn {
	// for {
	// if strings.HasPrefix(l.input[l.pos:], l.leftDelim) {
	// if l.pos > l.start {
	// l.emit(itemText)
	// }
	// return lexLeftDelim
	// }
	// if l.next() == eof {
	// break
	// }
	// }
	// // Correctly reached EOF.
	// if l.pos > l.start {
	// l.emit(itemText)
	// }
	// l.emit(itemEOF)
	return nil
}
<commit_msg>lex.go: Implement prelim tokenizing for sections<commit_after>// go-rst - A reStructuredText parser for Go
// 2014 (c) The go-rst Authors
// MIT Licensed. See LICENSE for details.
package parse

import (
	// "strings"
	// "unicode"
	"fmt"
	"unicode/utf8"
	"github.com/demizer/go-elog"
)

type itemElement int

const (
	itemEOF itemElement = iota
	itemError
	itemTitle
	itemSectionAdornment
	itemParagraph
)

var sectionAdornments = []rune{'!', '"', '#', '$', '\'', '%', '&', '(', ')', '*',
	'+', ',', '-', '.', '/', ':', ';', '<', '=', '>', '?', '@', '[', '\\',
	']', '^', '_', '`', '{', '|', '}', '~'}

func isSectionAdornment(r rune) bool {
	for _, a := range sectionAdornments {
		if a == r {
			return true
		}
	}
	return false
}

const eof = -1

type stateFn func(*lexer) stateFn

type item struct {
	typ itemElement
	pos Pos
	val string
}

func (i item) String() string {
	switch {
	case i.typ == itemEOF:
		return "EOF"
	case i.typ == itemError:
		return i.val
	/* case len(i.val) > 10:
		return fmt.Sprintf("%.10q...", i.val) */
	}
	return fmt.Sprintf("%q", i.val)
}

type lexer struct {
	name       string
	input      string
	state      stateFn
	pos        Pos
	start      Pos
	width      Pos
	lastPos    Pos
	items      chan item
}

func lex(name, input string) *lexer {
	l := &lexer{
		name:  name,
		input: input,
		items: make(chan item),
	}
	go l.run()
	return l
}

// emit passes an item back to the client.
func (l *lexer) emit(t itemElement) {
	l.items <- item{t, l.start, l.input[l.start:l.pos]}
	l.start = l.pos

}

func (l *lexer) backup() {
	l.pos -= l.width
}

func (l *lexer) current() rune {
	r, _ := utf8.DecodeRuneInString(l.input[l.pos:])
	return r
}

func (l *lexer) peek() rune {
	r := l.next()
	l.backup()
	return r
}

func (l *lexer) ignore() {
	l.pos += 1
	l.start = l.pos
}

// next returns the next rune in the input.
func (l *lexer) next() rune {
	if int(l.pos) >= len(l.input) {
		l.width = 0
		return eof
	}
	r, w := utf8.DecodeRuneInString(l.input[l.pos:])
	l.width = Pos(w)
	l.pos += l.width
	return r
}

// nextItem returns the next item from the input.
func (l *lexer) nextItem() item {
	item := <-l.items
	l.lastPos = item.pos
	return item

}

func (l *lexer) run() {
	for l.state = lexParagraph; l.state != nil; {
		l.state = l.state(l)
	}
}

// isSpace reports whether r is a space character.
func isSpace(r rune) bool {
	return r == ' ' || r == '\t'
}

// isEndOfLine reports whether r is an end-of-line character.
func isEndOfLine(r rune) bool {
	return r == '\r' || r == '\n'
}

func lexParagraph(l *lexer) stateFn {
	for {
		if len(l.input) == 0 {
			log.Debugln("\tEmit EOF!")
			l.emit(itemEOF)
			return nil
		}
		log.Debugf("\tlexParagraph: %q, Start: %d, Pos: %d\n",
			l.input[l.start:l.pos], l.start, l.pos)
		if isEndOfLine(l.current()) || isEndOfLine(l.peek()) {
			log.Debugln("\tFound newline!")
			if isSectionAdornment(rune(l.input[l.pos+2])) {
				log.Debugln("Transition lexSection...")
				return lexSection
			}
			l.ignore()
		}
		if l.next() == eof {
			break
		}
	}
	// Correctly reached EOF.
	if l.pos > l.start {
		log.Debugln("\tEmit Paragraph!")
		l.emit(itemParagraph)
	}
	l.emit(itemEOF)
	return nil
}

func lexSection(l *lexer) stateFn {
	var titleWidth, sectionWidth Pos
	if len(l.input) > 0 {
		log.Debugf("\tlexSection: %q, Pos: %d, titleWidth: %d, sectionWidth: %d\n",
			l.input[l.start:l.pos], l.pos, titleWidth, sectionWidth)
	}
	if isEndOfLine(l.peek()) {
		titleWidth = l.pos - l.start
		log.Debugf("\tEmit itemTitle: %d-%d\n", l.start, l.pos)
		l.emit(itemTitle)
		l.ignore()
	}
Loop:
	for {
		// log.Debugln(sectionWidth, titleWidth)
		switch r := l.next(); {
		case isSectionAdornment(r):
			if len(l.input) > 0 {
				log.Debugf("\tlexSection: %q, Pos: %d, " +
				           "titleWidth: %d, sectionWidth: %d\n",
					   l.input[l.start:l.pos], l.pos,
					   titleWidth, sectionWidth)
			}
			sectionWidth += 1
		case isEndOfLine(r):
			if sectionWidth == titleWidth {
				l.backup()
				log.Debugf("\tEmit itemSectionAdornment: " +
					   "%d-%d\n", l.start, l.pos)
				l.emit(itemSectionAdornment)
				l.ignore()
				break Loop
			}
		}
	}
	log.Debugln("Transition lexParagraph...")
	return lexParagraph
}
<|endoftext|>