<commit_before>// Use cgo to interface with nflog
//
// Docs: http://www.netfilter.org/projects/libnetfilter_log/doxygen/index.html
//
// Debian packages needed:
//   apt-get install iptables-dev linux-libc-dev libnetfilter-log-dev

package main

import (
	"log"
	"net"
	"reflect"
	"syscall"
	"unsafe"
)

/*
#cgo LDFLAGS: -lnetfilter_log
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <stdlib.h>
#include <libnetfilter_log/libnetfilter_log.h>
#include <inttypes.h>

// A record of each packet
typedef struct {
	char *payload;
	int payload_len;
	u_int32_t seq;
} packet;

// Max number of packets to collect at once
#define MAX_PACKETS (16*1024)

// A load of packets with count
typedef struct {
	int index;
	packet pkt[MAX_PACKETS];
} packets;

// Process the incoming packet putting pointers to the data to be handled by Go
static int _processPacket(struct nflog_g_handle *gh, struct nfgenmsg *nfmsg, struct nflog_data *nfd, void *data) {
	packets *ps = (packets *)data;
	if (ps->index >= MAX_PACKETS) {
		return 1;
	}
	packet *p = &ps->pkt[ps->index++];
	p->payload = 0;
	p->payload_len = nflog_get_payload(nfd, &p->payload);
	p->seq = 0;
	nflog_get_seq(nfd, &p->seq);
	return 0;
 }

// Register the callback - can't be done from Go
//
// We have to register a C function _processPacket
static int _callback_register(struct nflog_g_handle *gh, packets *data) {
	return nflog_callback_register(gh, _processPacket, data);
}
*/
import "C"

const (
	RecvBufferSize  = 4 * 1024 * 1024
	NflogBufferSize = 4 * 1024 * 1024
	MaxQueueLogs    = C.MAX_PACKETS - 1
)

// NfLog
type NfLog struct {
	// Main nflog_handle
	h *C.struct_nflog_handle
	// File descriptor for socket operations
	fd C.int
	// Group handle
	gh *C.struct_nflog_g_handle
	// The multicast address
	McastGroup int
	// The next expected sequence number
	seq uint32
	// Errors
	errors int64
	// Flavour of IP we are expecting, 4 or 6
	IpVersion byte
	// Mask for the IP
	Mask net.IPMask
	// Whether to apply the mask or not
	UseMask bool
	// Are we account the source or the destination address
	Direction IpDirection
	// Flavour of IP packet we are decoding
	IpPacket *IpPacketInfo
	// Accounting
	a *Accounting
	// Quit the loop
	quit chan struct{}
	// Buffer for accumulated packets
	addPackets []AddPacket
	// Index of this in nflogs - uses to pass to C instead of a Go pointer
	index int
	// Pointer to the packets
	packets *C.packets
}

// An array of NfLog pointers
type NfLogs [16]*NfLog

// A global array of all NfLogs in use
var nflogs NfLogs

// Count all active NfLogs
func (ns NfLogs) Count() int {
	active := 0
	for _, nflog := range ns {
		if nflog != nil {
			active++
		}
	}
	return active
}

// Stop all active NfLogs
func (ns NfLogs) Stop() {
	for _, nflog := range ns {
		if nflog != nil {
			nflog.Close()
		}
	}
}

// Create a new NfLog
//
// McastGroup is that specified in ip[6]tables
// IPv6 is a flag to say if it is IPv6 or not
// Direction is to monitor the source address or the dest address
func NewNfLog(McastGroup int, IpVersion byte, Direction IpDirection, MaskBits int, a *Accounting) *NfLog {
	h := C.nflog_open()
	if h == nil {
		log.Fatalf("Failed to open NFLOG: %s", strerror())
	}
	if *Debug {
		log.Println("Binding nfnetlink_log to AF_INET")
	}
	if C.nflog_bind_pf(h, C.AF_INET) < 0 {
		log.Fatalf("nflog_bind_pf failed: %s", strerror())
	}

	nflog := &NfLog{
		h:          h,
		fd:         C.nflog_fd(h),
		McastGroup: McastGroup,
		IpVersion:  IpVersion,
		Direction:  Direction,
		a:          a,
		quit:       make(chan struct{}),
		packets:    (*C.packets)(C.malloc(C.sizeof_packets)),
	}
	for i := range nflogs {
		if nflogs[i] == nil {
			nflog.index = i
			nflogs[i] = nflog
			goto found
		}
	}
	log.Fatal("Too many filters")
found:
	switch IpVersion {
	case 4:
		nflog.IpPacket = Ip4Packet
	case 6:
		nflog.IpPacket = Ip6Packet
	default:
		log.Fatalf("Bad IP version %d", IpVersion)
	}
	addrBits := 8 * nflog.IpPacket.AddrLen
	nflog.UseMask = MaskBits < addrBits
	nflog.Mask = net.CIDRMask(MaskBits, addrBits)
	nflog.makeGroup(McastGroup, nflog.IpPacket.HeaderSize)
	// Start the background process
	go nflog.Loop()
	return nflog
}

// Process a packet
func (nflog *NfLog) ProcessPacket(packet []byte, seq uint32) AddPacket {
	// log.Printf("Packet %#v", packet)
	// Peek the IP Version out of the header
	ip_version := packet[IpVersion] >> IpVersionShift & IpVersionMask
	// log.Printf("Received %d: size %d, IPv%d", seq, payload_len, ip_version)
	if seq != 0 && seq != nflog.seq {
		nflog.errors++
		log.Printf("%d missing packets detected, %d to %d", seq-nflog.seq, seq, nflog.seq)
	}
	nflog.seq = seq + 1
	if ip_version != nflog.IpVersion {
		nflog.errors++
		log.Printf("Bad IP version: %d", ip_version)
		return AddPacket{Length: -1}
	}
	i := nflog.IpPacket
	if len(packet) < i.HeaderSize {
		nflog.errors++
		log.Printf("Short IPv%d packet %d/%d bytes", ip_version, len(packet), i.HeaderSize)
		return AddPacket{Length: -1}
	}

	var addr net.IP
	if nflog.Direction {
		addr = i.Src(packet)
	} else {
		addr = i.Dst(packet)
	}

	// Mask the address
	if nflog.UseMask {
		addr = addr.Mask(nflog.Mask)
	}

	return AddPacket{
		Direction: nflog.Direction,
		Addr:      string(addr),
		Length:    i.Length(packet),
	}
}

// Receive data from nflog stored in nflog.packets
func (nflog *NfLog) processPackets(addPackets []AddPacket) []AddPacket {
	n := int(nflog.packets.index)
	if n >= C.MAX_PACKETS {
		log.Printf("Packets buffer overflowed")
	}

	var packet []byte
	sliceHeader := (*reflect.SliceHeader)((unsafe.Pointer(&packet)))

	for i := 0; i < n; i++ {
		p := &nflog.packets.pkt[i]

		// Get the packet into a []byte
		// NB if the C data goes away then BAD things will happen!
		// So don't keep slices from this after returning from this function
		sliceHeader.Cap = int(p.payload_len)
		sliceHeader.Len = int(p.payload_len)
		sliceHeader.Data = uintptr(unsafe.Pointer(p.payload))

		// Process the packet
		newAddPacket := nflog.ProcessPacket(packet, uint32(p.seq))
		if newAddPacket.Length >= 0 {
			addPackets = append(addPackets, newAddPacket)
		}
	}
	sliceHeader = nil
	packet = nil
	return addPackets
}

// Current nflog error
func strerror() error {
	return syscall.Errno(C.nflog_errno)
}

// Connects to the group specified with the size
func (nflog *NfLog) makeGroup(group, size int) {
	if *Debug {
		log.Printf("Binding this socket to group %d", group)
	}
	gh := C.nflog_bind_group(nflog.h, (C.u_int16_t)(group))
	if gh == nil {
		log.Fatalf("nflog_bind_group failed: %s", strerror())
	}
	nflog.gh = gh

	// Set the maximum amount of logs in buffer for this group
	if C.nflog_set_qthresh(gh, MaxQueueLogs) < 0 {
		log.Fatalf("nflog_set_qthresh failed: %s", strerror())
	}

	// Set local sequence numbering to detect missing packets
	if C.nflog_set_flags(gh, C.NFULNL_CFG_F_SEQ) < 0 {
		log.Fatalf("nflog_set_flags failed: %s", strerror())
	}

	// Set buffer size large
	if C.nflog_set_nlbufsiz(gh, NflogBufferSize) < 0 {
		log.Fatalf("nflog_set_nlbufsiz: %s", strerror())
	}

	// Set timeout
	// Doesn't seem to make any difference and don't know the unit
	// if C.nflog_set_timeout(gh, NflogTimeout) < 0 {
	// 	log.Fatalf("nflog_set_timeout: %s", strerror())
	// }

	if *Debug {
		log.Printf("Setting copy_packet mode to %d bytes", size)
	}
	if C.nflog_set_mode(gh, C.NFULNL_COPY_PACKET, (C.uint)(size)) < 0 {
		log.Fatalf("nflog_set_mode failed: %s", strerror())
	}

	// Register the callback now we are set up
	//
	// Note that we pass an index into an array, not a pointer to
	// the nflog - it isn't a good idea for C to hold pointers to
	// go objects which might move
	C._callback_register(gh, nflog.packets)
}

// Receive packets in a loop until quit
func (nflog *NfLog) Loop() {
	buflen := C.size_t(RecvBufferSize)
	pbuf := C.malloc(buflen)
	if pbuf == nil {
		log.Fatal("No memory for malloc")
	}
	defer C.free(pbuf)
	for {
		nr := C.recv(nflog.fd, pbuf, buflen, 0)
		select {
		case <-nflog.quit:
			return
		default:
		}
		if nr < 0 {
			log.Printf("Recvfrom failed: %s", strerror())
			nflog.errors++
		} else {
			// Handle messages in packet reusing memory
			ps := <-nflog.a.returnAddPackets
			nflog.packets.index = 0
			C.nflog_handle_packet(nflog.h, (*C.char)(pbuf), (C.int)(nr))
			nflog.a.processAddPackets <- nflog.processPackets(ps[:0])
		}
	}

}

// Close the NfLog down
func (nflog *NfLog) Close() {
	close(nflog.quit)
	// Sometimes hangs and doesn't seem to be necessary
	// if *Debug {
	// 	log.Printf("Unbinding socket %d from group %d", nflog.fd, nflog.McastGroup)
	// }
	// if C.nflog_unbind_group(nflog.gh) < 0 {
	// 	log.Printf("nflog_unbind_group(%d) failed: %s", nflog.McastGroup, strerror())
	// }
	if *Debug {
		log.Printf("Closing nflog socket %d group %d", nflog.fd, nflog.McastGroup)
	}
	if C.nflog_close(nflog.h) < 0 {
		log.Printf("nflog_close failed: %s", strerror())
	}
	// Mark this index as no longer in use
	nflogs[nflog.index] = nil
	C.free(unsafe.Pointer(nflog.packets))
}
<commit_msg>Reduce NflogBufferSize to below 128k so it works for massive speedup<commit_after>// Use cgo to interface with nflog
//
// Docs: http://www.netfilter.org/projects/libnetfilter_log/doxygen/index.html
//
// Debian packages needed:
//   apt-get install iptables-dev linux-libc-dev libnetfilter-log-dev

package main

import (
	"log"
	"net"
	"reflect"
	"syscall"
	"unsafe"
)

/*
#cgo LDFLAGS: -lnetfilter_log
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <stdlib.h>
#include <libnetfilter_log/libnetfilter_log.h>
#include <inttypes.h>

// A record of each packet
typedef struct {
	char *payload;
	int payload_len;
	u_int32_t seq;
} packet;

// Max number of packets to collect at once
#define MAX_PACKETS (16*1024)

// A load of packets with count
typedef struct {
	int index;
	packet pkt[MAX_PACKETS];
} packets;

// Process the incoming packet putting pointers to the data to be handled by Go
static int _processPacket(struct nflog_g_handle *gh, struct nfgenmsg *nfmsg, struct nflog_data *nfd, void *data) {
	packets *ps = (packets *)data;
	if (ps->index >= MAX_PACKETS) {
		return 1;
	}
	packet *p = &ps->pkt[ps->index++];
	p->payload = 0;
	p->payload_len = nflog_get_payload(nfd, &p->payload);
	p->seq = 0;
	nflog_get_seq(nfd, &p->seq);
	return 0;
 }

// Register the callback - can't be done from Go
//
// We have to register a C function _processPacket
static int _callback_register(struct nflog_g_handle *gh, packets *data) {
	return nflog_callback_register(gh, _processPacket, data);
}
*/
import "C"

const (
	RecvBufferSize  = 4 * 1024 * 1024
	NflogBufferSize = 128 * 1024 // Must be <= 128k (checked in kernel source)
	NflogTimeout    = 100        // Timeout before sending data in 1/100th second
	MaxQueueLogs    = C.MAX_PACKETS - 1
)

// NfLog
type NfLog struct {
	// Main nflog_handle
	h *C.struct_nflog_handle
	// File descriptor for socket operations
	fd C.int
	// Group handle
	gh *C.struct_nflog_g_handle
	// The multicast address
	McastGroup int
	// The next expected sequence number
	seq uint32
	// Errors
	errors int64
	// Flavour of IP we are expecting, 4 or 6
	IpVersion byte
	// Mask for the IP
	Mask net.IPMask
	// Whether to apply the mask or not
	UseMask bool
	// Are we account the source or the destination address
	Direction IpDirection
	// Flavour of IP packet we are decoding
	IpPacket *IpPacketInfo
	// Accounting
	a *Accounting
	// Quit the loop
	quit chan struct{}
	// Buffer for accumulated packets
	addPackets []AddPacket
	// Index of this in nflogs - uses to pass to C instead of a Go pointer
	index int
	// Pointer to the packets
	packets *C.packets
}

// An array of NfLog pointers
type NfLogs [16]*NfLog

// A global array of all NfLogs in use
var nflogs NfLogs

// Count all active NfLogs
func (ns NfLogs) Count() int {
	active := 0
	for _, nflog := range ns {
		if nflog != nil {
			active++
		}
	}
	return active
}

// Stop all active NfLogs
func (ns NfLogs) Stop() {
	for _, nflog := range ns {
		if nflog != nil {
			nflog.Close()
		}
	}
}

// Create a new NfLog
//
// McastGroup is that specified in ip[6]tables
// IPv6 is a flag to say if it is IPv6 or not
// Direction is to monitor the source address or the dest address
func NewNfLog(McastGroup int, IpVersion byte, Direction IpDirection, MaskBits int, a *Accounting) *NfLog {
	h := C.nflog_open()
	if h == nil {
		log.Fatalf("Failed to open NFLOG: %s", strerror())
	}
	if *Debug {
		log.Println("Binding nfnetlink_log to AF_INET")
	}
	if C.nflog_bind_pf(h, C.AF_INET) < 0 {
		log.Fatalf("nflog_bind_pf failed: %s", strerror())
	}

	nflog := &NfLog{
		h:          h,
		fd:         C.nflog_fd(h),
		McastGroup: McastGroup,
		IpVersion:  IpVersion,
		Direction:  Direction,
		a:          a,
		quit:       make(chan struct{}),
		packets:    (*C.packets)(C.malloc(C.sizeof_packets)),
	}
	for i := range nflogs {
		if nflogs[i] == nil {
			nflog.index = i
			nflogs[i] = nflog
			goto found
		}
	}
	log.Fatal("Too many filters")
found:
	switch IpVersion {
	case 4:
		nflog.IpPacket = Ip4Packet
	case 6:
		nflog.IpPacket = Ip6Packet
	default:
		log.Fatalf("Bad IP version %d", IpVersion)
	}
	addrBits := 8 * nflog.IpPacket.AddrLen
	nflog.UseMask = MaskBits < addrBits
	nflog.Mask = net.CIDRMask(MaskBits, addrBits)
	nflog.makeGroup(McastGroup, nflog.IpPacket.HeaderSize)
	// Start the background process
	go nflog.Loop()
	return nflog
}

// Process a packet
func (nflog *NfLog) ProcessPacket(packet []byte, seq uint32) AddPacket {
	// log.Printf("Packet %#v", packet)
	// Peek the IP Version out of the header
	ip_version := packet[IpVersion] >> IpVersionShift & IpVersionMask
	// log.Printf("Received %d: size %d, IPv%d", seq, payload_len, ip_version)
	if seq != 0 && seq != nflog.seq {
		nflog.errors++
		log.Printf("%d missing packets detected, %d to %d", seq-nflog.seq, seq, nflog.seq)
	}
	nflog.seq = seq + 1
	if ip_version != nflog.IpVersion {
		nflog.errors++
		log.Printf("Bad IP version: %d", ip_version)
		return AddPacket{Length: -1}
	}
	i := nflog.IpPacket
	if len(packet) < i.HeaderSize {
		nflog.errors++
		log.Printf("Short IPv%d packet %d/%d bytes", ip_version, len(packet), i.HeaderSize)
		return AddPacket{Length: -1}
	}

	var addr net.IP
	if nflog.Direction {
		addr = i.Src(packet)
	} else {
		addr = i.Dst(packet)
	}

	// Mask the address
	if nflog.UseMask {
		addr = addr.Mask(nflog.Mask)
	}

	return AddPacket{
		Direction: nflog.Direction,
		Addr:      string(addr),
		Length:    i.Length(packet),
	}
}

// Receive data from nflog stored in nflog.packets
func (nflog *NfLog) processPackets(addPackets []AddPacket) []AddPacket {
	n := int(nflog.packets.index)
	if n >= C.MAX_PACKETS {
		log.Printf("Packets buffer overflowed")
	}

	var packet []byte
	sliceHeader := (*reflect.SliceHeader)((unsafe.Pointer(&packet)))

	for i := 0; i < n; i++ {
		p := &nflog.packets.pkt[i]

		// Get the packet into a []byte
		// NB if the C data goes away then BAD things will happen!
		// So don't keep slices from this after returning from this function
		sliceHeader.Cap = int(p.payload_len)
		sliceHeader.Len = int(p.payload_len)
		sliceHeader.Data = uintptr(unsafe.Pointer(p.payload))

		// Process the packet
		newAddPacket := nflog.ProcessPacket(packet, uint32(p.seq))
		if newAddPacket.Length >= 0 {
			addPackets = append(addPackets, newAddPacket)
		}
	}
	sliceHeader = nil
	packet = nil
	return addPackets
}

// Current nflog error
func strerror() error {
	return syscall.Errno(C.nflog_errno)
}

// Connects to the group specified with the size
func (nflog *NfLog) makeGroup(group, size int) {
	if *Debug {
		log.Printf("Binding this socket to group %d", group)
	}
	gh := C.nflog_bind_group(nflog.h, (C.u_int16_t)(group))
	if gh == nil {
		log.Fatalf("nflog_bind_group failed: %s", strerror())
	}
	nflog.gh = gh

	// Set the maximum amount of logs in buffer for this group
	if C.nflog_set_qthresh(gh, MaxQueueLogs) < 0 {
		log.Fatalf("nflog_set_qthresh failed: %s", strerror())
	}

	// Set local sequence numbering to detect missing packets
	if C.nflog_set_flags(gh, C.NFULNL_CFG_F_SEQ) < 0 {
		log.Fatalf("nflog_set_flags failed: %s", strerror())
	}

	// Set buffer size large
	if C.nflog_set_nlbufsiz(gh, NflogBufferSize) < 0 {
		log.Fatalf("nflog_set_nlbufsiz: %s", strerror())
	}

	// Set timeout
	if C.nflog_set_timeout(gh, NflogTimeout) < 0 {
		log.Fatalf("nflog_set_timeout: %s", strerror())
	}

	if *Debug {
		log.Printf("Setting copy_packet mode to %d bytes", size)
	}
	if C.nflog_set_mode(gh, C.NFULNL_COPY_PACKET, (C.uint)(size)) < 0 {
		log.Fatalf("nflog_set_mode failed: %s", strerror())
	}

	// Register the callback now we are set up
	//
	// Note that we pass a block of memory allocated by C.malloc -
	// it isn't a good idea for C to hold pointers to go objects
	// which might move
	C._callback_register(gh, nflog.packets)
}

// Receive packets in a loop until quit
func (nflog *NfLog) Loop() {
	buflen := C.size_t(RecvBufferSize)
	pbuf := C.malloc(buflen)
	if pbuf == nil {
		log.Fatal("No memory for malloc")
	}
	defer C.free(pbuf)
	for {
		nr := C.recv(nflog.fd, pbuf, buflen, 0)
		select {
		case <-nflog.quit:
			return
		default:
		}
		if nr < 0 {
			log.Printf("Recvfrom failed: %s", strerror())
			nflog.errors++
		} else {
			// Handle messages in packet reusing memory
			ps := <-nflog.a.returnAddPackets
			nflog.packets.index = 0
			C.nflog_handle_packet(nflog.h, (*C.char)(pbuf), (C.int)(nr))
			nflog.a.processAddPackets <- nflog.processPackets(ps[:0])
		}
	}

}

// Close the NfLog down
func (nflog *NfLog) Close() {
	close(nflog.quit)
	// Sometimes hangs and doesn't seem to be necessary
	// if *Debug {
	// 	log.Printf("Unbinding socket %d from group %d", nflog.fd, nflog.McastGroup)
	// }
	// if C.nflog_unbind_group(nflog.gh) < 0 {
	// 	log.Printf("nflog_unbind_group(%d) failed: %s", nflog.McastGroup, strerror())
	// }
	if *Debug {
		log.Printf("Closing nflog socket %d group %d", nflog.fd, nflog.McastGroup)
	}
	if C.nflog_close(nflog.h) < 0 {
		log.Printf("nflog_close failed: %s", strerror())
	}
	// Mark this index as no longer in use
	nflogs[nflog.index] = nil
	C.free(unsafe.Pointer(nflog.packets))
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"io"
	"os"
	"text/tabwriter"

	"google.golang.org/grpc"

	"github.com/pachyderm/pachyderm"
	pfscmds "github.com/pachyderm/pachyderm/src/pfs/cmds"
	deploycmds "github.com/pachyderm/pachyderm/src/pkg/deploy/cmds"
	ppscmds "github.com/pachyderm/pachyderm/src/pps/cmds"
	"github.com/spf13/cobra"
	"go.pedge.io/pb/go/google/protobuf"
	"go.pedge.io/pkg/cobra"
	"go.pedge.io/proto/version"
	"golang.org/x/net/context"
)

func PachctlCmd(address string) (*cobra.Command, error) {
	rootCmd := &cobra.Command{
		Use: os.Args[0],
		Long: `Access the Pachyderm API.

Envronment variables:
  ADDRESS=0.0.0.0:650, the server to connect to.
`,
	}
	pfsCmds := pfscmds.Cmds(address)
	for _, cmd := range pfsCmds {
		rootCmd.AddCommand(cmd)
	}
	ppsCmds, err := ppscmds.Cmds(address)
	if err != nil {
		return nil, err
	}
	for _, cmd := range ppsCmds {
		rootCmd.AddCommand(cmd)
	}

	deployCmds := deploycmds.Cmds()
	for _, cmd := range deployCmds {
		rootCmd.AddCommand(cmd)
	}
	version := &cobra.Command{
		Use:   "version",
		Short: "Return version information.",
		Long:  "Return version information.",
		Run: pkgcobra.RunFixedArgs(0, func(args []string) error {
			versionClient, err := getVersionAPIClient(address)
			if err != nil {
				return err
			}
			version, err := versionClient.GetVersion(context.Background(), &google_protobuf.Empty{})
			if err != nil {
				return err
			}
			writer := tabwriter.NewWriter(os.Stdout, 20, 1, 3, ' ', 0)
			printVerisonHeader(writer)
			printVersion(writer, "pachctl", pachyderm.Version)
			printVersion(writer, "pachd", version)
			return writer.Flush()
		}),
	}
	rootCmd.AddCommand(version)
	return rootCmd, nil
}

func getVersionAPIClient(address string) (protoversion.APIClient, error) {
	clientConn, err := grpc.Dial(address, grpc.WithInsecure())
	if err != nil {
		return nil, err
	}
	return protoversion.NewAPIClient(clientConn), nil
}

func printVerisonHeader(w io.Writer) {
	fmt.Fprintf(w, "COMPONENT\tVERSION\t\n")
}

func printVersion(w io.Writer, component string, version *protoversion.Version) {
	fmt.Fprintf(
		w,
		"%s\t%d.%d.%d(%s)\t\n",
		component,
		version.Major,
		version.Minor,
		version.Micro,
		version.Additional,
	)
}
<commit_msg>Fix docstring.`<commit_after>package cmd

import (
	"fmt"
	"io"
	"os"
	"text/tabwriter"

	"google.golang.org/grpc"

	"github.com/pachyderm/pachyderm"
	pfscmds "github.com/pachyderm/pachyderm/src/pfs/cmds"
	deploycmds "github.com/pachyderm/pachyderm/src/pkg/deploy/cmds"
	ppscmds "github.com/pachyderm/pachyderm/src/pps/cmds"
	"github.com/spf13/cobra"
	"go.pedge.io/pb/go/google/protobuf"
	"go.pedge.io/pkg/cobra"
	"go.pedge.io/proto/version"
	"golang.org/x/net/context"
)

func PachctlCmd(address string) (*cobra.Command, error) {
	rootCmd := &cobra.Command{
		Use: os.Args[0],
		Long: `Access the Pachyderm API.

Envronment variables:
  ADDRESS=0.0.0.0:30650, the server to connect to.
`,
	}
	pfsCmds := pfscmds.Cmds(address)
	for _, cmd := range pfsCmds {
		rootCmd.AddCommand(cmd)
	}
	ppsCmds, err := ppscmds.Cmds(address)
	if err != nil {
		return nil, err
	}
	for _, cmd := range ppsCmds {
		rootCmd.AddCommand(cmd)
	}

	deployCmds := deploycmds.Cmds()
	for _, cmd := range deployCmds {
		rootCmd.AddCommand(cmd)
	}
	version := &cobra.Command{
		Use:   "version",
		Short: "Return version information.",
		Long:  "Return version information.",
		Run: pkgcobra.RunFixedArgs(0, func(args []string) error {
			versionClient, err := getVersionAPIClient(address)
			if err != nil {
				return err
			}
			version, err := versionClient.GetVersion(context.Background(), &google_protobuf.Empty{})
			if err != nil {
				return err
			}
			writer := tabwriter.NewWriter(os.Stdout, 20, 1, 3, ' ', 0)
			printVerisonHeader(writer)
			printVersion(writer, "pachctl", pachyderm.Version)
			printVersion(writer, "pachd", version)
			return writer.Flush()
		}),
	}
	rootCmd.AddCommand(version)
	return rootCmd, nil
}

func getVersionAPIClient(address string) (protoversion.APIClient, error) {
	clientConn, err := grpc.Dial(address, grpc.WithInsecure())
	if err != nil {
		return nil, err
	}
	return protoversion.NewAPIClient(clientConn), nil
}

func printVerisonHeader(w io.Writer) {
	fmt.Fprintf(w, "COMPONENT\tVERSION\t\n")
}

func printVersion(w io.Writer, component string, version *protoversion.Version) {
	fmt.Fprintf(
		w,
		"%s\t%d.%d.%d(%s)\t\n",
		component,
		version.Major,
		version.Minor,
		version.Micro,
		version.Additional,
	)
}
<|endoftext|>


<commit_before>package networkwebsockets

import (
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
	"text/template"
	"time"

	tls "github.com/richtr/go-tls-srp"
)

var (
	// Regular expression matchers

	serviceNameRegexStr = "[A-Za-z0-9/\\+=\\*\\._-]{1,255}"

	peerIdRegexStr = "[A-Za-z0-9]{4,255}"

	isValidLocalRequest = regexp.MustCompile(fmt.Sprintf("^((/control)?/network/%s/%s)$", serviceNameRegexStr, peerIdRegexStr))

	isValidControlRequest = regexp.MustCompile(fmt.Sprintf("^(/control/network/%s/%s)$", serviceNameRegexStr, peerIdRegexStr))

	isValidProxyRequest = regexp.MustCompile(fmt.Sprintf("^/%s$", serviceNameRegexStr))

	isValidServiceName = regexp.MustCompile(fmt.Sprintf("^%s$", serviceNameRegexStr))

	// TLS-SRP configuration components

	// We deliberately only use a weak salt because we don't persistently store TLS-SRP credential data
	Salt = []byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}

	serviceTab = CredentialsStore(map[string]string{})
)

type NetworkWebSocket_Service struct {
	Host string
	Port int

	ProxyPort int

	// All Named Web Socket channels that this service manages
	Channels map[string]*NetworkWebSocket

	discoveryBrowser *DiscoveryBrowser

	done chan int // blocks until .Stop() is called on this service
}

func NewNetworkWebSocketService(host string, port int) *NetworkWebSocket_Service {
	if host == "" {
		hostname, err := os.Hostname()
		if err != nil {
			log.Printf("Could not determine device hostname: %v\n", err)
			return nil
		}
		host = hostname
	}

	if port <= 1024 || port >= 65534 {
		port = 9009
	}

	service := &NetworkWebSocket_Service{
		Host: host,
		Port: port,

		ProxyPort: 0,

		Channels: make(map[string]*NetworkWebSocket),

		discoveryBrowser: NewDiscoveryBrowser(),

		done: make(chan int),
	}

	return service
}

func (service *NetworkWebSocket_Service) Start() <-chan int {
	// Start mDNS/DNS-SD Network Web Socket discovery service
	go service.StartDiscoveryBrowser(10)

	// Start HTTP/Network Web Socket creation server
	go service.StartHTTPServer()

	// Start TLS-SRP Network Web Socket (wss) proxy server
	go service.StartProxyServer()

	return service.StopNotify()
}

func (service *NetworkWebSocket_Service) StartHTTPServer() {
	// Create a new custom http server multiplexer
	serveMux := http.NewServeMux()

	// Serve the test console
	serveMux.HandleFunc("/", service.serveConsoleTemplate)

	// Serve network web socket creation endpoints for localhost clients
	serveMux.HandleFunc("/network/", service.serveLocalWSCreator)

	// Serve network web socket control endpoints for localhost clients
	serveMux.HandleFunc("/control/", service.serveLocalWSCreator)

	// Listen and on loopback address + port
	listener, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", service.Port))
	if err != nil {
		log.Fatal("Could not serve web server. ", err)
	}

	log.Printf("Serving Named Web Socket Creator Proxy at address [ ws://localhost:%d/ ]", service.Port)

	http.Serve(listener, serveMux)
}

func (service *NetworkWebSocket_Service) StartProxyServer() {
	// Create a new custom http server multiplexer
	serveMux := http.NewServeMux()

	// Serve secure network web socket proxy endpoints for network clients
	serveMux.HandleFunc("/", service.serveProxyWSCreator)

	// Generate random server salt for use in TLS-SRP data storage
	var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
	b := make([]rune, 32)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	srpSaltKey := string(b)

	tlsServerConfig := &tls.Config{
		SRPLookup:   serviceTab,
		SRPSaltKey:  srpSaltKey,
		SRPSaltSize: len(Salt),
	}

	// Listen on all addresses + port
	tlsSrpListener, err := tls.Listen("tcp", ":0", tlsServerConfig)
	if err != nil {
		log.Fatal("Could not serve proxy server. ", err)
	}

	// Obtain and store the port of the proxy endpoint
	_, port, err := net.SplitHostPort(tlsSrpListener.Addr().String())
	if err != nil {
		log.Fatal("Could not determine bound port of proxy server. ", err)
	}

	service.ProxyPort, _ = strconv.Atoi(port)

	log.Printf("Serving Named Web Socket Network Proxy at address [ wss://%s:%d/ ]", service.Host, service.ProxyPort)

	http.Serve(tlsSrpListener, serveMux)
}

func (service *NetworkWebSocket_Service) StartDiscoveryBrowser(timeoutSeconds int) {
	defer service.discoveryBrowser.Shutdown()

	log.Printf("Listening for Named Web Socket services on the local network...")

	for !service.discoveryBrowser.closed {
		service.discoveryBrowser.Browse(service, timeoutSeconds)
	}
}

func (service *NetworkWebSocket_Service) serveConsoleTemplate(w http.ResponseWriter, r *http.Request) {

	w.Header().Set("Content-Type", "text/html; charset=utf-8")

	// Only allow console access from localhost
	if isRequestFromLocalHost := service.checkRequestIsFromLocalHost(r.Host); !isRequestFromLocalHost {
		http.Error(w, fmt.Sprintf("Named WebSockets Test Console is only accessible from the local machine (i.e http://localhost:%d/console)", service.Port), 403)
		return
	}

	if r.Method != "GET" {
		http.Error(w, "Method Not Allowed", 405)
		return
	}

	if r.URL.Path == "/" {
		fmt.Fprint(w, "
A Named WebSockets Proxy is running on this host
")
		return
	}

	if r.URL.Path != "/console" {
		http.Error(w, "Not Found", 404)
		return
	}

	consoleHTML, err := Asset("_templates/console.html")
	if err != nil {
		// Asset was not found.
		http.Error(w, "Not Found", 404)
		return
	}

	t := template.Must(template.New("console").Parse(string(consoleHTML)))
	if t == nil {
		http.Error(w, "Internal Server Error", 501)
		return
	}

	t.Execute(w, service.Port)
}

func (service *NetworkWebSocket_Service) serveLocalWSCreator(w http.ResponseWriter, r *http.Request) {
	// Only allow access from localhost to all services
	if isRequestFromLocalHost := service.checkRequestIsFromLocalHost(r.Host); !isRequestFromLocalHost {
		http.Error(w, fmt.Sprintf("This interface is only accessible from the local machine on this port (%d)", service.Port), 403)
		return
	}

	if r.Method != "GET" {
		http.Error(w, "Method Not Allowed", 405)
		return
	}

	if isValidRequest := isValidLocalRequest.MatchString(r.URL.Path); !isValidRequest {
		http.Error(w, "Not Found", 404)
		return
	}

	if isValidWSUpgradeRequest := strings.ToLower(r.Header.Get("Upgrade")); isValidWSUpgradeRequest != "websocket" {
		http.Error(w, "Bad Request", 400)
		return
	}

	isControl := isValidControlRequest.MatchString(r.URL.Path)

	pathParts := strings.Split(r.URL.Path, "/")

	peerId := pathParts[len(pathParts)-1]
	serviceName := pathParts[len(pathParts)-2]

	if isValid := isValidServiceName.MatchString(serviceName); !isValid {
		http.Error(w, "Not Found", 404)
		return
	}

	// Resolve to web socket connection channel
	sock := service.getServiceByName(serviceName)

	if sock == nil {
		sock = NewNetworkWebSocket(service, serviceName, isControl)
	}

	// Handle websocket connection
	if isControl {
		sock.serveControl(w, r, peerId)
	} else {
		sock.servePeer(w, r, peerId)
	}
}

func (service *NetworkWebSocket_Service) serveProxyWSCreator(w http.ResponseWriter, r *http.Request) {

	if r.Method != "GET" {
		http.Error(w, "Method Not Allowed", 405)
		return
	}

	if isValidRequest := isValidProxyRequest.MatchString(r.URL.Path); !isValidRequest {
		http.Error(w, "Not Found", 404)
		return
	}

	if isValidWSUpgradeRequest := strings.ToLower(r.Header.Get("Upgrade")); isValidWSUpgradeRequest != "websocket" {
		http.Error(w, "Bad Request", 400)
		return
	}

	// Resolve servicePath to an active named websocket service
	for _, sock := range service.Channels {
		if sock.proxyPath == r.URL.Path {
			// Generate a new id for this proxy connection
			rand.Seed(time.Now().UTC().UnixNano())
			peerId := fmt.Sprintf("%d", rand.Int())

			sock.serveProxy(w, r, peerId)
			return
		}
	}

	http.Error(w, "Not Found", 404)
	return
}

// Check whether we know the given service name
func (service *NetworkWebSocket_Service) getServiceByName(serviceName string) *NetworkWebSocket {
	for _, sock := range service.Channels {
		if sock.serviceName == serviceName {
			return sock
		}
	}
	return nil
}

// Check whether a DNS-SD derived Network Web Socket hash is owned by the current proxy instance
func (service *NetworkWebSocket_Service) isOwnProxyService(serviceRecord *NetworkWebSocket_DNSRecord) bool {
	for _, sock := range service.Channels {
		if sock.serviceHash == serviceRecord.Hash_Base64 {
			return true
		}
	}
	return false
}

// Check whether a DNS-SD derived Network Web Socket hash is currently connected as a service
func (service *NetworkWebSocket_Service) isActiveProxyService(serviceRecord *NetworkWebSocket_DNSRecord) bool {
	for _, sock := range service.Channels {
		for _, proxy := range sock.proxies {
			if proxy.Hash_Base64 == serviceRecord.Hash_Base64 {
				return true
			}
		}
	}
	return false
}

// Stop stops the server gracefully, and shuts down the running goroutine.
// Stop should be called after a Start(s), otherwise it will block forever.
func (service *NetworkWebSocket_Service) Stop() {
	service.done <- 1
}

// StopNotify returns a channel that receives a empty integer
// when the server is stopped.
func (service *NetworkWebSocket_Service) StopNotify() <-chan int { return service.done }

//
// HELPER FUNCTIONS
//

func (service *NetworkWebSocket_Service) checkRequestIsFromLocalHost(host string) bool {
	allowedLocalHosts := map[string]bool{
		fmt.Sprintf("localhost:%d", service.Port):        true,
		fmt.Sprintf("127.0.0.1:%d", service.Port):        true,
		fmt.Sprintf("::1:%d", service.Port):              true,
		fmt.Sprintf("%s:%d", service.Host, service.Port): true,
	}

	if allowedLocalHosts[host] {
		return true
	}

	return false
}

/** Simple in-memory storage table for TLS-SRP usernames/passwords **/

type CredentialsStore map[string]string

func (cs CredentialsStore) Lookup(user string) (v, s []byte, grp tls.SRPGroup, err error) {
	grp = tls.SRPGroup4096

	p := cs[user]
	if p == "" {
		return nil, nil, grp, nil
	}

	v = tls.SRPVerifier(user, p, Salt, grp)
	return v, Salt, grp, nil
}
<commit_msg>Make channel lookup exportable<commit_after>package networkwebsockets

import (
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
	"text/template"
	"time"

	tls "github.com/richtr/go-tls-srp"
)

var (
	// Regular expression matchers

	serviceNameRegexStr = "[A-Za-z0-9/\\+=\\*\\._-]{1,255}"

	peerIdRegexStr = "[A-Za-z0-9]{4,255}"

	isValidLocalRequest = regexp.MustCompile(fmt.Sprintf("^((/control)?/network/%s/%s)$", serviceNameRegexStr, peerIdRegexStr))

	isValidControlRequest = regexp.MustCompile(fmt.Sprintf("^(/control/network/%s/%s)$", serviceNameRegexStr, peerIdRegexStr))

	isValidProxyRequest = regexp.MustCompile(fmt.Sprintf("^/%s$", serviceNameRegexStr))

	isValidServiceName = regexp.MustCompile(fmt.Sprintf("^%s$", serviceNameRegexStr))

	// TLS-SRP configuration components

	// We deliberately only use a weak salt because we don't persistently store TLS-SRP credential data
	Salt = []byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}

	serviceTab = CredentialsStore(map[string]string{})
)

type NetworkWebSocket_Service struct {
	Host string
	Port int

	ProxyPort int

	// All Named Web Socket channels that this service manages
	Channels map[string]*NetworkWebSocket

	discoveryBrowser *DiscoveryBrowser

	done chan int // blocks until .Stop() is called on this service
}

func NewNetworkWebSocketService(host string, port int) *NetworkWebSocket_Service {
	if host == "" {
		hostname, err := os.Hostname()
		if err != nil {
			log.Printf("Could not determine device hostname: %v\n", err)
			return nil
		}
		host = hostname
	}

	if port <= 1024 || port >= 65534 {
		port = 9009
	}

	service := &NetworkWebSocket_Service{
		Host: host,
		Port: port,

		ProxyPort: 0,

		Channels: make(map[string]*NetworkWebSocket),

		discoveryBrowser: NewDiscoveryBrowser(),

		done: make(chan int),
	}

	return service
}

func (service *NetworkWebSocket_Service) Start() <-chan int {
	// Start mDNS/DNS-SD Network Web Socket discovery service
	go service.StartDiscoveryBrowser(10)

	// Start HTTP/Network Web Socket creation server
	go service.StartHTTPServer()

	// Start TLS-SRP Network Web Socket (wss) proxy server
	go service.StartProxyServer()

	return service.StopNotify()
}

func (service *NetworkWebSocket_Service) StartHTTPServer() {
	// Create a new custom http server multiplexer
	serveMux := http.NewServeMux()

	// Serve the test console
	serveMux.HandleFunc("/", service.serveConsoleTemplate)

	// Serve network web socket creation endpoints for localhost clients
	serveMux.HandleFunc("/network/", service.serveLocalWSCreator)

	// Serve network web socket control endpoints for localhost clients
	serveMux.HandleFunc("/control/", service.serveLocalWSCreator)

	// Listen and on loopback address + port
	listener, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", service.Port))
	if err != nil {
		log.Fatal("Could not serve web server. ", err)
	}

	log.Printf("Serving Named Web Socket Creator Proxy at address [ ws://localhost:%d/ ]", service.Port)

	http.Serve(listener, serveMux)
}

func (service *NetworkWebSocket_Service) StartProxyServer() {
	// Create a new custom http server multiplexer
	serveMux := http.NewServeMux()

	// Serve secure network web socket proxy endpoints for network clients
	serveMux.HandleFunc("/", service.serveProxyWSCreator)

	// Generate random server salt for use in TLS-SRP data storage
	var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
	b := make([]rune, 32)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	srpSaltKey := string(b)

	tlsServerConfig := &tls.Config{
		SRPLookup:   serviceTab,
		SRPSaltKey:  srpSaltKey,
		SRPSaltSize: len(Salt),
	}

	// Listen on all addresses + port
	tlsSrpListener, err := tls.Listen("tcp", ":0", tlsServerConfig)
	if err != nil {
		log.Fatal("Could not serve proxy server. ", err)
	}

	// Obtain and store the port of the proxy endpoint
	_, port, err := net.SplitHostPort(tlsSrpListener.Addr().String())
	if err != nil {
		log.Fatal("Could not determine bound port of proxy server. ", err)
	}

	service.ProxyPort, _ = strconv.Atoi(port)

	log.Printf("Serving Named Web Socket Network Proxy at address [ wss://%s:%d/ ]", service.Host, service.ProxyPort)

	http.Serve(tlsSrpListener, serveMux)
}

func (service *NetworkWebSocket_Service) StartDiscoveryBrowser(timeoutSeconds int) {
	defer service.discoveryBrowser.Shutdown()

	log.Printf("Listening for Named Web Socket services on the local network...")

	for !service.discoveryBrowser.closed {
		service.discoveryBrowser.Browse(service, timeoutSeconds)
	}
}

func (service *NetworkWebSocket_Service) serveConsoleTemplate(w http.ResponseWriter, r *http.Request) {

	w.Header().Set("Content-Type", "text/html; charset=utf-8")

	// Only allow console access from localhost
	if isRequestFromLocalHost := service.checkRequestIsFromLocalHost(r.Host); !isRequestFromLocalHost {
		http.Error(w, fmt.Sprintf("Named WebSockets Test Console is only accessible from the local machine (i.e http://localhost:%d/console)", service.Port), 403)
		return
	}

	if r.Method != "GET" {
		http.Error(w, "Method Not Allowed", 405)
		return
	}

	if r.URL.Path == "/" {
		fmt.Fprint(w, "
A Named WebSockets Proxy is running on this host
")
		return
	}

	if r.URL.Path != "/console" {
		http.Error(w, "Not Found", 404)
		return
	}

	consoleHTML, err := Asset("_templates/console.html")
	if err != nil {
		// Asset was not found.
		http.Error(w, "Not Found", 404)
		return
	}

	t := template.Must(template.New("console").Parse(string(consoleHTML)))
	if t == nil {
		http.Error(w, "Internal Server Error", 501)
		return
	}

	t.Execute(w, service.Port)
}

func (service *NetworkWebSocket_Service) serveLocalWSCreator(w http.ResponseWriter, r *http.Request) {
	// Only allow access from localhost to all services
	if isRequestFromLocalHost := service.checkRequestIsFromLocalHost(r.Host); !isRequestFromLocalHost {
		http.Error(w, fmt.Sprintf("This interface is only accessible from the local machine on this port (%d)", service.Port), 403)
		return
	}

	if r.Method != "GET" {
		http.Error(w, "Method Not Allowed", 405)
		return
	}

	if isValidRequest := isValidLocalRequest.MatchString(r.URL.Path); !isValidRequest {
		http.Error(w, "Not Found", 404)
		return
	}

	if isValidWSUpgradeRequest := strings.ToLower(r.Header.Get("Upgrade")); isValidWSUpgradeRequest != "websocket" {
		http.Error(w, "Bad Request", 400)
		return
	}

	isControl := isValidControlRequest.MatchString(r.URL.Path)

	pathParts := strings.Split(r.URL.Path, "/")

	peerId := pathParts[len(pathParts)-1]
	serviceName := pathParts[len(pathParts)-2]

	if isValid := isValidServiceName.MatchString(serviceName); !isValid {
		http.Error(w, "Not Found", 404)
		return
	}

	// Resolve to web socket connection channel
	sock := service.GetChannelByName(serviceName)

	if sock == nil {
		sock = NewNetworkWebSocket(service, serviceName, isControl)
	}

	// Handle websocket connection
	if isControl {
		sock.serveControl(w, r, peerId)
	} else {
		sock.servePeer(w, r, peerId)
	}
}

func (service *NetworkWebSocket_Service) serveProxyWSCreator(w http.ResponseWriter, r *http.Request) {

	if r.Method != "GET" {
		http.Error(w, "Method Not Allowed", 405)
		return
	}

	if isValidRequest := isValidProxyRequest.MatchString(r.URL.Path); !isValidRequest {
		http.Error(w, "Not Found", 404)
		return
	}

	if isValidWSUpgradeRequest := strings.ToLower(r.Header.Get("Upgrade")); isValidWSUpgradeRequest != "websocket" {
		http.Error(w, "Bad Request", 400)
		return
	}

	// Resolve servicePath to an active named websocket service
	for _, sock := range service.Channels {
		if sock.proxyPath == r.URL.Path {
			// Generate a new id for this proxy connection
			rand.Seed(time.Now().UTC().UnixNano())
			peerId := fmt.Sprintf("%d", rand.Int())

			sock.serveProxy(w, r, peerId)
			return
		}
	}

	http.Error(w, "Not Found", 404)
	return
}

// Check whether we know the given service name
func (service *NetworkWebSocket_Service) GetChannelByName(serviceName string) *NetworkWebSocket {
	for _, sock := range service.Channels {
		if sock.serviceName == serviceName {
			return sock
		}
	}
	return nil
}

// Check whether a DNS-SD derived Network Web Socket hash is owned by the current proxy instance
func (service *NetworkWebSocket_Service) isOwnProxyService(serviceRecord *NetworkWebSocket_DNSRecord) bool {
	for _, sock := range service.Channels {
		if sock.serviceHash == serviceRecord.Hash_Base64 {
			return true
		}
	}
	return false
}

// Check whether a DNS-SD derived Network Web Socket hash is currently connected as a service
func (service *NetworkWebSocket_Service) isActiveProxyService(serviceRecord *NetworkWebSocket_DNSRecord) bool {
	for _, sock := range service.Channels {
		for _, proxy := range sock.proxies {
			if proxy.Hash_Base64 == serviceRecord.Hash_Base64 {
				return true
			}
		}
	}
	return false
}

// Stop stops the server gracefully, and shuts down the running goroutine.
// Stop should be called after a Start(s), otherwise it will block forever.
func (service *NetworkWebSocket_Service) Stop() {
	service.done <- 1
}

// StopNotify returns a channel that receives a empty integer
// when the server is stopped.
func (service *NetworkWebSocket_Service) StopNotify() <-chan int { return service.done }

//
// HELPER FUNCTIONS
//

func (service *NetworkWebSocket_Service) checkRequestIsFromLocalHost(host string) bool {
	allowedLocalHosts := map[string]bool{
		fmt.Sprintf("localhost:%d", service.Port):        true,
		fmt.Sprintf("127.0.0.1:%d", service.Port):        true,
		fmt.Sprintf("::1:%d", service.Port):              true,
		fmt.Sprintf("%s:%d", service.Host, service.Port): true,
	}

	if allowedLocalHosts[host] {
		return true
	}

	return false
}

/** Simple in-memory storage table for TLS-SRP usernames/passwords **/

type CredentialsStore map[string]string

func (cs CredentialsStore) Lookup(user string) (v, s []byte, grp tls.SRPGroup, err error) {
	grp = tls.SRPGroup4096

	p := cs[user]
	if p == "" {
		return nil, nil, grp, nil
	}

	v = tls.SRPVerifier(user, p, Salt, grp)
	return v, Salt, grp, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This benchmark measures time spent on local (non-synced) operations.
//
// All tests use 9-byte keys.
// "Tiny" benchmarks use testStruct values with empty []byte.
// "Huge" benchmarks use testStruct values with 100K-byte []byte.
// Scan and Exec benchmarks iterate once through all b.N records.
//
// To run: jiri go test v.io/v23/syncbase/nosql -test.bench=. -test.run=Benchmark -alsologtostderr=false -stderrthreshold=3
//
// Results (2016-01-15, MacBook Pro, 2.8 GHz Intel Core i7):
// BenchmarkTinyPut-8      	    1000	   1260983 ns/op
// BenchmarkTinyGet-8      	    2000	    772075 ns/op
// BenchmarkTinyDelete-8   	    2000	   1497883 ns/op
// BenchmarkTinyScan-8     	      50	  25266715 ns/op
// BenchmarkTinyExec-8     	     300	   6076609 ns/op
// BenchmarkTinyWatchPuts-8	      10	 179642079 ns/op
// BenchmarkHugePut-8      	    1000	   2916679 ns/op
// BenchmarkHugeGet-8      	    1000	   1982792 ns/op
// BenchmarkHugeDelete-8   	    1000	   1606726 ns/op
// BenchmarkHugeScan-8     	      10	 133327996 ns/op
// BenchmarkHugeExec-8     	      20	 104448117 ns/op
// BenchmarkHugeWatchPuts-8	       3	 384569733 ns/op
package nosql_test

import (
	"fmt"
	"math/rand"
	"testing"

	"v.io/v23/context"
	wire "v.io/v23/services/syncbase/nosql"
	"v.io/v23/services/watch"
	"v.io/v23/syncbase"
	"v.io/v23/syncbase/nosql"
	_ "v.io/x/ref/runtime/factories/generic"
	tu "v.io/x/ref/services/syncbase/testutil"
)

// prepare creates hierarchy "a/d/tb" and returns some handles along with a
// cleanup function.
func prepare(b *testing.B) (*context.T, nosql.Database, nosql.Table, func()) {
	ctx, sName, cleanup := tu.SetupOrDie(nil)
	s := syncbase.NewService(sName)
	a := s.App("a")
	if err := a.Create(ctx, nil); err != nil {
		b.Fatalf("can't create app: %v", err)
	}
	d := a.NoSQLDatabase("d", nil)
	if err := d.Create(ctx, nil); err != nil {
		b.Fatalf("can't create database: %v", err)
	}
	tb := d.Table("tb")
	if err := tb.Create(ctx, nil); err != nil {
		b.Fatalf("can't create table: %v", err)
	}
	return ctx, d, tb, func() {
		b.StopTimer()
		cleanup()
		b.StartTimer()
	}
}

type blobStruct struct {
	String  string
	BlobRef wire.BlobRef
}

type testStruct struct {
	A string
	B int
	C *blobStruct
	D []byte
}

func makeKey(i int) string {
	return fmt.Sprintf("%09d", i)
}

func makeTestStruct() interface{} {
	return testStruct{A: "hello, world!", B: 42}
}

func makeTestStruct100K() interface{} {
	var byteSlice []byte
	r := rand.New(rand.NewSource(23917))
	for i := 0; i < 100*1000; i++ {
		byteSlice = append(byteSlice, byte(r.Intn(256)))
	}
	return testStruct{A: "hello, world!", B: 42, D: byteSlice}
}

func writeRows(b *testing.B, ctx *context.T, tb nosql.Table, value interface{}) {
	writeRowsCustom(b, ctx, tb, value, b.N)
}

func writeRowsCustom(b *testing.B, ctx *context.T, tb nosql.Table, value interface{}, n int) {
	for i := 0; i < n; i++ {
		if err := tb.Put(ctx, makeKey(i), value); err != nil {
			b.Fatal(err)
		}
	}
}

func runPutBenchmark(b *testing.B, value interface{}) {
	ctx, _, tb, cleanup := prepare(b)
	defer cleanup()
	b.ResetTimer()
	writeRows(b, ctx, tb, value)
}

func runGetBenchmark(b *testing.B, value interface{}) {
	ctx, _, tb, cleanup := prepare(b)
	defer cleanup()
	writeRows(b, ctx, tb, value)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		var got testStruct
		if err := tb.Get(ctx, makeKey(i), &got); err != nil {
			b.Fatalf("tb.Get failed: %v", err)
		}
	}
}

func runDeleteBenchmark(b *testing.B, value interface{}) {
	ctx, _, tb, cleanup := prepare(b)
	defer cleanup()
	writeRows(b, ctx, tb, value)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if err := tb.Delete(ctx, makeKey(i)); err != nil {
			b.Fatalf("tb.Delete failed: %v", err)
		}
	}
}

const numRows = 100

// Measures how long it takes to process 'numRows' rows using scan.
func runScanBenchmark(b *testing.B, value interface{}) {
	ctx, _, tb, cleanup := prepare(b)
	defer cleanup()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		b.StopTimer()
		// TODO(sadovsky): Write rows just once, and clear any read caches on every
		// iteration.
		writeRowsCustom(b, ctx, tb, value, numRows)
		b.StartTimer()
		s := tb.Scan(ctx, nosql.Prefix(""))
		var got testStruct
		for s.Advance() {
			s.Value(&got)
		}
		if s.Err() != nil {
			b.Fatalf("stream error: %s", s.Err())
		}
	}
}

// Measures how long it takes to process 'numRows' rows using exec.
func runExecBenchmark(b *testing.B, value interface{}) {
	ctx, d, tb, cleanup := prepare(b)
	defer cleanup()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		b.StopTimer()
		// TODO(sadovsky): Write rows just once, and clear any read caches on every
		// iteration.
		writeRowsCustom(b, ctx, tb, value, numRows)
		b.StartTimer()
		_, s, err := d.Exec(ctx, "select v from tb")
		if err != nil {
			b.Fatalf("exec error: %s", s.Err())
		}
		for s.Advance() {
			s.Result()
		}
		if s.Err() != nil {
			b.Fatalf("stream error: %s", s.Err())
		}
	}
}

// Measures how long it takes to put and get notified about 'numRows' rows.
func runWatchPutsBenchmark(b *testing.B, value interface{}) {
	ctx, d, tb, cleanup := prepare(b)
	defer cleanup()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		w, err := d.Watch(ctx, "tb", "", watch.ResumeMarker("now"))
		if err != nil {
			b.Fatalf("watch error: %v", err)
		}
		done := make(chan struct{})
		go func() {
			seen := 0
			for seen < numRows && w.Advance() {
				seen++
				w.Change()
			}
			if w.Err() != nil {
				b.Fatalf("stream error: %v", w.Err())
			}
			close(done)
		}()
		writeRowsCustom(b, ctx, tb, value, numRows)
		<-done
		w.Cancel()
	}
}

// Measures how long it takes to put and get notified about a single value.
func runWatchOnePutBenchmark(b *testing.B, value interface{}) {
	ctx, d, tb, cleanup := prepare(b)
	defer cleanup()
	w, err := d.Watch(ctx, "tb", "", watch.ResumeMarker("now"))
	row := make(chan struct{})
	if err != nil {
		b.Fatalf("watch error: %v", err)
	}
	go func() {
		seen := 0
		for seen < b.N && w.Advance() {
			seen++
			w.Change()
			row <- struct{}{}
		}
		if w.Err() != nil {
			b.Fatalf("stream error: %v", w.Err())
		}
	}()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		writeRowsCustom(b, ctx, tb, value, 1)
		<-row
	}
	w.Cancel()
}

func BenchmarkTinyPut(b *testing.B) {
	runPutBenchmark(b, makeTestStruct())
}

func BenchmarkTinyGet(b *testing.B) {
	runGetBenchmark(b, makeTestStruct())
}

func BenchmarkTinyDelete(b *testing.B) {
	runDeleteBenchmark(b, makeTestStruct())
}

func BenchmarkTinyScan(b *testing.B) {
	runScanBenchmark(b, makeTestStruct())
}

func BenchmarkTinyExec(b *testing.B) {
	runExecBenchmark(b, makeTestStruct())
}

func BenchmarkTinyWatchPuts(b *testing.B) {
	runWatchPutsBenchmark(b, makeTestStruct())
}

func BenchmarkTinyWatchOnePut(b *testing.B) {
	runWatchOnePutBenchmark(b, makeTestStruct())
}

func BenchmarkHugePut(b *testing.B) {
	runPutBenchmark(b, makeTestStruct100K())
}

func BenchmarkHugeGet(b *testing.B) {
	runGetBenchmark(b, makeTestStruct100K())
}

func BenchmarkHugeDelete(b *testing.B) {
	runDeleteBenchmark(b, makeTestStruct100K())
}

func BenchmarkHugeScan(b *testing.B) {
	runScanBenchmark(b, makeTestStruct100K())
}

func BenchmarkHugeExec(b *testing.B) {
	runExecBenchmark(b, makeTestStruct100K())
}

func BenchmarkHugeWatchPuts(b *testing.B) {
	runWatchPutsBenchmark(b, makeTestStruct100K())
}

func BenchmarkHugeWatchOnePut(b *testing.B) {
	runWatchOnePutBenchmark(b, makeTestStruct100K())
}
<commit_msg>TBR syncbase benchmarks: disable watch bm's, which hang occasionally<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This benchmark measures time spent on local (non-synced) operations.
//
// All tests use 9-byte keys.
// "Tiny" benchmarks use testStruct values with empty []byte.
// "Huge" benchmarks use testStruct values with 100K-byte []byte.
// Scan and Exec benchmarks iterate once through all b.N records.
//
// To run: jiri go test v.io/v23/syncbase/nosql -test.bench=. -test.run=Benchmark -alsologtostderr=false -stderrthreshold=3
//
// Results (2016-01-15, MacBook Pro, 2.8 GHz Intel Core i7):
// BenchmarkTinyPut-8      	    1000	   1260983 ns/op
// BenchmarkTinyGet-8      	    2000	    772075 ns/op
// BenchmarkTinyDelete-8   	    2000	   1497883 ns/op
// BenchmarkTinyScan-8     	      50	  25266715 ns/op
// BenchmarkTinyExec-8     	     300	   6076609 ns/op
// BenchmarkTinyWatchPuts-8	      10	 179642079 ns/op
// BenchmarkHugePut-8      	    1000	   2916679 ns/op
// BenchmarkHugeGet-8      	    1000	   1982792 ns/op
// BenchmarkHugeDelete-8   	    1000	   1606726 ns/op
// BenchmarkHugeScan-8     	      10	 133327996 ns/op
// BenchmarkHugeExec-8     	      20	 104448117 ns/op
// BenchmarkHugeWatchPuts-8	       3	 384569733 ns/op
package nosql_test

import (
	"fmt"
	"math/rand"
	"testing"

	"v.io/v23/context"
	wire "v.io/v23/services/syncbase/nosql"
	"v.io/v23/services/watch"
	"v.io/v23/syncbase"
	"v.io/v23/syncbase/nosql"
	_ "v.io/x/ref/runtime/factories/generic"
	tu "v.io/x/ref/services/syncbase/testutil"
)

// prepare creates hierarchy "a/d/tb" and returns some handles along with a
// cleanup function.
func prepare(b *testing.B) (*context.T, nosql.Database, nosql.Table, func()) {
	ctx, sName, cleanup := tu.SetupOrDie(nil)
	s := syncbase.NewService(sName)
	a := s.App("a")
	if err := a.Create(ctx, nil); err != nil {
		b.Fatalf("can't create app: %v", err)
	}
	d := a.NoSQLDatabase("d", nil)
	if err := d.Create(ctx, nil); err != nil {
		b.Fatalf("can't create database: %v", err)
	}
	tb := d.Table("tb")
	if err := tb.Create(ctx, nil); err != nil {
		b.Fatalf("can't create table: %v", err)
	}
	return ctx, d, tb, func() {
		b.StopTimer()
		cleanup()
		b.StartTimer()
	}
}

type blobStruct struct {
	String  string
	BlobRef wire.BlobRef
}

type testStruct struct {
	A string
	B int
	C *blobStruct
	D []byte
}

func makeKey(i int) string {
	return fmt.Sprintf("%09d", i)
}

func makeTestStruct() interface{} {
	return testStruct{A: "hello, world!", B: 42}
}

func makeTestStruct100K() interface{} {
	var byteSlice []byte
	r := rand.New(rand.NewSource(23917))
	for i := 0; i < 100*1000; i++ {
		byteSlice = append(byteSlice, byte(r.Intn(256)))
	}
	return testStruct{A: "hello, world!", B: 42, D: byteSlice}
}

func writeRows(b *testing.B, ctx *context.T, tb nosql.Table, value interface{}) {
	writeRowsCustom(b, ctx, tb, value, b.N)
}

func writeRowsCustom(b *testing.B, ctx *context.T, tb nosql.Table, value interface{}, n int) {
	for i := 0; i < n; i++ {
		if err := tb.Put(ctx, makeKey(i), value); err != nil {
			b.Fatal(err)
		}
	}
}

func runPutBenchmark(b *testing.B, value interface{}) {
	ctx, _, tb, cleanup := prepare(b)
	defer cleanup()
	b.ResetTimer()
	writeRows(b, ctx, tb, value)
}

func runGetBenchmark(b *testing.B, value interface{}) {
	ctx, _, tb, cleanup := prepare(b)
	defer cleanup()
	writeRows(b, ctx, tb, value)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		var got testStruct
		if err := tb.Get(ctx, makeKey(i), &got); err != nil {
			b.Fatalf("tb.Get failed: %v", err)
		}
	}
}

func runDeleteBenchmark(b *testing.B, value interface{}) {
	ctx, _, tb, cleanup := prepare(b)
	defer cleanup()
	writeRows(b, ctx, tb, value)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if err := tb.Delete(ctx, makeKey(i)); err != nil {
			b.Fatalf("tb.Delete failed: %v", err)
		}
	}
}

const numRows = 100

// Measures how long it takes to process 'numRows' rows using scan.
func runScanBenchmark(b *testing.B, value interface{}) {
	ctx, _, tb, cleanup := prepare(b)
	defer cleanup()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		b.StopTimer()
		// TODO(sadovsky): Write rows just once, and clear any read caches on every
		// iteration.
		writeRowsCustom(b, ctx, tb, value, numRows)
		b.StartTimer()
		s := tb.Scan(ctx, nosql.Prefix(""))
		var got testStruct
		for s.Advance() {
			s.Value(&got)
		}
		if s.Err() != nil {
			b.Fatalf("stream error: %s", s.Err())
		}
	}
}

// Measures how long it takes to process 'numRows' rows using exec.
func runExecBenchmark(b *testing.B, value interface{}) {
	ctx, d, tb, cleanup := prepare(b)
	defer cleanup()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		b.StopTimer()
		// TODO(sadovsky): Write rows just once, and clear any read caches on every
		// iteration.
		writeRowsCustom(b, ctx, tb, value, numRows)
		b.StartTimer()
		_, s, err := d.Exec(ctx, "select v from tb")
		if err != nil {
			b.Fatalf("exec error: %s", s.Err())
		}
		for s.Advance() {
			s.Result()
		}
		if s.Err() != nil {
			b.Fatalf("stream error: %s", s.Err())
		}
	}
}

// Measures how long it takes to put and get notified about 'numRows' rows.
func runWatchPutsBenchmark(b *testing.B, value interface{}) {
	b.Skip("Hangs on occasion, for unknown reasons - v.io/i/1134")
	ctx, d, tb, cleanup := prepare(b)
	defer cleanup()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		w, err := d.Watch(ctx, "tb", "", watch.ResumeMarker("now"))
		if err != nil {
			b.Fatalf("watch error: %v", err)
		}
		done := make(chan struct{})
		go func() {
			seen := 0
			for seen < numRows && w.Advance() {
				seen++
				w.Change()
			}
			if w.Err() != nil {
				b.Fatalf("stream error: %v", w.Err())
			}
			close(done)
		}()
		writeRowsCustom(b, ctx, tb, value, numRows)
		<-done
		w.Cancel()
	}
}

// Measures how long it takes to put and get notified about a single value.
func runWatchOnePutBenchmark(b *testing.B, value interface{}) {
	b.Skip("Hangs on occasion, for unknown reasons - v.io/i/1134")
	ctx, d, tb, cleanup := prepare(b)
	defer cleanup()
	w, err := d.Watch(ctx, "tb", "", watch.ResumeMarker("now"))
	row := make(chan struct{})
	if err != nil {
		b.Fatalf("watch error: %v", err)
	}
	go func() {
		seen := 0
		for seen < b.N && w.Advance() {
			seen++
			w.Change()
			row <- struct{}{}
		}
		if w.Err() != nil {
			b.Fatalf("stream error: %v", w.Err())
		}
	}()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		writeRowsCustom(b, ctx, tb, value, 1)
		<-row
	}
	w.Cancel()
}

func BenchmarkTinyPut(b *testing.B) {
	runPutBenchmark(b, makeTestStruct())
}

func BenchmarkTinyGet(b *testing.B) {
	runGetBenchmark(b, makeTestStruct())
}

func BenchmarkTinyDelete(b *testing.B) {
	runDeleteBenchmark(b, makeTestStruct())
}

func BenchmarkTinyScan(b *testing.B) {
	runScanBenchmark(b, makeTestStruct())
}

func BenchmarkTinyExec(b *testing.B) {
	runExecBenchmark(b, makeTestStruct())
}

func BenchmarkTinyWatchPuts(b *testing.B) {
	runWatchPutsBenchmark(b, makeTestStruct())
}

func BenchmarkTinyWatchOnePut(b *testing.B) {
	runWatchOnePutBenchmark(b, makeTestStruct())
}

func BenchmarkHugePut(b *testing.B) {
	runPutBenchmark(b, makeTestStruct100K())
}

func BenchmarkHugeGet(b *testing.B) {
	runGetBenchmark(b, makeTestStruct100K())
}

func BenchmarkHugeDelete(b *testing.B) {
	runDeleteBenchmark(b, makeTestStruct100K())
}

func BenchmarkHugeScan(b *testing.B) {
	runScanBenchmark(b, makeTestStruct100K())
}

func BenchmarkHugeExec(b *testing.B) {
	runExecBenchmark(b, makeTestStruct100K())
}

func BenchmarkHugeWatchPuts(b *testing.B) {
	runWatchPutsBenchmark(b, makeTestStruct100K())
}

func BenchmarkHugeWatchOnePut(b *testing.B) {
	runWatchOnePutBenchmark(b, makeTestStruct100K())
}
<|endoftext|>


<commit_before>package clusterf
/*
 * Internal service state, maintained from config changes.
 */

import (
    "qmsk.net/clusterf/config"
    "log"
)

type Service struct {
    Name        string

    Frontend    *config.ServiceFrontend
    Backends    map[string]config.ServiceBackend

    driverFrontend  *ipvsFrontend
    driverBackends  map[string]*ipvsBackend
}

func newService(name string) *Service {
    return &Service{
        Name:           name,
        Backends:       make(map[string]config.ServiceBackend),

        driverBackends: make(map[string]*ipvsBackend),
    }
}

func (self *Service) driverError(err error) {
    log.Printf("cluster:Service %s: Error: %s\n", self.Name, err)
}

/* Configuration actions */
func (self *Service) configFrontend(action config.Action, frontendConfig *config.ConfigServiceFrontend) {
    frontend := frontendConfig.Frontend

    log.Printf("clusterf:Service %s: Frontend: %s %+v <- %+v\n", self.Name, action, frontend, self.Frontend)

    switch action {
    case config.NewConfig:
        self.Frontend = &frontend

    case config.SetConfig:
        if self.Frontend == nil {
            self.newFrontend(frontend)
        } else if *self.Frontend != frontend {
            self.setFrontend(frontend)
        }

        self.Frontend = &frontend

    case config.DelConfig:
        self.delFrontend()

        self.Frontend = nil
    }
}

func (self *Service) configBackend(backendName string, action config.Action, backendConfig *config.ConfigServiceBackend) {
    log.Printf("clusterf:Service %s: Backend %s: %s %+v <- %+v\n", self.Name, backendName, action, backendConfig.Backend, self.Backends[backendName])

    switch action {
    case config.NewConfig:
        self.Backends[backendName] = backendConfig.Backend

    case config.SetConfig:
        if self.Backends[backendName] != backendConfig.Backend {
            self.setBackend(backendName, backendConfig.Backend)

            self.Backends[backendName] = backendConfig.Backend
        }

    case config.DelConfig:
        self.delBackend(backendName)

        delete(self.Backends, backendName)
    }
}

// Synchronize state to IPVS
func (self *Service) sync(driver *IPVSDriver) {
    self.driverFrontend = driver.newFrontend()

    if self.Frontend != nil {
        // also adds backends
        self.newFrontend(*self.Frontend)
    }
}

/* Frontend actions */
func (self *Service) newFrontend(frontend config.ServiceFrontend) {
    log.Printf("clusterf:Service %s: new Frontend: %+v\n", self.Name, frontend)

    if err := self.driverFrontend.add(frontend); err != nil {
        self.driverError(err)
    } else {
        for backendName, backend := range self.Backends {
            self.newBackend(backendName, backend)
        }
    }
}

func (self *Service) setFrontend(frontend config.ServiceFrontend) {
    log.Printf("clusterf:Service %s: set Frontend: %+v\n", self.Name, frontend)

    if self.Frontend != nil {
        // TODO: smoother setup-before-teardown transition..?
        self.delFrontend()
    }

    self.newFrontend(frontend)
}

func (self *Service) delFrontend() {
    log.Printf("clusterf:Service %s: del Frontend: %+v\n", self.Name, self.Frontend)

    // del'ing the frontend will also remove all backend state
    if err := self.driverFrontend.del(); err != nil {
        self.driverError(err)
    }

    // XXX: clear self.driverBackend[*]?
}

/* Backend actions */
func (self *Service) newBackend(backendName string, backend config.ServiceBackend) {
    log.Printf("clusterf:Service %s: new Backend %s: %+v\n", self.Name, backendName, backend)

    self.driverBackends[backendName] = self.driverFrontend.newBackend()

    if err := self.driverBackends[backendName].add(backend); err != nil {
        self.driverError(err)
    }
}

func (self *Service) setBackend(backendName string, backend config.ServiceBackend) {
    log.Printf("clusterf:Service %s: set Backend %s: %+v\n", self.Name, backendName, backend)

    if driverBackend := self.driverBackends[backendName]; driverBackend == nil {
        self.newBackend(backendName, backend)
    } else if err := driverBackend.set(backend); err != nil {
        self.driverError(err)
    }
}

func (self *Service) delBackend(backendName string) {
    log.Printf("clusterf:Service %s: del Backend %s: %+v\n", self.Name, backendName, self.Backends[backendName])

    if err := self.driverBackends[backendName].del(); err != nil {
        self.driverError(err)
    }

    delete(self.driverBackends, backendName)
}
<commit_msg>service: crashfix configBackend() without active Frontend<commit_after>package clusterf
/*
 * Internal service state, maintained from config changes.
 */

import (
    "qmsk.net/clusterf/config"
    "log"
)

type Service struct {
    Name        string

    Frontend    *config.ServiceFrontend
    Backends    map[string]config.ServiceBackend

    driverFrontend  *ipvsFrontend
    driverBackends  map[string]*ipvsBackend
}

func newService(name string) *Service {
    return &Service{
        Name:           name,
        Backends:       make(map[string]config.ServiceBackend),

        driverBackends: make(map[string]*ipvsBackend),
    }
}

func (self *Service) driverError(err error) {
    log.Printf("cluster:Service %s: Error: %s\n", self.Name, err)
}

/* Configuration actions */
func (self *Service) configFrontend(action config.Action, frontendConfig *config.ConfigServiceFrontend) {
    frontend := frontendConfig.Frontend

    log.Printf("clusterf:Service %s: Frontend: %s %+v <- %+v\n", self.Name, action, frontend, self.Frontend)

    switch action {
    case config.NewConfig:
        self.Frontend = &frontend

    case config.SetConfig:
        if self.Frontend == nil {
            self.newFrontend(frontend)
        } else if *self.Frontend != frontend {
            self.setFrontend(frontend)
        }

        self.Frontend = &frontend

    case config.DelConfig:
        self.delFrontend()

        self.Frontend = nil
    }
}

func (self *Service) configBackend(backendName string, action config.Action, backendConfig *config.ConfigServiceBackend) {
    log.Printf("clusterf:Service %s: Backend %s: %s %+v <- %+v\n", self.Name, backendName, action, backendConfig.Backend, self.Backends[backendName])

    switch action {
    case config.NewConfig:
        self.Backends[backendName] = backendConfig.Backend

    case config.SetConfig:
        if self.Backends[backendName] == backendConfig.Backend {
            return
        }

        if self.Frontend != nil {
            self.setBackend(backendName, backendConfig.Backend)
        }

        self.Backends[backendName] = backendConfig.Backend

    case config.DelConfig:
        if self.Frontend != nil {
            self.delBackend(backendName)
        }

        delete(self.Backends, backendName)
    }
}

// Synchronize state to IPVS
func (self *Service) sync(driver *IPVSDriver) {
    self.driverFrontend = driver.newFrontend()

    if self.Frontend != nil {
        // also adds backends
        self.newFrontend(*self.Frontend)
    }
}

/* Frontend actions */
func (self *Service) newFrontend(frontend config.ServiceFrontend) {
    log.Printf("clusterf:Service %s: new Frontend: %+v\n", self.Name, frontend)

    if err := self.driverFrontend.add(frontend); err != nil {
        self.driverError(err)
    } else {
        for backendName, backend := range self.Backends {
            self.newBackend(backendName, backend)
        }
    }
}

func (self *Service) setFrontend(frontend config.ServiceFrontend) {
    log.Printf("clusterf:Service %s: set Frontend: %+v\n", self.Name, frontend)

    if self.Frontend != nil {
        // TODO: smoother setup-before-teardown transition..?
        self.delFrontend()
    }

    self.newFrontend(frontend)
}

func (self *Service) delFrontend() {
    log.Printf("clusterf:Service %s: del Frontend: %+v\n", self.Name, self.Frontend)

    // del'ing the frontend will also remove all backend state
    if err := self.driverFrontend.del(); err != nil {
        self.driverError(err)
    }

    // XXX: clear self.driverBackend[*]?
}

/* Backend actions */
func (self *Service) newBackend(backendName string, backend config.ServiceBackend) {
    log.Printf("clusterf:Service %s: new Backend %s: %+v\n", self.Name, backendName, backend)

    self.driverBackends[backendName] = self.driverFrontend.newBackend()

    if err := self.driverBackends[backendName].add(backend); err != nil {
        self.driverError(err)
    }
}

func (self *Service) setBackend(backendName string, backend config.ServiceBackend) {
    log.Printf("clusterf:Service %s: set Backend %s: %+v\n", self.Name, backendName, backend)

    if driverBackend := self.driverBackends[backendName]; driverBackend == nil {
        self.newBackend(backendName, backend)
    } else if err := driverBackend.set(backend); err != nil {
        self.driverError(err)
    }
}

func (self *Service) delBackend(backendName string) {
    log.Printf("clusterf:Service %s: del Backend %s: %+v\n", self.Name, backendName, self.Backends[backendName])

    if err := self.driverBackends[backendName].del(); err != nil {
        self.driverError(err)
    }

    delete(self.driverBackends, backendName)
}
<|endoftext|>


<commit_before>/**
 * Copyright 2013-2016 Seagate Technology LLC.
 *
 * This Source Code Form is subject to the terms of the Mozilla
 * Public License, v. 2.0. If a copy of the MPL was not
 * distributed with this file, You can obtain one at
 * https://mozilla.org/MP:/2.0/.
 *
 * This program is distributed in the hope that it will be useful,
 * but is provided AS-IS, WITHOUT ANY WARRANTY; including without
 * the implied warranty of MERCHANTABILITY, NON-INFRINGEMENT or
 * FITNESS FOR A PARTICULAR PURPOSE. See the Mozilla Public
 * License for more details.
 *
 * See www.openkinetic.org for more project information
 */

package kinetic

import (
	"crypto/tls"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	kproto "github.com/Kinetic/kinetic-go/proto"
	proto "github.com/golang/protobuf/proto"
)

const (
	defaultConnectionTimeout = 20 * time.Second
	defaultRequestTimeout    = 60 * time.Second
)

var (
	connectionTimeout = defaultConnectionTimeout
	requestTimeout    = defaultRequestTimeout
)

func newMessage(t kproto.Message_AuthType) *kproto.Message {
	msg := &kproto.Message{
		AuthType: t.Enum(),
	}
	if msg.GetAuthType() == kproto.Message_HMACAUTH {
		msg.HmacAuth = &kproto.Message_HMACauth{}
	}

	return msg
}

func newCommand(t kproto.Command_MessageType) *kproto.Command {
	return &kproto.Command{
		Header: &kproto.Command_Header{
			MessageType: t.Enum(),
		},
	}
}

type networkService struct {
	rxMu           sync.Mutex
	txMu           sync.Mutex
	mapMu          sync.Mutex
	conn           net.Conn
	clusterVersion int64                      // Cluster version
	seq            int64                      // Operation sequence ID
	connID         int64                      // current connection ID
	option         ClientOptions              // current connection operation
	hmap           map[int64]*ResponseHandler // Message handler map
	fatal          bool                       // Network has fatal failure
	fatalError     error                      // Network fatal error details
	device         Log                        // Store device information from handshake package
}

func newNetworkService(op ClientOptions) (*networkService, error) {
	var conn net.Conn
	var err error

	if op.Timeout > 0 {
		connectionTimeout = time.Duration(op.Timeout) * time.Millisecond
	}
	if op.RequestTimeout > 0 {
		requestTimeout = time.Duration(op.RequestTimeout) * time.Millisecond
	}

	target := fmt.Sprintf("%s:%d", op.Host, op.Port)
	if op.UseSSL {
		// TODO: Need to enable verify certification later
		config := tls.Config{InsecureSkipVerify: true}
		d := &net.Dialer{Timeout: connectionTimeout}
		conn, err = tls.DialWithDialer(d, "tcp", target, &config)
	} else {
		conn, err = net.DialTimeout("tcp", target, connectionTimeout)
	}

	if err != nil {
		klog.Error("Can't establish connection to ", op.Host, err)
		return nil, err
	}

	ns := &networkService{
		conn:           conn,
		clusterVersion: 0,
		seq:            0,
		connID:         -1,
		option:         op,
		hmap:           make(map[int64]*ResponseHandler),
		fatal:          false,
		fatalError:     nil,
	}

	ns.rxMu.Lock()
	// Do the handshake.
	// Device Configuration and Limits from handshake will be stored in networkService.device
	_, _, _, err = ns.receive()
	ns.rxMu.Unlock()

	if err != nil {
		klog.Error("Can't establish connection to %s", op.Host)
		return nil, err
	}

	klog.Debugf("Connected to %s:%d", op.Host, op.Port)
	klog.Debugf("\tVendor: %s", ns.device.Configuration.Vendor)
	klog.Debugf("\tModel: %s", ns.device.Configuration.Model)
	klog.Debugf("\tWorldWideName: %s", ns.device.Configuration.WorldWideName)
	klog.Debugf("\tSerial Number: %s", ns.device.Configuration.SerialNumber)
	klog.Debugf("\tFirmware Version: %s", ns.device.Configuration.Version)
	klog.Debugf("\tKinetic Protocol Version: %s", ns.device.Configuration.ProtocolVersion)
	klog.Debugf("\tPort: %d", ns.device.Configuration.Port)
	klog.Debugf("\tTlsPort: %d", ns.device.Configuration.TLSPort)
	klog.Debugf("\tCurrentPowerLevel : %s", ns.device.Configuration.CurrentPowerLevel.String())

	return ns, nil
}

// When client network service has error, call error handling
// from all Messagehandler current in Queue.
func (ns *networkService) clientError(s Status, mh *ResponseHandler) {
	ns.mapMu.Lock()
	for ack, h := range ns.hmap {
		h.fail(s)
		delete(ns.hmap, ack)
	}
	ns.mapMu.Unlock()

	if mh != nil {
		mh.fail(s)
	}
}

func (ns *networkService) listen() error {
	if ns.fatal {
		return errors.New("Can't listen, network service has fatal error: " + ns.fatalError.Error())
	}

	ns.mapMu.Lock()
	if len(ns.hmap) == 0 {
		ns.mapMu.Unlock()
		return nil
	}
	ns.mapMu.Unlock()

	ns.rxMu.Lock()
	msg, cmd, value, err := ns.receive()
	ns.rxMu.Unlock()
	if err != nil {
		klog.Error("Network Service listen error")
		return err
	}

	if cmd.GetHeader() != nil {
		klog.Debug("Kinetic response received ", cmd.GetHeader().GetMessageType().String(),
			", AckSeq = ", cmd.GetHeader().GetAckSequence(),
			", Code = ", cmd.GetStatus().GetCode())
	} else if msg.GetAuthType() == kproto.Message_UNSOLICITEDSTATUS {
		klog.Debug("Kinetic UNSOLICITEDSTATUS : ",
			"Code = ", cmd.GetStatus().GetCode(),
			", StatusMessage = ", cmd.GetStatus().GetStatusMessage())
	}

	// For UNSOLICITEDSTATUS, command may not have Header or AckSequence, set the ack to -1 so
	// no ResponseHandler will be found from hmap table.
	var ack int64 = -1
	if cmd.GetHeader() != nil {
		ack = cmd.GetHeader().GetAckSequence()
	}

	ns.mapMu.Lock()
	h, ok := ns.hmap[ack]
	ns.mapMu.Unlock()
	if ok == false {
		// It's high chance this is an UNSOLICITEDSTATUS message, display the Status.
		klog.Errorf("Couldn't find a handler for acksequence %d, status=%s", ack, getStatusFromProto(cmd).String())
		// This is an unexpected packet. Each listen() call expect remove one ResponseHandler from hmap.
		// So need to fire another listen() to make sure ResponseHandler in hmap got chance to exit.
		// Either by receive correct packet, or network read failure.
		go ns.listen()
		return nil
	}

	h.handle(cmd, value)

	ns.mapMu.Lock()
	delete(ns.hmap, ack)
	ns.mapMu.Unlock()

	return nil
}

// submit will send the message to kinetic device, insert ResponseHandler for this message sequence number.
// ResponseHandler can be nil if the message no require for Ack, eg batch PUT / DELETE.
func (ns *networkService) submit(msg *kproto.Message, cmd *kproto.Command, value []byte, h *ResponseHandler) error {
	if ns.fatal {
		return errors.New("Can't submit, network service has fatal error: " + ns.fatalError.Error())
	}

	ns.txMu.Lock()

	cmd.GetHeader().ConnectionID = &ns.connID
	cmd.GetHeader().Sequence = &ns.seq
	cmd.GetHeader().ClusterVersion = &ns.clusterVersion

	cmdBytes, err := proto.Marshal(cmd)
	if err != nil {
		klog.Error("Error marshl Kinetic Command")
		s := Status{Code: ClientInternalError, ErrorMsg: "Error marshl Kinetic Command"}
		ns.clientError(s, h)
		return err
	}
	msg.CommandBytes = cmdBytes[:]

	if msg.GetAuthType() == kproto.Message_HMACAUTH {
		msg.GetHmacAuth().Identity = &ns.option.User
		msg.GetHmacAuth().Hmac = computeHmac(msg.CommandBytes, ns.option.Hmac)
	}

	klog.Debug("Kinetic message send ", cmd.GetHeader().GetMessageType().String(), " Seq = ", ns.seq)

	err = ns.send(msg, value)

	if err != nil {
		return err
	}

	if h != nil {
		ns.mapMu.Lock()
		ns.hmap[ns.seq] = h
		ns.mapMu.Unlock()
	}

	ns.seq++
	ns.txMu.Unlock()

	return nil
}

func (ns *networkService) send(msg *kproto.Message, value []byte) error {
	msgBytes, err := proto.Marshal(msg)
	if err != nil {
		s := Status{Code: ClientInternalError, ErrorMsg: "Error marshl Kinetic Message"}
		ns.clientError(s, nil)
		return err
	}

	// Set timeout for send packet
	ns.conn.SetWriteDeadline(time.Now().Add(requestTimeout))

	// Construct message header 9 bytes
	header := make([]byte, 9)
	header[0] = 'F' // Magic number
	binary.BigEndian.PutUint32(header[1:5], uint32(len(msgBytes)))
	binary.BigEndian.PutUint32(header[5:9], uint32(len(value)))

	packet := append(header, msgBytes...)
	if value != nil && len(value) > 0 {
		packet = append(packet, value...)
	}

	_, err = ns.conn.Write(packet)
	if err != nil {
		klog.Error("Network I/O write error, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O write error, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return err
	}

	return nil
}

func (ns *networkService) receive() (*kproto.Message, *kproto.Command, []byte, error) {
	// Set timeout for receive packet
	ns.conn.SetReadDeadline(time.Now().Add(requestTimeout))

	header := make([]byte, 9)

	_, err := io.ReadFull(ns.conn, header[0:])
	if err != nil {
		klog.Error("Network I/O read error, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	magic := header[0]
	if magic != 'F' {
		klog.Error("Network I/O read error Header wrong magic")
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error Header wrong magic"}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = errors.New("Wrong magic number")
		return nil, nil, nil, errors.New("Network I/O read error Header wrong magic")
	}

	protoLen := int(binary.BigEndian.Uint32(header[1:5]))
	valueLen := int(binary.BigEndian.Uint32(header[5:9]))

	protoBuf := make([]byte, protoLen)
	_, err = io.ReadFull(ns.conn, protoBuf)
	if err != nil {
		klog.Error("Network I/O read error receive Kinetic Header, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error receive Kinetic Header, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	msg := &kproto.Message{}
	err = proto.Unmarshal(protoBuf, msg)
	if err != nil {
		klog.Error("Network I/O read error receive Kinetic Header, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error reaceive Kinetic Message, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	if msg.GetAuthType() == kproto.Message_HMACAUTH && validateHmac(msg, ns.option.Hmac) == false {
		klog.Error("Response HMAC mismatch")
		s := Status{Code: ClientResponseHMACError, ErrorMsg: "Response HMAC mismatch"}
		ns.clientError(s, nil)
		return nil, nil, nil, err
	}

	cmd := &kproto.Command{}
	err = proto.Unmarshal(msg.CommandBytes, cmd)
	if err != nil {
		klog.Error("Network I/O read error parsing Kinetic Command, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error parsing Kinetic Command, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	if cmd.Header != nil && cmd.Header.ConnectionID != nil {
		if ns.connID < 0 {
			// This is handshake packet
			ns.device = getLogFromProto(cmd)

			// Only update client cluster version during Handshake
			if cmd.Header.ClusterVersion != nil {
				ns.clusterVersion = cmd.GetHeader().GetClusterVersion()
			}
		}
		ns.connID = cmd.GetHeader().GetConnectionID()
	}

	if valueLen > 0 {
		valueBuf := make([]byte, valueLen)
		_, err = io.ReadFull(ns.conn, valueBuf)
		if err != nil {
			klog.Error("Network I/O read error parsing Kinetic Value, " + err.Error())
			s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error parsing Kinetic Value, " + err.Error()}
			ns.clientError(s, nil)
			ns.fatal = true
			ns.fatalError = err
			return nil, nil, nil, err
		}

		return msg, cmd, valueBuf, nil
	}

	return msg, cmd, nil, nil
}

func (ns *networkService) close() {
	ns.conn.Close()
	klog.Debugf("Connection to %s closed", ns.option.Host)
}
<commit_msg>export default connection and operation timeout constant<commit_after>/**
 * Copyright 2013-2016 Seagate Technology LLC.
 *
 * This Source Code Form is subject to the terms of the Mozilla
 * Public License, v. 2.0. If a copy of the MPL was not
 * distributed with this file, You can obtain one at
 * https://mozilla.org/MP:/2.0/.
 *
 * This program is distributed in the hope that it will be useful,
 * but is provided AS-IS, WITHOUT ANY WARRANTY; including without
 * the implied warranty of MERCHANTABILITY, NON-INFRINGEMENT or
 * FITNESS FOR A PARTICULAR PURPOSE. See the Mozilla Public
 * License for more details.
 *
 * See www.openkinetic.org for more project information
 */

package kinetic

import (
	"crypto/tls"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	kproto "github.com/Kinetic/kinetic-go/proto"
	proto "github.com/golang/protobuf/proto"
)

const (
	// DefaultConnectionTimeout is default timeout in seconds
	// to make network connections to Kinetic devices.
	DefaultConnectionTimeout = 10 * time.Second
	// DefaultRequestTimeout is default timeout in seconds
	// for any Kinetic API call. If Kinetic device no response
	// within this time period, timeout error will trigger.
	DefaultRequestTimeout = 50 * time.Second
)

var (
	connectionTimeout = DefaultConnectionTimeout
	requestTimeout    = DefaultRequestTimeout
)

func newMessage(t kproto.Message_AuthType) *kproto.Message {
	msg := &kproto.Message{
		AuthType: t.Enum(),
	}
	if msg.GetAuthType() == kproto.Message_HMACAUTH {
		msg.HmacAuth = &kproto.Message_HMACauth{}
	}

	return msg
}

func newCommand(t kproto.Command_MessageType) *kproto.Command {
	return &kproto.Command{
		Header: &kproto.Command_Header{
			MessageType: t.Enum(),
		},
	}
}

type networkService struct {
	rxMu           sync.Mutex
	txMu           sync.Mutex
	mapMu          sync.Mutex
	conn           net.Conn
	clusterVersion int64                      // Cluster version
	seq            int64                      // Operation sequence ID
	connID         int64                      // current connection ID
	option         ClientOptions              // current connection operation
	hmap           map[int64]*ResponseHandler // Message handler map
	fatal          bool                       // Network has fatal failure
	fatalError     error                      // Network fatal error details
	device         Log                        // Store device information from handshake package
}

func newNetworkService(op ClientOptions) (*networkService, error) {
	var conn net.Conn
	var err error

	if op.Timeout > 0 {
		connectionTimeout = time.Duration(op.Timeout) * time.Second
	}
	if op.RequestTimeout > 0 {
		requestTimeout = time.Duration(op.RequestTimeout) * time.Second
	}

	target := fmt.Sprintf("%s:%d", op.Host, op.Port)
	if op.UseSSL {
		// TODO: Need to enable verify certification later
		config := tls.Config{InsecureSkipVerify: true}
		d := &net.Dialer{Timeout: connectionTimeout}
		conn, err = tls.DialWithDialer(d, "tcp", target, &config)
	} else {
		conn, err = net.DialTimeout("tcp", target, connectionTimeout)
	}

	if err != nil {
		klog.Error("Can't establish connection to ", op.Host, err)
		return nil, err
	}

	ns := &networkService{
		conn:           conn,
		clusterVersion: 0,
		seq:            0,
		connID:         -1,
		option:         op,
		hmap:           make(map[int64]*ResponseHandler),
		fatal:          false,
		fatalError:     nil,
	}

	ns.rxMu.Lock()
	// Do the handshake.
	// Device Configuration and Limits from handshake will be stored in networkService.device
	_, _, _, err = ns.receive()
	ns.rxMu.Unlock()

	if err != nil {
		klog.Error("Can't establish connection to %s", op.Host)
		return nil, err
	}

	klog.Debugf("Connected to %s:%d", op.Host, op.Port)
	klog.Debugf("    Vendor: %s", ns.device.Configuration.Vendor)
	klog.Debugf("    Model: %s", ns.device.Configuration.Model)
	klog.Debugf("    WorldWideName: %s", ns.device.Configuration.WorldWideName)
	klog.Debugf("    Serial Number: %s", ns.device.Configuration.SerialNumber)
	klog.Debugf("    Firmware Version: %s", ns.device.Configuration.Version)
	klog.Debugf("    Kinetic Protocol Version: %s", ns.device.Configuration.ProtocolVersion)
	klog.Debugf("    Port: %d", ns.device.Configuration.Port)
	klog.Debugf("    TlsPort: %d", ns.device.Configuration.TLSPort)
	klog.Debugf("    CurrentPowerLevel : %s", ns.device.Configuration.CurrentPowerLevel.String())
	klog.Debugf("    Connection Timeout : %d s", connectionTimeout/time.Second)
	klog.Debugf("    Operation Timeout : %d s", requestTimeout/time.Second)

	return ns, nil
}

// When client network service has error, call error handling
// from all Messagehandler current in Queue.
func (ns *networkService) clientError(s Status, mh *ResponseHandler) {
	ns.mapMu.Lock()
	for ack, h := range ns.hmap {
		h.fail(s)
		delete(ns.hmap, ack)
	}
	ns.mapMu.Unlock()

	if mh != nil {
		mh.fail(s)
	}
}

func (ns *networkService) listen() error {
	if ns.fatal {
		return errors.New("Can't listen, network service has fatal error: " + ns.fatalError.Error())
	}

	ns.mapMu.Lock()
	if len(ns.hmap) == 0 {
		ns.mapMu.Unlock()
		return nil
	}
	ns.mapMu.Unlock()

	ns.rxMu.Lock()
	msg, cmd, value, err := ns.receive()
	ns.rxMu.Unlock()
	if err != nil {
		klog.Error("Network Service listen error")
		return err
	}

	if cmd.GetHeader() != nil {
		klog.Debug("Kinetic response received ", cmd.GetHeader().GetMessageType().String(),
			", AckSeq = ", cmd.GetHeader().GetAckSequence(),
			", Code = ", cmd.GetStatus().GetCode())
	} else if msg.GetAuthType() == kproto.Message_UNSOLICITEDSTATUS {
		klog.Debug("Kinetic UNSOLICITEDSTATUS : ",
			"Code = ", cmd.GetStatus().GetCode(),
			", StatusMessage = ", cmd.GetStatus().GetStatusMessage())
	}

	// For UNSOLICITEDSTATUS, command may not have Header or AckSequence, set the ack to -1 so
	// no ResponseHandler will be found from hmap table.
	var ack int64 = -1
	if cmd.GetHeader() != nil {
		ack = cmd.GetHeader().GetAckSequence()
	}

	ns.mapMu.Lock()
	h, ok := ns.hmap[ack]
	ns.mapMu.Unlock()
	if ok == false {
		// It's high chance this is an UNSOLICITEDSTATUS message, display the Status.
		klog.Errorf("Couldn't find a handler for acksequence %d, status=%s", ack, getStatusFromProto(cmd).String())
		// This is an unexpected packet. Each listen() call expect remove one ResponseHandler from hmap.
		// So need to fire another listen() to make sure ResponseHandler in hmap got chance to exit.
		// Either by receive correct packet, or network read failure.
		go ns.listen()
		return nil
	}

	h.handle(cmd, value)

	ns.mapMu.Lock()
	delete(ns.hmap, ack)
	ns.mapMu.Unlock()

	return nil
}

// submit will send the message to kinetic device, insert ResponseHandler for this message sequence number.
// ResponseHandler can be nil if the message no require for Ack, eg batch PUT / DELETE.
func (ns *networkService) submit(msg *kproto.Message, cmd *kproto.Command, value []byte, h *ResponseHandler) error {
	if ns.fatal {
		return errors.New("Can't submit, network service has fatal error: " + ns.fatalError.Error())
	}

	ns.txMu.Lock()

	cmd.GetHeader().ConnectionID = &ns.connID
	cmd.GetHeader().Sequence = &ns.seq
	cmd.GetHeader().ClusterVersion = &ns.clusterVersion

	cmdBytes, err := proto.Marshal(cmd)
	if err != nil {
		klog.Error("Error marshl Kinetic Command")
		s := Status{Code: ClientInternalError, ErrorMsg: "Error marshl Kinetic Command"}
		ns.clientError(s, h)
		return err
	}
	msg.CommandBytes = cmdBytes[:]

	if msg.GetAuthType() == kproto.Message_HMACAUTH {
		msg.GetHmacAuth().Identity = &ns.option.User
		msg.GetHmacAuth().Hmac = computeHmac(msg.CommandBytes, ns.option.Hmac)
	}

	klog.Debug("Kinetic message send ", cmd.GetHeader().GetMessageType().String(), " Seq = ", ns.seq)

	err = ns.send(msg, value)

	if err != nil {
		return err
	}

	if h != nil {
		ns.mapMu.Lock()
		ns.hmap[ns.seq] = h
		ns.mapMu.Unlock()
	}

	ns.seq++
	ns.txMu.Unlock()

	return nil
}

func (ns *networkService) send(msg *kproto.Message, value []byte) error {
	msgBytes, err := proto.Marshal(msg)
	if err != nil {
		s := Status{Code: ClientInternalError, ErrorMsg: "Error marshl Kinetic Message"}
		ns.clientError(s, nil)
		return err
	}

	// Set timeout for send packet
	ns.conn.SetWriteDeadline(time.Now().Add(requestTimeout))

	// Construct message header 9 bytes
	header := make([]byte, 9)
	header[0] = 'F' // Magic number
	binary.BigEndian.PutUint32(header[1:5], uint32(len(msgBytes)))
	binary.BigEndian.PutUint32(header[5:9], uint32(len(value)))

	packet := append(header, msgBytes...)
	if value != nil && len(value) > 0 {
		packet = append(packet, value...)
	}

	_, err = ns.conn.Write(packet)
	if err != nil {
		klog.Error("Network I/O write error, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O write error, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return err
	}

	return nil
}

func (ns *networkService) receive() (*kproto.Message, *kproto.Command, []byte, error) {
	// Set timeout for receive packet
	ns.conn.SetReadDeadline(time.Now().Add(requestTimeout))

	header := make([]byte, 9)

	_, err := io.ReadFull(ns.conn, header[0:])
	if err != nil {
		klog.Error("Network I/O read error, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	magic := header[0]
	if magic != 'F' {
		klog.Error("Network I/O read error Header wrong magic")
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error Header wrong magic"}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = errors.New("Wrong magic number")
		return nil, nil, nil, errors.New("Network I/O read error Header wrong magic")
	}

	protoLen := int(binary.BigEndian.Uint32(header[1:5]))
	valueLen := int(binary.BigEndian.Uint32(header[5:9]))

	protoBuf := make([]byte, protoLen)
	_, err = io.ReadFull(ns.conn, protoBuf)
	if err != nil {
		klog.Error("Network I/O read error receive Kinetic Header, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error receive Kinetic Header, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	msg := &kproto.Message{}
	err = proto.Unmarshal(protoBuf, msg)
	if err != nil {
		klog.Error("Network I/O read error receive Kinetic Header, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error reaceive Kinetic Message, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	if msg.GetAuthType() == kproto.Message_HMACAUTH && validateHmac(msg, ns.option.Hmac) == false {
		klog.Error("Response HMAC mismatch")
		s := Status{Code: ClientResponseHMACError, ErrorMsg: "Response HMAC mismatch"}
		ns.clientError(s, nil)
		return nil, nil, nil, err
	}

	cmd := &kproto.Command{}
	err = proto.Unmarshal(msg.CommandBytes, cmd)
	if err != nil {
		klog.Error("Network I/O read error parsing Kinetic Command, " + err.Error())
		s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error parsing Kinetic Command, " + err.Error()}
		ns.clientError(s, nil)
		ns.fatal = true
		ns.fatalError = err
		return nil, nil, nil, err
	}

	if cmd.Header != nil && cmd.Header.ConnectionID != nil {
		if ns.connID < 0 {
			// This is handshake packet
			ns.device = getLogFromProto(cmd)

			// Only update client cluster version during Handshake
			if cmd.Header.ClusterVersion != nil {
				ns.clusterVersion = cmd.GetHeader().GetClusterVersion()
			}
		}
		ns.connID = cmd.GetHeader().GetConnectionID()
	}

	if valueLen > 0 {
		valueBuf := make([]byte, valueLen)
		_, err = io.ReadFull(ns.conn, valueBuf)
		if err != nil {
			klog.Error("Network I/O read error parsing Kinetic Value, " + err.Error())
			s := Status{Code: ClientIOError, ErrorMsg: "Network I/O read error parsing Kinetic Value, " + err.Error()}
			ns.clientError(s, nil)
			ns.fatal = true
			ns.fatalError = err
			return nil, nil, nil, err
		}

		return msg, cmd, valueBuf, nil
	}

	return msg, cmd, nil, nil
}

func (ns *networkService) close() {
	ns.conn.Close()
	klog.Debugf("Connection to %s closed", ns.option.Host)
}
<|endoftext|>


<commit_before>/*
 *    Copyright (C) 2017 Christian Muehlhaeuser
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU Affero General Public License as published
 *    by the Free Software Foundation, either version 3 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU Affero General Public License for more details.
 *
 *    You should have received a copy of the GNU Affero General Public License
 *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *    Authors:
 *      Christian Muehlhaeuser <muesli@gmail.com>
 */

// Package templatehelper provides a func-map of common template functions
package templatehelper

import (
	"encoding/json"
	htmlTemplate "html/template"
	"regexp"
	"strings"
	"text/template"
)

// FuncMap contains all the common string helpers
var (
	FuncMap = template.FuncMap{
		"JSON": func(values ...interface{}) htmlTemplate.JS {
			json, _ := json.Marshal(values)
			return htmlTemplate.JS(json)
		},
		"Left": func(values ...interface{}) string {
			return values[0].(string)[:values[1].(int)]
		},
		"Matches": func(values ...interface{}) bool {
			ok, _ := regexp.MatchString(values[1].(string), values[0].(string))
			return ok
		},
		"Mid": func(values ...interface{}) string {
			if len(values) > 2 {
				return values[0].(string)[values[1].(int):values[2].(int)]
			}
			return values[0].(string)[values[1].(int):]
		},
		"Right": func(values ...interface{}) string {
			return values[0].(string)[len(values[0].(string))-values[1].(int):]
		},
		"Last": func(values ...interface{}) string {
			return values[0].([]string)[len(values[0].([]string))-1]
		},
		// strings functions
		"Compare":      strings.Compare, // 1.5+ only
		"Contains":     strings.Contains,
		"ContainsAny":  strings.ContainsAny,
		"Count":        strings.Count,
		"EqualFold":    strings.EqualFold,
		"HasPrefix":    strings.HasPrefix,
		"HasSuffix":    strings.HasSuffix,
		"Index":        strings.Index,
		"IndexAny":     strings.IndexAny,
		"Join":         strings.Join,
		"LastIndex":    strings.LastIndex,
		"LastIndexAny": strings.LastIndexAny,
		"Repeat":       strings.Repeat,
		"Replace":      strings.Replace,
		"Split":        strings.Split,
		"SplitAfter":   strings.SplitAfter,
		"SplitAfterN":  strings.SplitAfterN,
		"SplitN":       strings.SplitN,
		"Title":        strings.Title,
		"ToLower":      strings.ToLower,
		"ToTitle":      strings.ToTitle,
		"ToUpper":      strings.ToUpper,
		"Trim":         strings.Trim,
		"TrimLeft":     strings.TrimLeft,
		"TrimPrefix":   strings.TrimPrefix,
		"TrimRight":    strings.TrimRight,
		"TrimSpace":    strings.TrimSpace,
		"TrimSuffix":   strings.TrimSuffix,
	}
)
<commit_msg>Make usage of Left/Mid/Right template helpers safer by checking string boundaries<commit_after>/*
 *    Copyright (C) 2017-2019 Christian Muehlhaeuser
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU Affero General Public License as published
 *    by the Free Software Foundation, either version 3 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU Affero General Public License for more details.
 *
 *    You should have received a copy of the GNU Affero General Public License
 *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *    Authors:
 *      Christian Muehlhaeuser <muesli@gmail.com>
 */

// Package templatehelper provides a func-map of common template functions
package templatehelper

import (
	"encoding/json"
	htmlTemplate "html/template"
	"regexp"
	"strings"
	"text/template"
)

// FuncMap contains all the common string helpers
var (
	FuncMap = template.FuncMap{
		"JSON": func(values ...interface{}) htmlTemplate.JS {
			json, _ := json.Marshal(values)
			return htmlTemplate.JS(json)
		},
		"Left": func(values ...interface{}) string {
			s := values[0].(string)
			n := values[1].(int)
			if n > len(s) {
				n = len(s)
			}

			return s[:n]
		},
		"Matches": func(values ...interface{}) bool {
			ok, _ := regexp.MatchString(values[1].(string), values[0].(string))
			return ok
		},
		"Mid": func(values ...interface{}) string {
			s := values[0].(string)
			l := values[1].(int)
			if l > len(s) {
				l = len(s)
			}

			if len(values) > 2 {
				r := values[2].(int)
				if r > len(s) {
					r = len(s)
				}
				return s[l:r]
			}
			return s[l:]
		},
		"Right": func(values ...interface{}) string {
			s := values[0].(string)
			n := len(s) - values[1].(int)
			if n < 0 {
				n = 0
			}

			return s[n:]
		},
		"Last": func(values ...interface{}) string {
			return values[0].([]string)[len(values[0].([]string))-1]
		},
		// strings functions
		"Compare":      strings.Compare, // 1.5+ only
		"Contains":     strings.Contains,
		"ContainsAny":  strings.ContainsAny,
		"Count":        strings.Count,
		"EqualFold":    strings.EqualFold,
		"HasPrefix":    strings.HasPrefix,
		"HasSuffix":    strings.HasSuffix,
		"Index":        strings.Index,
		"IndexAny":     strings.IndexAny,
		"Join":         strings.Join,
		"LastIndex":    strings.LastIndex,
		"LastIndexAny": strings.LastIndexAny,
		"Repeat":       strings.Repeat,
		"Replace":      strings.Replace,
		"Split":        strings.Split,
		"SplitAfter":   strings.SplitAfter,
		"SplitAfterN":  strings.SplitAfterN,
		"SplitN":       strings.SplitN,
		"Title":        strings.Title,
		"ToLower":      strings.ToLower,
		"ToTitle":      strings.ToTitle,
		"ToUpper":      strings.ToUpper,
		"Trim":         strings.Trim,
		"TrimLeft":     strings.TrimLeft,
		"TrimPrefix":   strings.TrimPrefix,
		"TrimRight":    strings.TrimRight,
		"TrimSpace":    strings.TrimSpace,
		"TrimSuffix":   strings.TrimSuffix,
	}
)
<|endoftext|>