<commit_before>package noise

import (
	"errors"
	"io"
)

type CipherState struct {
	cs CipherSuite
	c  Cipher
	k  [32]byte
	n  uint64
}

func (s *CipherState) Encrypt(out, ad, plaintext []byte) []byte {
	out = s.c.Encrypt(out, s.n, ad, plaintext)
	s.n++
	return out
}

func (s *CipherState) Decrypt(out, ad, ciphertext []byte) ([]byte, error) {
	out, err := s.c.Decrypt(out, s.n, ad, ciphertext)
	s.n++
	return out, err
}

type SymmetricState struct {
	CipherState
	hasK   bool
	hasPSK bool
	ck     []byte
	h      []byte
}

func (s *SymmetricState) InitializeSymmetric(handshakeName []byte) {
	h := s.cs.Hash()
	if len(handshakeName) <= h.Size() {
		s.h = make([]byte, h.Size())
		copy(s.h, handshakeName)
	} else {
		h.Write(handshakeName)
		s.h = h.Sum(nil)
	}
	s.ck = make([]byte, len(s.h))
	copy(s.ck, s.h)
}

func (s *SymmetricState) MixKey(dhOutput []byte) {
	s.n = 0
	s.hasK = true
	var hk []byte
	s.ck, hk = HKDF(s.cs.Hash, s.ck[:0], s.k[:0], s.ck, dhOutput)
	copy(s.k[:], hk)
	s.c = s.cs.Cipher(s.k)
}

func (s *SymmetricState) MixHash(data []byte) {
	h := s.cs.Hash()
	h.Write(s.h)
	h.Write(data)
	s.h = h.Sum(s.h[:0])
}

func (s *SymmetricState) MixPresharedKey(presharedKey []byte) {
	var temp []byte
	s.ck, temp = HKDF(s.cs.Hash, s.ck[:0], nil, s.ck, presharedKey)
	s.MixHash(temp)
	s.hasPSK = true
}

func (s *SymmetricState) EncryptAndHash(out, plaintext []byte) []byte {
	if !s.hasK {
		s.MixHash(plaintext)
		return append(out, plaintext...)
	}
	ciphertext := s.Encrypt(out, s.h, plaintext)
	s.MixHash(ciphertext[len(out):])
	return ciphertext
}

func (s *SymmetricState) DecryptAndHash(out, data []byte) ([]byte, error) {
	if !s.hasK {
		s.MixHash(data)
		return append(out, data...), nil
	}
	plaintext, err := s.Decrypt(out, s.h, data)
	if err != nil {
		return nil, err
	}
	s.MixHash(data)
	return plaintext, nil
}

func (s *SymmetricState) Split() (*CipherState, *CipherState) {
	s1, s2 := &CipherState{cs: s.cs}, &CipherState{cs: s.cs}
	hk1, hk2 := HKDF(s.cs.Hash, s1.k[:0], s2.k[:0], s.ck, nil)
	copy(s1.k[:], hk1)
	copy(s2.k[:], hk2)
	s1.c = s.cs.Cipher(s1.k)
	s2.c = s.cs.Cipher(s2.k)
	return s1, s2
}

type MessagePattern int

type HandshakePattern struct {
	Name                 string
	InitiatorPreMessages []MessagePattern
	ResponderPreMessages []MessagePattern
	Messages             [][]MessagePattern
}

const (
	MessagePatternS MessagePattern = iota
	MessagePatternE
	MessagePatternDHEE
	MessagePatternDHES
	MessagePatternDHSE
	MessagePatternDHSS
)

const MaxMsgLen = 65535

type HandshakeState struct {
	SymmetricState
	s               DHKey  // local static keypair
	e               DHKey  // local ephemeral keypair
	rs              []byte // remote party's static public key
	re              []byte // remote party's ephemeral public key
	messagePatterns [][]MessagePattern
	shouldWrite     bool
	msgIdx          int
	rng             io.Reader
}

type Config struct {
	CipherSuite      CipherSuite
	Random           io.Reader
	Pattern          HandshakePattern
	Initiator        bool
	Prologue         []byte
	PresharedKey     []byte
	StaticKeypair    DHKey
	EphemeralKeypair DHKey
	PeerStatic       []byte
	PeerEphemeral    []byte
}

func NewHandshakeState(c Config) *HandshakeState {
	hs := &HandshakeState{
		s:               c.StaticKeypair,
		e:               c.EphemeralKeypair,
		rs:              c.PeerStatic,
		re:              c.PeerEphemeral,
		messagePatterns: c.Pattern.Messages,
		shouldWrite:     c.Initiator,
		rng:             c.Random,
	}
	hs.SymmetricState.cs = c.CipherSuite
	namePrefix := "Noise_"
	if len(c.PresharedKey) > 0 {
		namePrefix = "NoisePSK_"
	}
	hs.InitializeSymmetric([]byte(namePrefix + c.Pattern.Name + "_" + string(hs.cs.Name())))
	hs.MixHash(c.Prologue)
	if len(c.PresharedKey) > 0 {
		hs.MixPresharedKey(c.PresharedKey)
	}
	for _, m := range c.Pattern.InitiatorPreMessages {
		switch {
		case c.Initiator && m == MessagePatternS:
			hs.MixHash(hs.s.Public)
		case c.Initiator && m == MessagePatternE:
			hs.MixHash(hs.e.Public)
		case !c.Initiator && m == MessagePatternS:
			hs.MixHash(hs.rs)
		case !c.Initiator && m == MessagePatternE:
			hs.MixHash(hs.re)
		}
	}
	for _, m := range c.Pattern.ResponderPreMessages {
		switch {
		case !c.Initiator && m == MessagePatternS:
			hs.MixHash(hs.s.Public)
		case !c.Initiator && m == MessagePatternE:
			hs.MixHash(hs.e.Public)
		case c.Initiator && m == MessagePatternS:
			hs.MixHash(hs.rs)
		case c.Initiator && m == MessagePatternE:
			hs.MixHash(hs.re)
		}
	}
	return hs
}

func (s *HandshakeState) WriteMessage(out, payload []byte) ([]byte, *CipherState, *CipherState) {
	if !s.shouldWrite {
		panic("noise: unexpected call to WriteMessage should be ReadMessage")
	}
	if s.msgIdx > len(s.messagePatterns)-1 {
		panic("noise: no handshake messages left")
	}
	if len(payload) > MaxMsgLen {
		panic("noise: message is too long")
	}

	for _, msg := range s.messagePatterns[s.msgIdx] {
		switch msg {
		case MessagePatternE:
			s.e = s.cs.GenerateKeypair(s.rng)
			out = append(out, s.e.Public...)
			s.MixHash(s.e.Public)
			if s.hasPSK {
				s.MixKey(s.e.Public)
			}
		case MessagePatternS:
			if len(s.s.Public) == 0 {
				panic("noise: invalid state, s.Public is nil")
			}
			out = s.EncryptAndHash(out, s.s.Public)
		case MessagePatternDHEE:
			s.MixKey(s.cs.DH(s.e.Private, s.re))
		case MessagePatternDHES:
			s.MixKey(s.cs.DH(s.e.Private, s.rs))
		case MessagePatternDHSE:
			s.MixKey(s.cs.DH(s.s.Private, s.re))
		case MessagePatternDHSS:
			s.MixKey(s.cs.DH(s.s.Private, s.rs))
		}
	}
	s.shouldWrite = false
	s.msgIdx++
	out = s.EncryptAndHash(out, payload)

	if s.msgIdx >= len(s.messagePatterns) {
		cs1, cs2 := s.Split()
		return out, cs1, cs2
	}

	return out, nil, nil
}

var ErrShortMessage = errors.New("noise: message is too short")

func (s *HandshakeState) ReadMessage(out, message []byte) ([]byte, *CipherState, *CipherState, error) {
	if s.shouldWrite {
		panic("noise: unexpected call to ReadMessage should be WriteMessage")
	}
	if s.msgIdx > len(s.messagePatterns)-1 {
		panic("noise: no handshake messages left")
	}

	var err error
	for _, msg := range s.messagePatterns[s.msgIdx] {
		switch msg {
		case MessagePatternE, MessagePatternS:
			expected := s.cs.DHLen()
			if msg == MessagePatternS && s.hasK {
				expected += 16
			}
			if len(message) < expected {
				return nil, nil, nil, ErrShortMessage
			}
			switch msg {
			case MessagePatternE:
				if cap(s.re) < s.cs.DHLen() {
					s.re = make([]byte, s.cs.DHLen())
				}
				s.re = s.re[:s.cs.DHLen()]
				copy(s.re, message)
				s.MixHash(s.re)
				if s.hasPSK {
					s.MixKey(s.re)
				}
			case MessagePatternS:
				if len(s.rs) > 0 {
					panic("noise: invalid state, rs is not nil")
				}
				s.rs, err = s.DecryptAndHash(s.rs[:0], message[:expected])
			}
			if err != nil {
				return nil, nil, nil, err
			}
			message = message[expected:]
		case MessagePatternDHEE:
			s.MixKey(s.cs.DH(s.e.Private, s.re))
		case MessagePatternDHES:
			s.MixKey(s.cs.DH(s.s.Private, s.re))
		case MessagePatternDHSE:
			s.MixKey(s.cs.DH(s.e.Private, s.rs))
		case MessagePatternDHSS:
			s.MixKey(s.cs.DH(s.s.Private, s.rs))
		}
	}
	s.shouldWrite = true
	s.msgIdx++
	out, err = s.DecryptAndHash(out, message)
	if err != nil {
		return nil, nil, nil, err
	}

	if s.msgIdx >= len(s.messagePatterns) {
		cs1, cs2 := s.Split()
		return out, cs1, cs2, nil
	}

	return out, nil, nil, nil
}
<commit_msg>Copy PeerEphemeral to avoid buffer corruption<commit_after>package noise

import (
	"errors"
	"io"
)

type CipherState struct {
	cs CipherSuite
	c  Cipher
	k  [32]byte
	n  uint64
}

func (s *CipherState) Encrypt(out, ad, plaintext []byte) []byte {
	out = s.c.Encrypt(out, s.n, ad, plaintext)
	s.n++
	return out
}

func (s *CipherState) Decrypt(out, ad, ciphertext []byte) ([]byte, error) {
	out, err := s.c.Decrypt(out, s.n, ad, ciphertext)
	s.n++
	return out, err
}

type SymmetricState struct {
	CipherState
	hasK   bool
	hasPSK bool
	ck     []byte
	h      []byte
}

func (s *SymmetricState) InitializeSymmetric(handshakeName []byte) {
	h := s.cs.Hash()
	if len(handshakeName) <= h.Size() {
		s.h = make([]byte, h.Size())
		copy(s.h, handshakeName)
	} else {
		h.Write(handshakeName)
		s.h = h.Sum(nil)
	}
	s.ck = make([]byte, len(s.h))
	copy(s.ck, s.h)
}

func (s *SymmetricState) MixKey(dhOutput []byte) {
	s.n = 0
	s.hasK = true
	var hk []byte
	s.ck, hk = HKDF(s.cs.Hash, s.ck[:0], s.k[:0], s.ck, dhOutput)
	copy(s.k[:], hk)
	s.c = s.cs.Cipher(s.k)
}

func (s *SymmetricState) MixHash(data []byte) {
	h := s.cs.Hash()
	h.Write(s.h)
	h.Write(data)
	s.h = h.Sum(s.h[:0])
}

func (s *SymmetricState) MixPresharedKey(presharedKey []byte) {
	var temp []byte
	s.ck, temp = HKDF(s.cs.Hash, s.ck[:0], nil, s.ck, presharedKey)
	s.MixHash(temp)
	s.hasPSK = true
}

func (s *SymmetricState) EncryptAndHash(out, plaintext []byte) []byte {
	if !s.hasK {
		s.MixHash(plaintext)
		return append(out, plaintext...)
	}
	ciphertext := s.Encrypt(out, s.h, plaintext)
	s.MixHash(ciphertext[len(out):])
	return ciphertext
}

func (s *SymmetricState) DecryptAndHash(out, data []byte) ([]byte, error) {
	if !s.hasK {
		s.MixHash(data)
		return append(out, data...), nil
	}
	plaintext, err := s.Decrypt(out, s.h, data)
	if err != nil {
		return nil, err
	}
	s.MixHash(data)
	return plaintext, nil
}

func (s *SymmetricState) Split() (*CipherState, *CipherState) {
	s1, s2 := &CipherState{cs: s.cs}, &CipherState{cs: s.cs}
	hk1, hk2 := HKDF(s.cs.Hash, s1.k[:0], s2.k[:0], s.ck, nil)
	copy(s1.k[:], hk1)
	copy(s2.k[:], hk2)
	s1.c = s.cs.Cipher(s1.k)
	s2.c = s.cs.Cipher(s2.k)
	return s1, s2
}

type MessagePattern int

type HandshakePattern struct {
	Name                 string
	InitiatorPreMessages []MessagePattern
	ResponderPreMessages []MessagePattern
	Messages             [][]MessagePattern
}

const (
	MessagePatternS MessagePattern = iota
	MessagePatternE
	MessagePatternDHEE
	MessagePatternDHES
	MessagePatternDHSE
	MessagePatternDHSS
)

const MaxMsgLen = 65535

type HandshakeState struct {
	SymmetricState
	s               DHKey  // local static keypair
	e               DHKey  // local ephemeral keypair
	rs              []byte // remote party's static public key
	re              []byte // remote party's ephemeral public key
	messagePatterns [][]MessagePattern
	shouldWrite     bool
	msgIdx          int
	rng             io.Reader
}

type Config struct {
	CipherSuite      CipherSuite
	Random           io.Reader
	Pattern          HandshakePattern
	Initiator        bool
	Prologue         []byte
	PresharedKey     []byte
	StaticKeypair    DHKey
	EphemeralKeypair DHKey
	PeerStatic       []byte
	PeerEphemeral    []byte
}

func NewHandshakeState(c Config) *HandshakeState {
	hs := &HandshakeState{
		s:               c.StaticKeypair,
		e:               c.EphemeralKeypair,
		rs:              c.PeerStatic,
		messagePatterns: c.Pattern.Messages,
		shouldWrite:     c.Initiator,
		rng:             c.Random,
	}
	if len(c.PeerEphemeral) > 0 {
		hs.re = make([]byte, len(c.PeerEphemeral))
		copy(hs.re, c.PeerEphemeral)
	}
	hs.SymmetricState.cs = c.CipherSuite
	namePrefix := "Noise_"
	if len(c.PresharedKey) > 0 {
		namePrefix = "NoisePSK_"
	}
	hs.InitializeSymmetric([]byte(namePrefix + c.Pattern.Name + "_" + string(hs.cs.Name())))
	hs.MixHash(c.Prologue)
	if len(c.PresharedKey) > 0 {
		hs.MixPresharedKey(c.PresharedKey)
	}
	for _, m := range c.Pattern.InitiatorPreMessages {
		switch {
		case c.Initiator && m == MessagePatternS:
			hs.MixHash(hs.s.Public)
		case c.Initiator && m == MessagePatternE:
			hs.MixHash(hs.e.Public)
		case !c.Initiator && m == MessagePatternS:
			hs.MixHash(hs.rs)
		case !c.Initiator && m == MessagePatternE:
			hs.MixHash(hs.re)
		}
	}
	for _, m := range c.Pattern.ResponderPreMessages {
		switch {
		case !c.Initiator && m == MessagePatternS:
			hs.MixHash(hs.s.Public)
		case !c.Initiator && m == MessagePatternE:
			hs.MixHash(hs.e.Public)
		case c.Initiator && m == MessagePatternS:
			hs.MixHash(hs.rs)
		case c.Initiator && m == MessagePatternE:
			hs.MixHash(hs.re)
		}
	}
	return hs
}

func (s *HandshakeState) WriteMessage(out, payload []byte) ([]byte, *CipherState, *CipherState) {
	if !s.shouldWrite {
		panic("noise: unexpected call to WriteMessage should be ReadMessage")
	}
	if s.msgIdx > len(s.messagePatterns)-1 {
		panic("noise: no handshake messages left")
	}
	if len(payload) > MaxMsgLen {
		panic("noise: message is too long")
	}

	for _, msg := range s.messagePatterns[s.msgIdx] {
		switch msg {
		case MessagePatternE:
			s.e = s.cs.GenerateKeypair(s.rng)
			out = append(out, s.e.Public...)
			s.MixHash(s.e.Public)
			if s.hasPSK {
				s.MixKey(s.e.Public)
			}
		case MessagePatternS:
			if len(s.s.Public) == 0 {
				panic("noise: invalid state, s.Public is nil")
			}
			out = s.EncryptAndHash(out, s.s.Public)
		case MessagePatternDHEE:
			s.MixKey(s.cs.DH(s.e.Private, s.re))
		case MessagePatternDHES:
			s.MixKey(s.cs.DH(s.e.Private, s.rs))
		case MessagePatternDHSE:
			s.MixKey(s.cs.DH(s.s.Private, s.re))
		case MessagePatternDHSS:
			s.MixKey(s.cs.DH(s.s.Private, s.rs))
		}
	}
	s.shouldWrite = false
	s.msgIdx++
	out = s.EncryptAndHash(out, payload)

	if s.msgIdx >= len(s.messagePatterns) {
		cs1, cs2 := s.Split()
		return out, cs1, cs2
	}

	return out, nil, nil
}

var ErrShortMessage = errors.New("noise: message is too short")

func (s *HandshakeState) ReadMessage(out, message []byte) ([]byte, *CipherState, *CipherState, error) {
	if s.shouldWrite {
		panic("noise: unexpected call to ReadMessage should be WriteMessage")
	}
	if s.msgIdx > len(s.messagePatterns)-1 {
		panic("noise: no handshake messages left")
	}

	var err error
	for _, msg := range s.messagePatterns[s.msgIdx] {
		switch msg {
		case MessagePatternE, MessagePatternS:
			expected := s.cs.DHLen()
			if msg == MessagePatternS && s.hasK {
				expected += 16
			}
			if len(message) < expected {
				return nil, nil, nil, ErrShortMessage
			}
			switch msg {
			case MessagePatternE:
				if cap(s.re) < s.cs.DHLen() {
					s.re = make([]byte, s.cs.DHLen())
				}
				s.re = s.re[:s.cs.DHLen()]
				copy(s.re, message)
				s.MixHash(s.re)
				if s.hasPSK {
					s.MixKey(s.re)
				}
			case MessagePatternS:
				if len(s.rs) > 0 {
					panic("noise: invalid state, rs is not nil")
				}
				s.rs, err = s.DecryptAndHash(s.rs[:0], message[:expected])
			}
			if err != nil {
				return nil, nil, nil, err
			}
			message = message[expected:]
		case MessagePatternDHEE:
			s.MixKey(s.cs.DH(s.e.Private, s.re))
		case MessagePatternDHES:
			s.MixKey(s.cs.DH(s.s.Private, s.re))
		case MessagePatternDHSE:
			s.MixKey(s.cs.DH(s.e.Private, s.rs))
		case MessagePatternDHSS:
			s.MixKey(s.cs.DH(s.s.Private, s.rs))
		}
	}
	s.shouldWrite = true
	s.msgIdx++
	out, err = s.DecryptAndHash(out, message)
	if err != nil {
		return nil, nil, nil, err
	}

	if s.msgIdx >= len(s.messagePatterns) {
		cs1, cs2 := s.Split()
		return out, cs1, cs2, nil
	}

	return out, nil, nil, nil
}
<|endoftext|>


<commit_before>package ddtxn

import (
	"errors"
	"flag"
	"log"
	"runtime/debug"
	"sync"
)

type TID uint64
type Key [16]byte
type Value interface{}

type Chunk struct {
	sync.RWMutex
	rows map[Key]*BRecord
	//padding [32]int64
}

var (
	ENOKEY = errors.New("doppel: no key")
	EABORT = errors.New("doppel: abort")
	ESTASH = errors.New("doppel: stash")
)

const (
	CHUNKS = 256
)

// Global data
type Store struct {
	store           []*Chunk
	NChunksAccessed []int64
	dd              []Key
	cand            *Candidates
}

func NewStore() *Store {
	x := make([]*OneStat, 0)
	sh := StatsHeap(x)
	s := &Store{
		store:           make([]*Chunk, CHUNKS),
		NChunksAccessed: make([]int64, CHUNKS),
		dd:              make([]Key, 0, 100),
		cand:            &Candidates{make(map[Key]*OneStat), &sh},
	}
	var bb byte

	for i := 0; i < CHUNKS; i++ {
		chunk := &Chunk{
			rows: make(map[Key]*BRecord),
		}
		bb = byte(i)
		s.store[bb] = chunk
	}
	return s
}

func (s *Store) getOrCreateTypedKey(k Key, v Value, kt KeyType) *BRecord {
	br, err := s.getKey(k)
	if err == ENOKEY {
		// Create key
		chunk := s.store[k[0]]
		var ok bool
		chunk.Lock()
		br, ok = chunk.rows[k]
		if !ok {
			br = MakeBR(k, v, kt)
			chunk.rows[k] = br
		}
		chunk.Unlock()
	}
	return br
}

func (s *Store) CreateKey(k Key, v Value, kt KeyType) *BRecord {
	chunk := s.store[k[0]]
	chunk.Lock()
	br := MakeBR(k, v, kt)
	chunk.rows[k] = br
	chunk.Unlock()
	return br
}

func (s *Store) CreateInt32Key(k Key, v int32, kt KeyType) *BRecord {
	chunk := s.store[k[0]]
	chunk.Lock()
	br := MakeBR(k, v, kt)
	chunk.rows[k] = br
	chunk.Unlock()
	return br
}

func (s *Store) SetInt32(br *BRecord, v int32, op KeyType) {
	switch op {
	case SUM:
		br.int_value += v
	case MAX:
		if v > br.int_value {
			br.int_value = v
		}
	}
}

func (s *Store) Set(br *BRecord, v Value, op KeyType) {
	switch op {
	case SUM:
		br.int_value += v.(int32)
	case MAX:
		x := v.(int32)
		if x > br.int_value {
			br.int_value = v.(int32)
		}
	case WRITE:
		br.value = v
	case LIST:
		if v != nil {
			br.AddOneToRecord(v.(Entry))
		}
	}
}

var UseRLocks = flag.Bool("rlock", true, "Use Rlocks\n")

func (s *Store) Get(k Key) (*BRecord, error) {
	return s.getKey(k)
}

func (s *Store) getKey(k Key) (*BRecord, error) {
	if len(k) == 0 {
		debug.PrintStack()
		log.Fatalf("[store] getKey(): Empty key\n")
	}
	//s.NChunksAccessed[k[0]]++
	if !*UseRLocks {
		x, err := s.getKeyStatic(k)
		return x, err
	}
	chunk := s.store[k[0]]
	if chunk == nil {
		log.Fatalf("[store] Didn't initialize chunk for key %v byte %v\n", k, k[0])
	}
	chunk.RLock()
	vr, ok := chunk.rows[k]
	if !ok {
		chunk.RUnlock()
		return vr, ENOKEY
	}
	chunk.RUnlock()
	return vr, nil
}

func (s *Store) getKeyStatic(k Key) (*BRecord, error) {
	if len(k) == 0 {
		log.Fatalf("[store] getKey(): Empty key\n")
	}
	chunk := s.store[k[0]]
	if chunk == nil {
		log.Fatalf("[store] Didn't initialize chunk for key %v byte %v\n", k, k[0])
	}
	vr, ok := chunk.rows[k]
	if !ok {
		return vr, ENOKEY
	}
	return vr, nil
}

func (s *Store) IsDD(k Key) bool {
	for i := 0; i < len(s.dd); i++ {
		if s.dd[i] == k {
			return true
		}
	}
	return false
}
<commit_msg>check<commit_after>package ddtxn

import (
	"errors"
	"flag"
	"log"
	"runtime/debug"
	"sync"
)

type TID uint64
type Key [16]byte
type Value interface{}

type Chunk struct {
	sync.RWMutex
	rows map[Key]*BRecord
	//padding [32]int64
}

var UseRLocks = flag.Bool("rlock", true, "Use Rlocks\n")

var (
	ENOKEY = errors.New("doppel: no key")
	EABORT = errors.New("doppel: abort")
	ESTASH = errors.New("doppel: stash")
)

const (
	CHUNKS = 256
)

// Global data
type Store struct {
	store           []*Chunk
	NChunksAccessed []int64
	dd              []Key
	cand            *Candidates
}

func NewStore() *Store {
	x := make([]*OneStat, 0)
	sh := StatsHeap(x)
	s := &Store{
		store:           make([]*Chunk, CHUNKS),
		NChunksAccessed: make([]int64, CHUNKS),
		dd:              make([]Key, 0, 100),
		cand:            &Candidates{make(map[Key]*OneStat), &sh},
	}
	var bb byte

	for i := 0; i < CHUNKS; i++ {
		chunk := &Chunk{
			rows: make(map[Key]*BRecord),
		}
		bb = byte(i)
		s.store[bb] = chunk
	}
	return s
}

func (s *Store) getOrCreateTypedKey(k Key, v Value, kt KeyType) *BRecord {
	br, err := s.getKey(k)
	if err == ENOKEY {
		if !*UseRLocks {
			log.Fatalf("Should have preallocated keys if not locking chunks\n")
		}
		// Create key
		chunk := s.store[k[0]]
		var ok bool
		chunk.Lock()
		br, ok = chunk.rows[k]
		if !ok {
			br = MakeBR(k, v, kt)
			chunk.rows[k] = br
		}
		chunk.Unlock()
	}
	return br
}

func (s *Store) CreateKey(k Key, v Value, kt KeyType) *BRecord {
	chunk := s.store[k[0]]
	br := MakeBR(k, v, kt)
	chunk.Lock()
	chunk.rows[k] = br
	chunk.Unlock()
	return br
}

func (s *Store) CreateInt32Key(k Key, v int32, kt KeyType) *BRecord {
	chunk := s.store[k[0]]
	br := MakeBR(k, v, kt)
	chunk.Lock()
	chunk.rows[k] = br
	chunk.Unlock()
	return br
}

func (s *Store) SetInt32(br *BRecord, v int32, op KeyType) {
	switch op {
	case SUM:
		br.int_value += v
	case MAX:
		if v > br.int_value {
			br.int_value = v
		}
	}
}

func (s *Store) Set(br *BRecord, v Value, op KeyType) {
	switch op {
	case SUM:
		br.int_value += v.(int32)
	case MAX:
		x := v.(int32)
		if x > br.int_value {
			br.int_value = v.(int32)
		}
	case WRITE:
		br.value = v
	case LIST:
		if v != nil {
			br.AddOneToRecord(v.(Entry))
		}
	}
}

func (s *Store) Get(k Key) (*BRecord, error) {
	return s.getKey(k)
}

func (s *Store) getKey(k Key) (*BRecord, error) {
	if len(k) == 0 {
		debug.PrintStack()
		log.Fatalf("[store] getKey(): Empty key\n")
	}
	//s.NChunksAccessed[k[0]]++
	if !*UseRLocks {
		x, err := s.getKeyStatic(k)
		return x, err
	}
	chunk := s.store[k[0]]
	if chunk == nil {
		log.Fatalf("[store] Didn't initialize chunk for key %v byte %v\n", k, k[0])
	}
	chunk.RLock()
	vr, ok := chunk.rows[k]
	if !ok {
		chunk.RUnlock()
		return vr, ENOKEY
	}
	chunk.RUnlock()
	return vr, nil
}

func (s *Store) getKeyStatic(k Key) (*BRecord, error) {
	if len(k) == 0 {
		log.Fatalf("[store] getKey(): Empty key\n")
	}
	chunk := s.store[k[0]]
	if chunk == nil {
		log.Fatalf("[store] Didn't initialize chunk for key %v byte %v\n", k, k[0])
	}
	vr, ok := chunk.rows[k]
	if !ok {
		return vr, ENOKEY
	}
	return vr, nil
}

func (s *Store) IsDD(k Key) bool {
	for i := 0; i < len(s.dd); i++ {
		if s.dd[i] == k {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package rest

func newStore() *store {
	s := new(store)
	s.dataMap = map[string][]interface{}{}
	s.idsMap = map[string]map[string]bool{}
	return s
}

type store struct {
	dataMap map[string][]interface{}
	idsMap  map[string]map[string]bool
}

func (self *store) AddId(name string, id string) {
	if nil == self.idsMap[name] {
		self.idsMap[name] = map[string]bool{}
	}

	self.idsMap[name][id] = true
}

func (self *store) empty(name string) {
	delete(self.idsMap, name)
}

func (self *store) AddRecord(name string, data interface{}) {
	self.dataMap[name] = append(self.dataMap[name], data)
}

func (self *store) all(name string) []interface{} {
	if ds, ok := self.dataMap[name]; ok {
		return ds
	}

	return nil
}

func (self *store) DataMap() map[string][]interface{} {
	return self.dataMap
}

func (self *store) fillByIds(ctx IContext) {
	count := len(self.idsMap)
	if 0 == count {
		return
	}

	for apiName, idsMap := range self.idsMap {
		if 0 == count {
			self.empty(apiName)
			return
		}

		api := ctx.Handler().Get(apiName)
		if api == nil {
			continue
		}

		for id, _ := range idsMap {
			api.Fill(ctx, id)
		}

		self.empty(apiName)
		self.fillByIds(ctx)
	}
}
<commit_msg>allow to store meta in store<commit_after>package rest

func newStore() *store {
	s := new(store)
	s.dataMap = map[string][]interface{}{}
	s.idsMap = map[string]map[string]bool{}
	s.meta = map[string]interface{}{}
	return s
}

type store struct {
	dataMap map[string][]interface{}
	idsMap  map[string]map[string]bool
	meta    map[string]interface{}
}

func (self *store) Meta() map[string]interface{} {
	return self.meta
}

func (self *store) AddMeta(name string, val interface{}) {
	self.meta[name] = val
}

func (self *store) AddId(name string, id string) {
	if nil == self.idsMap[name] {
		self.idsMap[name] = map[string]bool{}
	}

	self.idsMap[name][id] = true
}

func (self *store) empty(name string) {
	delete(self.idsMap, name)
}

func (self *store) AddRecord(name string, data interface{}) {
	self.dataMap[name] = append(self.dataMap[name], data)
}

func (self *store) all(name string) []interface{} {
	if ds, ok := self.dataMap[name]; ok {
		return ds
	}

	return nil
}

func (self *store) DataMap() map[string][]interface{} {
	return self.dataMap
}

func (self *store) fillByIds(ctx IContext) {
	count := len(self.idsMap)
	if 0 == count {
		return
	}

	for apiName, idsMap := range self.idsMap {
		if 0 == count {
			self.empty(apiName)
			return
		}

		api := ctx.Handler().Get(apiName)
		if api == nil {
			continue
		}

		for id, _ := range idsMap {
			api.Fill(ctx, id)
		}

		self.empty(apiName)
		self.fillByIds(ctx)
	}
}
<|endoftext|>


<commit_before>package voucher

// Suite is a suite of Checks, which
type Suite struct {
	checks map[string]Check
}

// Add adds a Check to the checks that can be run. Once a Check is added,
// it can be referenced by the name that was passed in when this function was called.
func (cs *Suite) Add(name string, check Check) {
	// Don't replace the existing check if it already exists.
	if !cs.Has(name) {
		cs.checks[name] = check
	}
}

// Has returns true if the passed check exists. Returns false if it does not.
func (cs *Suite) Has(name string) bool {
	return nil != cs.checks[name]
}

// Get returns the requested Check, or nil if one does not exist.
func (cs *Suite) Get(name string) (Check, error) {
	if cs.Has(name) {
		return cs.checks[name], nil
	}
	return nil, ErrNoCheck
}

// Run executes each of the Checks specified by the activeChecks parameter.
//
// For example, if a Suite has the "diy" and "nobody" tests, calling
//
//    Run(imageData)
//
// will run the "diy" and "nobody" tests.
//
// Run returns a []CheckResult with a CheckResult for each Check that was run.
func (cs *Suite) Run(imageData ImageData) []CheckResult {
	results := make([]CheckResult, 0, len(cs.checks))
	for name, check := range cs.checks {
		ok, err := check.Check(imageData)
		if err == nil {
			results = append(results, CheckResult{Name: name, Err: "", Success: ok, ImageData: imageData})
		} else {
			results = append(results, CheckResult{Name: name, Err: err.Error(), Success: false, ImageData: imageData})
		}
	}
	return results
}

// Attest runs through the passed []CheckResult and if a CheckResult is marked as successful,
// runs the CreateAttestion function in the Check corresponding to that CheckResult. Each
// CheckResult is updated with the details (or error) and the resulting []CheckResult is
// returned.
func (cs *Suite) Attest(metadataClient MetadataClient, results []CheckResult) []CheckResult {
	for i, result := range results {
		if result.Success {
			details, err := createAttestation(metadataClient, result)
			results[i].Details = details
			if nil == err {
				results[i].Attested = true
			} else {
				results[i].Err = err.Error()
			}
		}
	}
	return results
}

// RunAndAttest calls Run, followed by Attest, and returns the final []CheckResult.
func (cs *Suite) RunAndAttest(metadataClient MetadataClient, imageData ImageData) []CheckResult {
	results := cs.Run(imageData)
	return cs.Attest(metadataClient, results)
}

// createAttestation generates an attestation for the image Check described by CheckResult.
// That attestation is then added to the metadata server the MetadataClient is connected to.
func createAttestation(client MetadataClient, result CheckResult) (MetadataItem, error) {
	payload, err := client.NewPayloadBody(result.ImageData)
	if err != nil {
		return nil, err
	}

	attestationPayload := NewAttestationPayload(result.Name, payload)
	occ, err := client.AddAttestationToImage(result.ImageData, attestationPayload)
	return occ, err
}

// NewSuite creates a new Suite.
func NewSuite() *Suite {
	suite := new(Suite)
	suite.checks = make(map[string]Check)
	return suite
}
<commit_msg>adds concurrency to Voucher checks (#27)<commit_after>package voucher

// Suite is a suite of Checks, which
type Suite struct {
	checks map[string]Check
}

// Add adds a Check to the checks that can be run. Once a Check is added,
// it can be referenced by the name that was passed in when this function was called.
func (cs *Suite) Add(name string, check Check) {
	// Don't replace the existing check if it already exists.
	if !cs.Has(name) {
		cs.checks[name] = check
	}
}

// Has returns true if the passed check exists. Returns false if it does not.
func (cs *Suite) Has(name string) bool {
	return nil != cs.checks[name]
}

// Get returns the requested Check, or nil if one does not exist.
func (cs *Suite) Get(name string) (Check, error) {
	if cs.Has(name) {
		return cs.checks[name], nil
	}
	return nil, ErrNoCheck
}

// runner runs the passed check against the passed ImageData, and pushes results to the
// CheckResults channel.
func runner(name string, check Check, imageData ImageData, resultsChan chan CheckResult) {
	ok, err := check.Check(imageData)
	if err == nil {
		resultsChan <- CheckResult{Name: name, Err: "", Success: ok, ImageData: imageData}
	} else {
		resultsChan <- CheckResult{Name: name, Err: err.Error(), Success: false, ImageData: imageData}
	}
}

// Run executes each of the Checks specified by the activeChecks parameter.
//
// For example, if a Suite has the "diy" and "nobody" tests, calling
//
//    Run(imageData)
//
// will run the "diy" and "nobody" tests.
//
// Run returns a []CheckResult with a CheckResult for each Check that was run.
func (cs *Suite) Run(imageData ImageData) []CheckResult {
	results := make([]CheckResult, 0, len(cs.checks))
	resultsChan := make(chan CheckResult, len(cs.checks))
	defer close(resultsChan)

	for name, check := range cs.checks {
		go runner(name, check, imageData, resultsChan)
	}

	for range cs.checks {
		results = append(results, <-resultsChan)
	}

	return results
}

// Attest runs through the passed []CheckResult and if a CheckResult is marked as successful,
// runs the CreateAttestion function in the Check corresponding to that CheckResult. Each
// CheckResult is updated with the details (or error) and the resulting []CheckResult is
// returned.
func (cs *Suite) Attest(metadataClient MetadataClient, results []CheckResult) []CheckResult {
	for i, result := range results {
		if result.Success {
			details, err := createAttestation(metadataClient, result)
			results[i].Details = details
			if nil == err {
				results[i].Attested = true
			} else {
				results[i].Err = err.Error()
			}
		}
	}
	return results
}

// RunAndAttest calls Run, followed by Attest, and returns the final []CheckResult.
func (cs *Suite) RunAndAttest(metadataClient MetadataClient, imageData ImageData) []CheckResult {
	results := cs.Run(imageData)
	return cs.Attest(metadataClient, results)
}

// createAttestation generates an attestation for the image Check described by CheckResult.
// That attestation is then added to the metadata server the MetadataClient is connected to.
func createAttestation(client MetadataClient, result CheckResult) (MetadataItem, error) {
	payload, err := client.NewPayloadBody(result.ImageData)
	if err != nil {
		return nil, err
	}

	attestationPayload := NewAttestationPayload(result.Name, payload)
	occ, err := client.AddAttestationToImage(result.ImageData, attestationPayload)
	return occ, err
}

// NewSuite creates a new Suite.
func NewSuite() *Suite {
	suite := new(Suite)
	suite.checks = make(map[string]Check)
	return suite
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"text/template"
)

const data = `package {{.PkgName}}
{{if .DbrUsed}}import "github.com/gocraft/dbr"{{end}}

var fieldsNames = []string{ {{.FieldsNames}} }
{{if .AutoInc}}var auto_increment_field string = "{{.AutoInc}}"{{end}}

type {{.Table}} struct {
	{{range .Fields}}{{.Name}}{{/*tab*/}} {{.Type}}{{/*tab*/}} {{.Tag}}
	{{end}}
}

func New() *{{.Table}} {
	return new({{.Table}})
}

func NewSlice() []*{{.Table}} {
	return make([]*{{.Table}}, 0)
}

func FieldsNames() []string {
	return fieldsNames
}
{{if .AutoInc}}
func FieldsNamesWithOutAI() []string {
	var slice []string
	for _, iterator := range fieldsNames {
		if iterator == auto_increment_field {
			continue
		}
		slice = append(slice, iterator)
	}
	return slice
}
{{end}}
`

type Field struct {
	Name string
	Type string
	Tag  string
}

type TplData struct {
	PkgName     string
	Table       string
	Fields      []Field
	FieldsNames string
	AutoInc     string
	DbrUsed     bool
}

func CreateTableModel(path, table string, db *sql.DB, verbose bool) {
	var (
		name  string
		typ   string
		null  string
		key   string
		def   sql.NullString
		extra string
	)

	template_data := TplData{}
	template_data.PkgName = table
	template_data.Table = strings.Title(table)

	// get table columns info
	q := fmt.Sprintf("SHOW COLUMNS FROM %s", table)
	if rows, err := db.Query(q); err == nil {
		if verbose {
			fmt.Println("\tfields:")
		}
		for rows.Next() {
			err := rows.Scan(&name, &typ, &null, &key, &def, &extra)
			if err != nil {
				fmt.Fprintf(os.Stderr, "rows.Scan: %s\n", err.Error())
				continue
			}
			if verbose {
				fmt.Printf("\t\tname: `%s` type: %s null: %s key: %s def: %s extra: %s\n", name, typ, null, key, def.String, extra)
			}
			titled_name := strings.Title(name)
			if extra == "auto_increment" {
				template_data.AutoInc = titled_name
			}

			if typ == "tinyint(1)" { // bool need be first because next `strings.Contains(typ, "int")`
				if null == "YES" {
					template_data.DbrUsed = true
					typ = "dbr.NullBool"
				} else {
					typ = "bool"
				}
			} else if strings.Contains(typ, "int") {
				if null == "YES" {
					template_data.DbrUsed = true
					typ = "dbr.NullInt64"
				} else {
					typ = "int64"
				}
			} else if strings.Contains(typ, "float") {
				if null == "YES" {
					template_data.DbrUsed = true
					typ = "dbr.NullFloat64"
				} else {
					typ = "float64"
				}
			} else if strings.Contains(typ, "date") || strings.Contains(typ, "timestamp") {
				template_data.DbrUsed = true
				typ = "dbr.NullTime"
			} else {
				if null == "YES" {
					template_data.DbrUsed = true
					typ = "dbr.NullString"
				} else {
					typ = "string"
				}
			}

			tag := fmt.Sprintf("`db:\"%s\"`", name)
			if verbose {
				fmt.Printf("\t\t\t => %s %s %s\n", titled_name, typ, tag)
			}
			table_field := Field{titled_name, typ, tag}
			template_data.Fields = append(template_data.Fields, table_field)
			template_data.FieldsNames = fmt.Sprintf("%s, \"%s\"", template_data.FieldsNames, table_field.Name)
		}
		template_data.FieldsNames = strings.Trim(template_data.FieldsNames, ",")
	}
	t := template.Must(template.New("struct").Parse(data))

	fullPath := path + "/" + table
	fullFileName := fullPath + "/model.go"
	err := os.MkdirAll(fullPath, 0700)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
	}

	file, err := os.Create(fullFileName)
	defer file.Close()
	if err == nil {
		if err := t.Execute(file, template_data); err != nil {
			fmt.Errorf("%s", err.Error())
		}
		cmd := exec.Command("go", "fmt", fullFileName)
		err = cmd.Start()
		if err == nil {
			err = cmd.Wait()
		}
	}
}
<commit_msg>double, real, decimal data types<commit_after>package main

import (
	"database/sql"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"text/template"
)

const data = `package {{.PkgName}}
{{if .DbrUsed}}import "github.com/gocraft/dbr"{{end}}

var fieldsNames = []string{ {{.FieldsNames}} }
{{if .AutoInc}}var auto_increment_field string = "{{.AutoInc}}"{{end}}

type {{.Table}} struct {
	{{range .Fields}}{{.Name}}{{/*tab*/}} {{.Type}}{{/*tab*/}} {{.Tag}}
	{{end}}
}

func New() *{{.Table}} {
	return new({{.Table}})
}

func NewSlice() []*{{.Table}} {
	return make([]*{{.Table}}, 0)
}

func FieldsNames() []string {
	return fieldsNames
}
{{if .AutoInc}}
func FieldsNamesWithOutAI() []string {
	var slice []string
	for _, iterator := range fieldsNames {
		if iterator == auto_increment_field {
			continue
		}
		slice = append(slice, iterator)
	}
	return slice
}
{{end}}
`

type Field struct {
	Name string
	Type string
	Tag  string
}

type TplData struct {
	PkgName     string
	Table       string
	Fields      []Field
	FieldsNames string
	AutoInc     string
	DbrUsed     bool
}

func CreateTableModel(path, table string, db *sql.DB, verbose bool) {
	var (
		name  string
		typ   string
		null  string
		key   string
		def   sql.NullString
		extra string
	)

	template_data := TplData{}
	template_data.PkgName = table
	template_data.Table = strings.Title(table)

	// get table columns info
	q := fmt.Sprintf("SHOW COLUMNS FROM %s", table)
	if rows, err := db.Query(q); err == nil {
		if verbose {
			fmt.Println("\tfields:")
		}
		for rows.Next() {
			err := rows.Scan(&name, &typ, &null, &key, &def, &extra)
			if err != nil {
				fmt.Fprintf(os.Stderr, "rows.Scan: %s\n", err.Error())
				continue
			}
			if verbose {
				fmt.Printf("\t\tname: `%s` type: %s null: %s key: %s def: %s extra: %s\n", name, typ, null, key, def.String, extra)
			}
			titled_name := strings.Title(name)
			if extra == "auto_increment" {
				template_data.AutoInc = titled_name
			}

			if typ == "tinyint(1)" { // bool need be first because next `strings.Contains(typ, "int")`
				if null == "YES" {
					template_data.DbrUsed = true
					typ = "dbr.NullBool"
				} else {
					typ = "bool"
				}
			} else if strings.Contains(typ, "int") {
				if null == "YES" {
					template_data.DbrUsed = true
					typ = "dbr.NullInt64"
				} else {
					typ = "int64"
				}
			} else if strings.Contains(typ, "float") ||
				strings.Contains(typ, "decimal") ||
				strings.Contains(typ, "double") ||
				strings.Contains(typ, "real") {
				if null == "YES" {
					template_data.DbrUsed = true
					typ = "dbr.NullFloat64"
				} else {
					typ = "float64"
				}
			} else if strings.Contains(typ, "date") || strings.Contains(typ, "timestamp") {
				template_data.DbrUsed = true
				typ = "dbr.NullTime"
			} else {
				if null == "YES" {
					template_data.DbrUsed = true
					typ = "dbr.NullString"
				} else {
					typ = "string"
				}
			}

			tag := fmt.Sprintf("`db:\"%s\"`", name)
			if verbose {
				fmt.Printf("\t\t\t => %s %s %s\n", titled_name, typ, tag)
			}
			table_field := Field{titled_name, typ, tag}
			template_data.Fields = append(template_data.Fields, table_field)
			template_data.FieldsNames = fmt.Sprintf("%s, \"%s\"", template_data.FieldsNames, table_field.Name)
		}
		template_data.FieldsNames = strings.Trim(template_data.FieldsNames, ",")
	}
	t := template.Must(template.New("struct").Parse(data))

	fullPath := path + "/" + table
	fullFileName := fullPath + "/model.go"
	err := os.MkdirAll(fullPath, 0700)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
	}

	file, err := os.Create(fullFileName)
	defer file.Close()
	if err == nil {
		if err := t.Execute(file, template_data); err != nil {
			fmt.Errorf("%s", err.Error())
		}
		cmd := exec.Command("go", "fmt", fullFileName)
		err = cmd.Start()
		if err == nil {
			err = cmd.Wait()
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/csv"
	"fmt"
	"log"
	"math/big"
	"os"

	// Register hdb driver.
	_ "github.com/SAP/go-hdb/driver"
	// cli helper
	"github.com/mkideal/cli"
	// ini config
	"github.com/go-ini/ini"
	// internal
	"github.com/morxs/go-hana/utils"
)

type argT struct {
	cli.Helper
	ArgStart  string `cli:"*s" usage:"Start Date (SAP format)"`
	ArgEnd    string `cli:"*e" usage:"End Date (SAP format)"`
	ArgConfig string `cli:"c" usage:"Custom config file" dft:"config.ini"`
}

const (
	driverName = "hdb"
)

const (
	tcurrSQL = `select MANDT,
	   	KURST,
	   	FCURR,
	   	TCURR,
	   	GDATU,
	   	UKURS,
	   	FFACT,
	   	TFACT,
	   	DATUM
	   from
	   (
	   	select
	   	MANDT,
	   	KURST,
	   	FCURR,
	   	TCURR,
	   	GDATU,
		UKURS as "UKURS",
		FFACT as "FFACT",
		TFACT as "TFACT",
	   	cast(('99999999' - gdatu) as varchar(8)) as "DATUM"
	   	from z_wilmar1.tcurr
	   	where mandt = '777'
	   )
	   where "DATUM" between ? and ?`
)

const (
	File = "tcurr.csv"
)

func main() {
	cli.Run(new(argT), func(ctx *cli.Context) error {
		argv := ctx.Argv().(*argT)

		// read config file
		utils.WriteMsg("READ CONFIG")
		iniCfg, err := ini.Load(argv.ArgConfig)
		if err != nil {
			utils.WriteMsg("CONFIG")
			log.Fatal(err)
		}
		iniSection := iniCfg.Section("server")
		/*
			fmt.Println(iniSection)
			fmt.Println(iniSection.KeyStrings())
			fmt.Println(iniSection.Key("uid").String())
			fmt.Println(iniSection.GetKey("uid"))
		*/
		iniKeyUsername := iniSection.Key("uid").String()
		iniKeyPassword := iniSection.Key("pwd").String()
		iniKeyHost := iniSection.Key("host").String()
		iniKeyPort := iniSection.Key("port").String()
		hdbDsn := "hdb://" + iniKeyUsername + ":" + iniKeyPassword + "@" + iniKeyHost + ":" + iniKeyPort

		utils.WriteMsg("OPEN HDB")
		//fmt.Print("OPENDB...")
		db, err := sql.Open(driverName, hdbDsn)
		if err != nil {
			//fmt.Print("OPENDB")
			log.Fatal(err)
		}
		defer db.Close()

		if err := db.Ping(); err != nil {
			log.Fatal(err)
		}

		// create file
		utils.WriteMsg("CREATE FILE: " + File)
		file, err := os.Create(File)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()

		// try to query
		utils.WriteMsg("QUERY")
		rows, err := db.Query(tcurrSQL, argv.ArgStart, argv.ArgEnd)
		if err != nil {
			log.Fatal(err)
		}
		defer rows.Close()

		// prepare file
		utils.WriteMsg("WRITE CSV")
		w := csv.NewWriter(file)
		w.Comma = ';'

		fmt.Println(rows.Columns())
		// add header to file
		rs, _ := rows.Columns()
		var rec []string
		for i := 0; i < len(rs); i++ {
			rec = append(rec, rs[i])
		}
		w.Write(rec)

		for rows.Next() {
			var mandt, kurst, fcurr, tcurr, gdatu, datum string
			//var ukurs, ffact, tfact string
			//var ukurs2, tfact2, ffact2 []byte
			var ukurs, ffact, tfact []byte
			if err := rows.Scan(&mandt, &kurst, &fcurr, &tcurr, &gdatu, &ukurs, &ffact, &tfact, &datum); err != nil {
				//if err := rows.Scan(&mandt, &kurst, &fcurr, &tcurr, &gdatu, ukurs, ffact, tfact, &datum); err != nil {
				//fmt.Print("SCAN")
				utils.WriteMsg("SCAN")
				log.Fatal(err)
			}
			//encoded := base64.StdEncoding.EncodeToString(ukurs)
			//fmt.Println("....>", ukurs2, ffact2, tfact2)
			/*
				var z float64
				buf := bytes.NewReader(ukurs2)
				err := binary.Read(buf, binary.LittleEndian, &z)
				if err != nil {
					fmt.Println("binary.Read failed:", err)
				}
				fmt.Println(z)
			*/
			//fmt.Print(strconv.Itoa(ukurs[1]), strconv.Itoa(ukurs[2]))

			var bi big.Int
			var z float64
			var neg bool
			var i int

			var record []string

			record = append(record, mandt)
			record = append(record, kurst)
			record = append(record, fcurr)
			record = append(record, tcurr)
			record = append(record, gdatu)

			// ukurs
			neg, i = utils.DecodeDecimal(ukurs, &bi)
			//fmt.Println(neg, bi, i)
			z = utils.BigIntToFloat(neg, &bi, i)
			//fmt.Printf("%.4f", z)
			record = append(record, fmt.Sprintf("%.4f", z))

			// ffact
			neg, i = utils.DecodeDecimal(ffact, &bi)
			z = utils.BigIntToFloat(neg, &bi, i)
			record = append(record, fmt.Sprintf("%.4f", z))

			// tfact
			neg, i = utils.DecodeDecimal(tfact, &bi)
			z = utils.BigIntToFloat(neg, &bi, i)
			record = append(record, fmt.Sprintf("%.4f", z))

			//record = append(record, string(ffact2))
			/*
				record = append(record, string(ukurs[:bytes.IndexByte(ukurs, 0)]))
				record = append(record, string(ffact[:bytes.IndexByte(ffact, 0)]))
				record = append(record, string(tfact[:bytes.IndexByte(tfact, 0)]))
			*/
			/*
				record = append(record, strconv.FormatFloat(ukurs, 'f', 4, 64))
				record = append(record, strconv.FormatFloat(ffact, 'f', 4, 64))
				record = append(record, strconv.FormatFloat(tfact, 'f', 4, 64))
			*/
			record = append(record, datum)
			w.Write(record)
			//fmt.Println(record)
			//fmt.Printf("%s;%s;%s;%s;%s;%v;%v;%v;%s\n", mandt, kurst, fcurr, tcurr, gdatu, ukurs, ffact, tfact, datum)
		}
		w.Flush()

		if err := rows.Err(); err != nil {
			//fmt.Print("ROWS")
			utils.WriteMsg("ROWS")
			log.Fatal(err)
		}
		return nil
	})
}
<commit_msg>add new requirement to include only 'M' type<commit_after>package main

import (
	"database/sql"
	"encoding/csv"
	"fmt"
	"log"
	"math/big"
	"os"

	// Register hdb driver.
	_ "github.com/SAP/go-hdb/driver"
	// cli helper
	"github.com/mkideal/cli"
	// ini config
	"github.com/go-ini/ini"
	// internal
	"github.com/morxs/go-hana/utils"
)

type argT struct {
	cli.Helper
	ArgStart  string `cli:"*s" usage:"Start Date (SAP format)"`
	ArgEnd    string `cli:"*e" usage:"End Date (SAP format)"`
	ArgConfig string `cli:"c" usage:"Custom config file" dft:"config.ini"`
}

const (
	driverName = "hdb"
)

const (
	tcurrSQL = `select MANDT,
	   	KURST,
	   	FCURR,
	   	TCURR,
	   	GDATU,
	   	UKURS,
	   	FFACT,
	   	TFACT,
	   	DATUM
	   from
	   (
	   	select
	   	MANDT,
	   	KURST,
	   	FCURR,
	   	TCURR,
	   	GDATU,
		UKURS as "UKURS",
		FFACT as "FFACT",
		TFACT as "TFACT",
	   	cast(('99999999' - gdatu) as varchar(8)) as "DATUM"
	   	from z_wilmar1.tcurr
	   	where mandt = '777'
		and kurst = 'M'
	   )
	   where "DATUM" between ? and ?`
)

const (
	File = "tcurr.csv"
)

func main() {
	cli.Run(new(argT), func(ctx *cli.Context) error {
		argv := ctx.Argv().(*argT)

		// read config file
		utils.WriteMsg("READ CONFIG")
		iniCfg, err := ini.Load(argv.ArgConfig)
		if err != nil {
			utils.WriteMsg("CONFIG")
			log.Fatal(err)
		}
		iniSection := iniCfg.Section("server")
		/*
			fmt.Println(iniSection)
			fmt.Println(iniSection.KeyStrings())
			fmt.Println(iniSection.Key("uid").String())
			fmt.Println(iniSection.GetKey("uid"))
		*/
		iniKeyUsername := iniSection.Key("uid").String()
		iniKeyPassword := iniSection.Key("pwd").String()
		iniKeyHost := iniSection.Key("host").String()
		iniKeyPort := iniSection.Key("port").String()
		hdbDsn := "hdb://" + iniKeyUsername + ":" + iniKeyPassword + "@" + iniKeyHost + ":" + iniKeyPort

		utils.WriteMsg("OPEN HDB")
		//fmt.Print("OPENDB...")
		db, err := sql.Open(driverName, hdbDsn)
		if err != nil {
			//fmt.Print("OPENDB")
			log.Fatal(err)
		}
		defer db.Close()

		if err := db.Ping(); err != nil {
			log.Fatal(err)
		}

		// create file
		utils.WriteMsg("CREATE FILE: " + File)
		file, err := os.Create(File)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()

		// try to query
		utils.WriteMsg("QUERY")
		rows, err := db.Query(tcurrSQL, argv.ArgStart, argv.ArgEnd)
		if err != nil {
			log.Fatal(err)
		}
		defer rows.Close()

		// prepare file
		utils.WriteMsg("WRITE CSV")
		w := csv.NewWriter(file)
		w.Comma = ';'

		fmt.Println(rows.Columns())
		// add header to file
		rs, _ := rows.Columns()
		var rec []string
		for i := 0; i < len(rs); i++ {
			rec = append(rec, rs[i])
		}
		w.Write(rec)

		for rows.Next() {
			var mandt, kurst, fcurr, tcurr, gdatu, datum string
			//var ukurs, ffact, tfact string
			//var ukurs2, tfact2, ffact2 []byte
			var ukurs, ffact, tfact []byte
			if err := rows.Scan(&mandt, &kurst, &fcurr, &tcurr, &gdatu, &ukurs, &ffact, &tfact, &datum); err != nil {
				//if err := rows.Scan(&mandt, &kurst, &fcurr, &tcurr, &gdatu, ukurs, ffact, tfact, &datum); err != nil {
				//fmt.Print("SCAN")
				utils.WriteMsg("SCAN")
				log.Fatal(err)
			}
			//encoded := base64.StdEncoding.EncodeToString(ukurs)
			//fmt.Println("....>", ukurs2, ffact2, tfact2)
			/*
				var z float64
				buf := bytes.NewReader(ukurs2)
				err := binary.Read(buf, binary.LittleEndian, &z)
				if err != nil {
					fmt.Println("binary.Read failed:", err)
				}
				fmt.Println(z)
			*/
			//fmt.Print(strconv.Itoa(ukurs[1]), strconv.Itoa(ukurs[2]))

			var bi big.Int
			var z float64
			var neg bool
			var i int

			var record []string

			record = append(record, mandt)
			record = append(record, kurst)
			record = append(record, fcurr)
			record = append(record, tcurr)
			record = append(record, gdatu)

			// ukurs
			neg, i = utils.DecodeDecimal(ukurs, &bi)
			//fmt.Println(neg, bi, i)
			z = utils.BigIntToFloat(neg, &bi, i)
			//fmt.Printf("%.4f", z)
			record = append(record, fmt.Sprintf("%.4f", z))

			// ffact
			neg, i = utils.DecodeDecimal(ffact, &bi)
			z = utils.BigIntToFloat(neg, &bi, i)
			record = append(record, fmt.Sprintf("%.4f", z))

			// tfact
			neg, i = utils.DecodeDecimal(tfact, &bi)
			z = utils.BigIntToFloat(neg, &bi, i)
			record = append(record, fmt.Sprintf("%.4f", z))

			//record = append(record, string(ffact2))
			/*
				record = append(record, string(ukurs[:bytes.IndexByte(ukurs, 0)]))
				record = append(record, string(ffact[:bytes.IndexByte(ffact, 0)]))
				record = append(record, string(tfact[:bytes.IndexByte(tfact, 0)]))
			*/
			/*
				record = append(record, strconv.FormatFloat(ukurs, 'f', 4, 64))
				record = append(record, strconv.FormatFloat(ffact, 'f', 4, 64))
				record = append(record, strconv.FormatFloat(tfact, 'f', 4, 64))
			*/
			record = append(record, datum)
			w.Write(record)
			//fmt.Println(record)
			//fmt.Printf("%s;%s;%s;%s;%s;%v;%v;%v;%s\n", mandt, kurst, fcurr, tcurr, gdatu, ukurs, ffact, tfact, datum)
		}
		w.Flush()

		if err := rows.Err(); err != nil {
			//fmt.Print("ROWS")
			utils.WriteMsg("ROWS")
			log.Fatal(err)
		}
		return nil
	})
}
<|endoftext|>


<commit_before>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package sh

import "fmt"

// Token is the set of lexical tokens.
type Token int

// The list of all possible tokens.
const (
	ILLEGAL Token = -iota
	STOPPED
	EOF
	LIT
	COMMENT

	// POSIX Shell reserved words
	IF
	THEN
	ELIF
	ELSE
	FI
	WHILE
	FOR
	IN
	UNTIL
	DO
	DONE
	CASE
	ESAC

	SQUOTE // '
	DQUOTE // "
	BQUOTE // `

	NOT    // !
	LBRACE // {
	RBRACE // }

	// Bash reserved words
	FUNCTION
	DECLARE

	// Rest of tokens
	AND  // &
	LAND // &&
	OR   // |
	LOR  // ||

	ASSIGN // =
	DOLLAR // $
	LPAREN // (

	RPAREN     // )
	SEMICOLON  // ;
	DSEMICOLON // ;;
	COLON      // :

	LSS // <
	GTR // >
	SHL // <<
	SHR // >>

	ADD   // +
	SUB   // -
	REM   // %
	MUL   // *
	QUO   // /
	XOR   // ^
	INC   // ++
	DEC   // --
	POW   // **
	COMMA // ,
	NEQ   // !=
	LEQ   // <=
	GEQ   // >=

	PIPEALL  // |&
	RDRINOUT // <>
	DPLIN    // <&
	DPLOUT   // >&
	DHEREDOC // <<-
	WHEREDOC // <<<
	CMDIN    // <(

	CADD    // :+
	CSUB    // :-
	QUEST   // ?
	CQUEST  // :?
	CASSIGN // :=
	DREM    // %%
	HASH    // #
	DHASH   // ##

	DLPAREN // ((
	DRPAREN // ))
)

// Pos is the internal representation of a position within a source
// file.
type Pos struct {
	Line   int // line number, starting at 1
	Column int // column number, starting at 1
}

// Position describes an arbitrary position in a source file. Offsets,
// including column numbers, are in bytes.
type Position struct {
	Filename string
	Line     int // line number, starting at 1
	Column   int // column number, starting at 1
}

func (p Position) String() string {
	if p.Filename == "" {
		return fmt.Sprintf("%d:%d", p.Line, p.Column)
	}
	return fmt.Sprintf("%s:%d:%d", p.Filename, p.Line, p.Column)
}

func init() {
	for _, t := range regList {
		tokNames[t.tok] = t.str
	}
	for _, t := range arithmList {
		tokNames[t.tok] = t.str
	}
	for _, t := range paramList {
		tokNames[t.tok] = t.str
	}
}

type tokEntry struct {
	str string
	tok Token
}

var (
	tokNames = map[Token]string{
		ILLEGAL: `ILLEGAL`,
		STOPPED: `STOPPED`,
		EOF:     `EOF`,
		LIT:     `literal`,
		COMMENT: `comment`,

		IF:    "if",
		THEN:  "then",
		ELIF:  "elif",
		ELSE:  "else",
		FI:    "fi",
		WHILE: "while",
		FOR:   "for",
		IN:    "in",
		UNTIL: "until",
		DO:    "do",
		DONE:  "done",
		CASE:  "case",
		ESAC:  "esac",

		NOT:    "!",
		LBRACE: "{",
		RBRACE: "}",

		FUNCTION: "function",
		DECLARE:  "declare",

		DLPAREN: "((",
		DRPAREN: "))",
	}

	regList = []tokEntry{
		{"'", SQUOTE},
		{"\"", DQUOTE},
		{"`", BQUOTE},

		{"&", AND},
		{"&&", LAND},
		{"|", OR},
		{"||", LOR},

		{"$", DOLLAR},
		{"(", LPAREN},
		{")", RPAREN},
		{";", SEMICOLON},
		{";;", DSEMICOLON},

		{"<", LSS},
		{">", GTR},
		{"<<", SHL},
		{">>", SHR},
		{"|&", PIPEALL},
		{"<>", RDRINOUT},
		{"<&", DPLIN},
		{">&", DPLOUT},
		{"<<-", DHEREDOC},
		{"<<<", WHEREDOC},
		{"<(", CMDIN},
	}
	paramList = []tokEntry{
		{"+", ADD},
		{":+", CADD},
		{"-", SUB},
		{":-", CSUB},
		{"?", QUEST},
		{":?", CQUEST},
		{"=", ASSIGN},
		{":=", CASSIGN},
		{"%", REM},
		{"%%", DREM},
		{"#", HASH},
		{"##", DHASH},
	}
	arithmList = []tokEntry{
		{"!", NOT},
		{"=", ASSIGN},
		{"(", LPAREN},
		{")", RPAREN},
		{"&", AND},
		{"&&", LAND},
		{"|", OR},
		{"||", LOR},
		{"<", LSS},
		{">", GTR},
		{"<<", SHL},
		{">>", SHR},

		{"+", ADD},
		{"-", SUB},
		{"%", REM},
		{"*", MUL},
		{"/", QUO},
		{"^", XOR},
		{"++", INC},
		{"--", DEC},
		{"**", POW},
		{",", COMMA},
		{"!=", NEQ},
		{"<=", LEQ},
		{">=", GEQ},
		{"?", QUEST},
		{":", COLON},
	}
)

func (t Token) String() string {
	if s, e := tokNames[t]; e {
		return s
	}
	return string(t)
}

func (p *parser) doToken(tokList []tokEntry) Token {
	// In reverse, to not treat e.g. && as & two times
	for i := len(tokList) - 1; i >= 0; i-- {
		t := tokList[i]
		if p.readOnly(t.str) {
			return t.tok
		}
	}
	return ILLEGAL
}

func (p *parser) doRegToken() Token    { return p.doToken(regList) }
func (p *parser) doParamToken() Token  { return p.doToken(paramList) }
func (p *parser) doArithmToken() Token { return p.doToken(arithmList) }
<commit_msg>Simplify Token.String()<commit_after>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package sh

import "fmt"

// Token is the set of lexical tokens.
type Token int

// The list of all possible tokens.
const (
	ILLEGAL Token = -iota
	STOPPED
	EOF
	LIT
	COMMENT

	// POSIX Shell reserved words
	IF
	THEN
	ELIF
	ELSE
	FI
	WHILE
	FOR
	IN
	UNTIL
	DO
	DONE
	CASE
	ESAC

	SQUOTE // '
	DQUOTE // "
	BQUOTE // `

	NOT    // !
	LBRACE // {
	RBRACE // }

	// Bash reserved words
	FUNCTION
	DECLARE

	// Rest of tokens
	AND  // &
	LAND // &&
	OR   // |
	LOR  // ||

	ASSIGN // =
	DOLLAR // $
	LPAREN // (

	RPAREN     // )
	SEMICOLON  // ;
	DSEMICOLON // ;;
	COLON      // :

	LSS // <
	GTR // >
	SHL // <<
	SHR // >>

	ADD   // +
	SUB   // -
	REM   // %
	MUL   // *
	QUO   // /
	XOR   // ^
	INC   // ++
	DEC   // --
	POW   // **
	COMMA // ,
	NEQ   // !=
	LEQ   // <=
	GEQ   // >=

	PIPEALL  // |&
	RDRINOUT // <>
	DPLIN    // <&
	DPLOUT   // >&
	DHEREDOC // <<-
	WHEREDOC // <<<
	CMDIN    // <(

	CADD    // :+
	CSUB    // :-
	QUEST   // ?
	CQUEST  // :?
	CASSIGN // :=
	DREM    // %%
	HASH    // #
	DHASH   // ##

	DLPAREN // ((
	DRPAREN // ))
)

// Pos is the internal representation of a position within a source
// file.
type Pos struct {
	Line   int // line number, starting at 1
	Column int // column number, starting at 1
}

// Position describes an arbitrary position in a source file. Offsets,
// including column numbers, are in bytes.
type Position struct {
	Filename string
	Line     int // line number, starting at 1
	Column   int // column number, starting at 1
}

func (p Position) String() string {
	if p.Filename == "" {
		return fmt.Sprintf("%d:%d", p.Line, p.Column)
	}
	return fmt.Sprintf("%s:%d:%d", p.Filename, p.Line, p.Column)
}

func init() {
	for _, t := range regList {
		tokNames[t.tok] = t.str
	}
	for _, t := range arithmList {
		tokNames[t.tok] = t.str
	}
	for _, t := range paramList {
		tokNames[t.tok] = t.str
	}
}

type tokEntry struct {
	str string
	tok Token
}

var (
	tokNames = map[Token]string{
		ILLEGAL: `ILLEGAL`,
		STOPPED: `STOPPED`,
		EOF:     `EOF`,
		LIT:     `literal`,
		COMMENT: `comment`,

		IF:    "if",
		THEN:  "then",
		ELIF:  "elif",
		ELSE:  "else",
		FI:    "fi",
		WHILE: "while",
		FOR:   "for",
		IN:    "in",
		UNTIL: "until",
		DO:    "do",
		DONE:  "done",
		CASE:  "case",
		ESAC:  "esac",

		NOT:    "!",
		LBRACE: "{",
		RBRACE: "}",

		FUNCTION: "function",
		DECLARE:  "declare",

		DLPAREN: "((",
		DRPAREN: "))",
	}

	regList = []tokEntry{
		{"'", SQUOTE},
		{"\"", DQUOTE},
		{"`", BQUOTE},

		{"&", AND},
		{"&&", LAND},
		{"|", OR},
		{"||", LOR},

		{"$", DOLLAR},
		{"(", LPAREN},
		{")", RPAREN},
		{";", SEMICOLON},
		{";;", DSEMICOLON},

		{"<", LSS},
		{">", GTR},
		{"<<", SHL},
		{">>", SHR},
		{"|&", PIPEALL},
		{"<>", RDRINOUT},
		{"<&", DPLIN},
		{">&", DPLOUT},
		{"<<-", DHEREDOC},
		{"<<<", WHEREDOC},
		{"<(", CMDIN},
	}
	paramList = []tokEntry{
		{"+", ADD},
		{":+", CADD},
		{"-", SUB},
		{":-", CSUB},
		{"?", QUEST},
		{":?", CQUEST},
		{"=", ASSIGN},
		{":=", CASSIGN},
		{"%", REM},
		{"%%", DREM},
		{"#", HASH},
		{"##", DHASH},
	}
	arithmList = []tokEntry{
		{"!", NOT},
		{"=", ASSIGN},
		{"(", LPAREN},
		{")", RPAREN},
		{"&", AND},
		{"&&", LAND},
		{"|", OR},
		{"||", LOR},
		{"<", LSS},
		{">", GTR},
		{"<<", SHL},
		{">>", SHR},

		{"+", ADD},
		{"-", SUB},
		{"%", REM},
		{"*", MUL},
		{"/", QUO},
		{"^", XOR},
		{"++", INC},
		{"--", DEC},
		{"**", POW},
		{",", COMMA},
		{"!=", NEQ},
		{"<=", LEQ},
		{">=", GEQ},
		{"?", QUEST},
		{":", COLON},
	}
)

func (t Token) String() string { return tokNames[t] }

func (p *parser) doToken(tokList []tokEntry) Token {
	// In reverse, to not treat e.g. && as & two times
	for i := len(tokList) - 1; i >= 0; i-- {
		t := tokList[i]
		if p.readOnly(t.str) {
			return t.tok
		}
	}
	return ILLEGAL
}

func (p *parser) doRegToken() Token    { return p.doToken(regList) }
func (p *parser) doParamToken() Token  { return p.doToken(paramList) }
func (p *parser) doArithmToken() Token { return p.doToken(arithmList) }
<|endoftext|>


<commit_before>//
// token.go
//
// Created by Frederic DELBOS - fred@hyperboloide.com on Nov 10 2014.
// This file is subject to the terms and conditions defined in
// file 'LICENSE.txt', which is part of this source code package.
//

package main

import (
	"github.com/dchest/uniuri"
)

type Token struct {
	Key        string `json:"key"`
	Service    string `json:"service"`
	Identifier string `json:"identifier,omitempty"`
}

type TokenService struct {
	Service string
	cache   Cache
}

func (ts *TokenService) NewToken(service string) (*Token, error) {
	token := &Token{
		Key:        uniuri.New(),
		Service:    service,
		Identifier: "",
	}
	if err := ts.Set(token); err != nil {
		return nil, err
	}
	return token, nil
}

func (ts *TokenService) Get(key string) (*Token, error) {
	token := new(Token)
	if err := ts.cache.Get(key, token); err != nil {
		return nil, err
	}
	return token, nil
}

func (ts *TokenService) Del(key string) error {
	return ts.cache.Del(key)
}

func (ts *TokenService) Set(token *Token) error {
	return ts.cache.Set(token.Key, token)
}
<commit_msg>cleanning<commit_after>//
// token.go
//
// Created by Frederic DELBOS - fred@hyperboloide.com on Nov 10 2014.
//

package main

import (
	"github.com/dchest/uniuri"
)

type Token struct {
	Key        string `json:"key"`
	Service    string `json:"service"`
	Identifier string `json:"identifier,omitempty"`
}

type TokenService struct {
	Service string
	cache   Cache
}

func (ts *TokenService) NewToken(service string) (*Token, error) {
	token := &Token{
		Key:        uniuri.New(),
		Service:    service,
		Identifier: "",
	}
	if err := ts.Set(token); err != nil {
		return nil, err
	}
	return token, nil
}

func (ts *TokenService) Get(key string) (*Token, error) {
	token := new(Token)
	if err := ts.cache.Get(key, token); err != nil {
		return nil, err
	}
	return token, nil
}

func (ts *TokenService) Del(key string) error {
	return ts.cache.Del(key)
}

func (ts *TokenService) Set(token *Token) error {
	return ts.cache.Set(token.Key, token)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

func goimports(paths []string, verbose bool) error {
	args := []string{"-l", "-w"}
	for _, path := range paths {
		args = append(args, path)
	}
	if verbose {
		fmt.Println("goimports", strings.Join(args, " "))
	}
	cmd := exec.Command("goimports", args...)
	var outbuf bytes.Buffer
	cmd.Stdout = &outbuf
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}
	if outbuf.String() != "" {
		return fmt.Errorf("goimports -l -w:\n%s", strings.TrimSpace(outbuf.String()))
	}
	return nil
}

func gofmt(paths []string, verbose bool) error {
	args := []string{"-l", "-w", "-s"}
	for _, path := range paths {
		args = append(args, path)
	}
	if verbose {
		fmt.Println("gofmt", strings.Join(args, " "))
	}
	cmd := exec.Command("gofmt", args...)
	var outbuf bytes.Buffer
	cmd.Stdout = &outbuf
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}
	if outbuf.String() != "" {
		return fmt.Errorf("gofmt -l -w -s:\n%s", strings.TrimSpace(outbuf.String()))
	}
	return nil
}

func golint(paths []string, verbose bool) error {
	var pathsWithSubDirs []string
	for _, path := range paths {
		fi, err := os.Stat(path)
		if err != nil {
			return err
		}
		if fi.IsDir() {
			pathsWithSubDirs = append(pathsWithSubDirs, filepath.Join(path, "..."))
		} else {
			pathsWithSubDirs = append(pathsWithSubDirs, path)
		}

	}
	if verbose {
		fmt.Println("golint", strings.Join(pathsWithSubDirs, " "))
	}
	cmd := exec.Command("golint", pathsWithSubDirs...)
	var outbuf bytes.Buffer
	cmd.Stdout = &outbuf
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}
	if outbuf.String() != "" {
		return fmt.Errorf("golint:\n%s", strings.TrimSpace(outbuf.String()))
	}
	return nil
}

func govet(paths []string, verbose bool) error {
	if verbose {
		fmt.Println("go vet", strings.Join(paths, " "))
	}
	args := []string{"vet"}
	args = append(args, paths...)
	cmd := exec.Command("go", args...)
	var errbuf bytes.Buffer
	cmd.Stdout = os.Stdout
	cmd.Stderr = &errbuf
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("%s: go vet:\n%s", err,
			strings.TrimSpace(errbuf.String()))
	}
	return nil
}

func gotest(path string, coverage, verbose bool) error {
	args := []string{"test"}
	if coverage {
		args = append(args, "-cover")
	}
	args = append(args, path)
	if verbose {
		fmt.Println("go", strings.Join(args, " "))
	}
	cmd := exec.Command("go", args...)
	var errbuf bytes.Buffer
	cmd.Stdout = os.Stdout
	cmd.Stderr = &errbuf
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("%s: go test:\n%s", err,
			strings.TrimSpace(errbuf.String()))
	}
	return nil
}
<commit_msg>propagate -v flag to `go test`<commit_after>package main

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

func goimports(paths []string, verbose bool) error {
	args := []string{"-l", "-w"}
	for _, path := range paths {
		args = append(args, path)
	}
	if verbose {
		fmt.Println("goimports", strings.Join(args, " "))
	}
	cmd := exec.Command("goimports", args...)
	var outbuf bytes.Buffer
	cmd.Stdout = &outbuf
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}
	if outbuf.String() != "" {
		return fmt.Errorf("goimports -l -w:\n%s", strings.TrimSpace(outbuf.String()))
	}
	return nil
}

func gofmt(paths []string, verbose bool) error {
	args := []string{"-l", "-w", "-s"}
	for _, path := range paths {
		args = append(args, path)
	}
	if verbose {
		fmt.Println("gofmt", strings.Join(args, " "))
	}
	cmd := exec.Command("gofmt", args...)
	var outbuf bytes.Buffer
	cmd.Stdout = &outbuf
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}
	if outbuf.String() != "" {
		return fmt.Errorf("gofmt -l -w -s:\n%s", strings.TrimSpace(outbuf.String()))
	}
	return nil
}

func golint(paths []string, verbose bool) error {
	var pathsWithSubDirs []string
	for _, path := range paths {
		fi, err := os.Stat(path)
		if err != nil {
			return err
		}
		if fi.IsDir() {
			pathsWithSubDirs = append(pathsWithSubDirs, filepath.Join(path, "..."))
		} else {
			pathsWithSubDirs = append(pathsWithSubDirs, path)
		}

	}
	if verbose {
		fmt.Println("golint", strings.Join(pathsWithSubDirs, " "))
	}
	cmd := exec.Command("golint", pathsWithSubDirs...)
	var outbuf bytes.Buffer
	cmd.Stdout = &outbuf
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}
	if outbuf.String() != "" {
		return fmt.Errorf("golint:\n%s", strings.TrimSpace(outbuf.String()))
	}
	return nil
}

func govet(paths []string, verbose bool) error {
	if verbose {
		fmt.Println("go vet", strings.Join(paths, " "))
	}
	args := []string{"vet"}
	args = append(args, paths...)
	cmd := exec.Command("go", args...)
	var errbuf bytes.Buffer
	cmd.Stdout = os.Stdout
	cmd.Stderr = &errbuf
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("%s: go vet:\n%s", err,
			strings.TrimSpace(errbuf.String()))
	}
	return nil
}

func gotest(path string, coverage, verbose bool) error {
	args := []string{"test"}
	if coverage {
		args = append(args, "-cover")
	}
	if verbose {
		args = append(args, "-v")
	}
	args = append(args, path)
	if verbose {
		fmt.Println("go", strings.Join(args, " "))
	}
	cmd := exec.Command("go", args...)
	var errbuf bytes.Buffer
	cmd.Stdout = os.Stdout
	cmd.Stderr = &errbuf
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("%s: go test:\n%s", err,
			strings.TrimSpace(errbuf.String()))
	}
	return nil
}
<|endoftext|>


<commit_before>package sq

import (
	"encoding/binary"
	"fmt"
	"os"
	"path"
	"sync"
	"syscall"
	"unsafe"
)

var (
	encoder           = binary.LittleEndian
	blank             = struct{}{}
	ChannelNameLenErr = fmt.Errorf("channel name cannot exceed %d characters", MAX_CHANNEL_NAME_SIZE)
	ChannelExistsErr  = fmt.Errorf("channel already exists")
	ChannelCreateErr  = fmt.Errorf("channel count not be created")
	pageSize          = os.Getpagesize()
)

type addChannelWork struct {
	err     error
	channel *Channel
	config  *ChannelConfiguration
	c       chan *addChannelWork
}

type Topic struct {
	path         string
	state        *State
	states       *States
	segmentSize  int
	channelsLock sync.RWMutex
	channels     map[string]*Channel
	segment      *Segment
	segmentsLock sync.RWMutex
	segments     map[uint64]*Segment
	addChannel   chan *addChannelWork
	messageAdded chan struct{}
	dataLock     sync.RWMutex
}

func OpenTopic(name string, config *TopicConfiguration) (*Topic, error) {
	t := &Topic{
		path:         path.Join(config.path, name),
		channels:     make(map[string]*Channel),
		segments:     make(map[uint64]*Segment),
		addChannel:   make(chan *addChannelWork),
		segmentSize:  config.segmentSize,
		messageAdded: make(chan struct{}, 64),
	}
	err := loadStates(t)
	if err != nil {
		return nil, err
	}

	if id := t.state.segmentId; id == 0 {
		if err := t.expand(); err != nil {
			return nil, err
		}
		if err := t.states.syncTopic(); err != nil {
			return nil, err
		}
	} else {
		t.findWritePosition(openSegment(t, id, false))
		t.segments[id] = t.segment
	}
	go t.worker()
	return t, nil
}

func (t *Topic) Write(data []byte) error {
	length := len(data)

	t.dataLock.Lock()
	start := int(t.state.offset)
	dataStart := start + 4
	dataEnd := dataStart + length

	// do we have enough space in the current segment?
	if dataEnd > t.segmentSize {
		if err := t.expand(); err != nil {
			return err
		}
		start = int(SEGMENT_HEADER_SIZE)
		dataStart = start + 4
		dataEnd = dataStart + length
	}

	// write the message length
	encoder.PutUint32(t.segment.data[start:], uint32(length))
	// write the message
	copy(t.segment.data[dataStart:], data)
	// location of the next write
	t.state.offset = uint32(dataEnd)
	t.dataLock.Unlock()

	// sync the part of the data file we just wrote
	from := start / pageSize * pageSize
	to := dataStart + length - from
	_, _, errno := syscall.Syscall(syscall.SYS_MSYNC, uintptr(unsafe.Pointer(&t.segment.data[from])), uintptr(to), syscall.MS_SYNC)
	if errno != 0 {
		return syscall.Errno(errno)
	}

	// notify channels that a new message is waiting
	t.messageAdded <- blank

	return nil
}

func (t *Topic) Channel(name string, config *ChannelConfiguration) (*Channel, error) {
	if len(name) > MAX_CHANNEL_NAME_SIZE {
		return nil, ChannelNameLenErr
	}
	if config == nil {
		config = ConfigureChannel()
	}
	config.name = name
	res := &addChannelWork{
		c:      make(chan *addChannelWork),
		config: config,
	}
	t.addChannel <- res
	<-res.c
	return res.channel, res.err
}

func (t *Topic) expand() error {
	segment := newSegment(t)
	t.segmentsLock.Lock()
	t.segments[segment.id] = segment
	t.segmentsLock.Unlock()
	if t.segment != nil {
		// create a pointer to the next segment id
		t.segment.nextId = segment.id
		t.segment.size = t.state.offset
		if err := t.segment.syncHeader(); err != nil {
			return err
		}
		t.channelsLock.RLock()
		noChannels := len(t.states.channels) == 0
		t.channelsLock.RUnlock()
		if noChannels {
			t.deleteSegment(t.segment.id)
		}
	}
	t.segment = segment
	t.state.offset = SEGMENT_HEADER_SIZE
	t.state.segmentId = segment.id
	return nil
}

func (t *Topic) worker() {
	for {
		select {
		case work := <-t.addChannel:
			work.channel, work.err = t.createChannel(work.config)
			work.c <- work
		case <-t.messageAdded:
			t.channelsLock.RLock()
			for _, c := range t.channels {
				c.notify()
			}
			t.channelsLock.RUnlock()
		}
	}
}

func (t *Topic) createChannel(config *ChannelConfiguration) (*Channel, error) {
	t.channelsLock.Lock()
	_, exists := t.channels[config.name]
	if exists {
		t.channelsLock.Unlock()
		if config.temp {
			// try again until we've created one
			return t.createChannel(config)
		}
		return nil, ChannelExistsErr
	}
	defer t.channelsLock.Unlock()

	c := newChannel(t, config)
	if config.temp {
		c.state = new(State)
	} else {
		c.state = t.states.getOrCreate(c.name)
		if c.state == nil {
			return nil, ChannelCreateErr
		}
	}

	// if we have a temp channel, or a new channel, set the position to the
	// writer's current position
	if config.temp || c.state.segmentId == 0 {
		t.dataLock.RLock()
		c.state.offset = t.state.offset
		c.state.segmentId = t.state.segmentId
		t.dataLock.RUnlock()
	}
	t.channels[c.name] = c
	return c, nil
}

func (t *Topic) read(channel *Channel) []byte {
	state := channel.state

	// no need to lock the channel here since the state of the channel is only
	// ever changed after this point (either to move to the next segment or message)
	t.dataLock.RLock()
	isCurrentSegment := state.segmentId == t.state.segmentId
	isCurrentOffset := state.offset == t.state.offset
	t.dataLock.RUnlock()

	// are we fully caught up with the topic?
	if isCurrentSegment && isCurrentOffset {
		return nil
	}

	segment := t.loadSegment(state.segmentId)
	// If we aren't on the current segment, we might be at the end of an old
	// segment. We can safely read old segments without locking since the only
	// writer (the topic), is done with it
	if !isCurrentSegment {
		if state.offset >= segment.size {
			previousId := state.segmentId
			segment = t.loadSegment(segment.nextId)
			channel.changeSegment(segment)
			if t.isSegmentUsable(previousId) == false {
				t.deleteSegment(previousId)
			}
		}
	}

	l := encoder.Uint32(segment.data[state.offset:])
	start := state.offset + 4
	end := start + uint32(l)
	return segment.data[start:end]
}

func (t *Topic) loadSegment(id uint64) *Segment {
	t.segmentsLock.RLock()
	segment := t.segments[id]
	t.segmentsLock.RUnlock()
	if segment != nil {
		return segment
	}

	t.segmentsLock.Lock()
	defer t.segmentsLock.Unlock()
	if segment := t.segments[id]; segment != nil {
		return segment
	}
	segment = openSegment(t, id, false)
	t.segments[id] = segment
	return segment
}

// A segment is usable if any channel references it or an earlier one
// How we check this depends on whether or not the channel is active
// (currently consuming message). For an active channel, we ask the channel,
// for an inactive channel, we can query the state direclty knowing that
// no one can be writing to it at the same time
func (t *Topic) isSegmentUsable(id uint64) bool {
	t.channelsLock.RLock()
	defer t.channelsLock.RUnlock()
	for name, state := range t.states.channels {
		if channel, active := t.channels[name]; active {
			if channel.isSegmentUsable(id) == true {
				return true
			}
		} else if state.isSegmentUsable(id) == true {
			return true
		}
	}
	return false
}

// The topic's persisted state could be behind from the actual state. However,
// we can determine the actual state based on the data which is guaranteed to be
// up to date (the next reference id and the data itself)
func (t *Topic) findWritePosition(segment *Segment) {
	for segment.nextId != 0 {
		segment = openSegment(t, segment.nextId, false)
	}
	t.segment = segment
	t.state.segmentId = segment.id

	offset := SEGMENT_HEADER_SIZE
	for offset < uint32(t.segmentSize) {
		l := encoder.Uint32(segment.data[offset:])
		if l == 0 {
			break
		}
		offset += 4 + l
	}
	t.state.offset = offset
}

func (t *Topic) deleteSegment(id uint64) {
	t.segmentsLock.Lock()
	segment := t.segments[id]
	delete(t.segments, id)
	t.segmentsLock.Unlock()
	if segment != nil {
		segment.delete()
	}
}
<commit_msg>don't break when message is larger than maximum message size (return an error)<commit_after>package sq

import (
	"encoding/binary"
	"fmt"
	"os"
	"path"
	"sync"
	"syscall"
	"unsafe"
)

var (
	encoder            = binary.LittleEndian
	blank              = struct{}{}
	ChannelNameLenErr  = fmt.Errorf("channel name cannot exceed %d characters", MAX_CHANNEL_NAME_SIZE)
	ChannelExistsErr   = fmt.Errorf("channel already exists")
	ChannelCreateErr   = fmt.Errorf("channel count not be created")
	MessageTooLargeErr = fmt.Errorf("message too large")
	pageSize           = os.Getpagesize()
)

type addChannelWork struct {
	err     error
	channel *Channel
	config  *ChannelConfiguration
	c       chan *addChannelWork
}

type Topic struct {
	path         string
	state        *State
	states       *States
	segmentSize  int
	channelsLock sync.RWMutex
	channels     map[string]*Channel
	segment      *Segment
	segmentsLock sync.RWMutex
	segments     map[uint64]*Segment
	addChannel   chan *addChannelWork
	messageAdded chan struct{}
	dataLock     sync.RWMutex
}

func OpenTopic(name string, config *TopicConfiguration) (*Topic, error) {
	t := &Topic{
		path:         path.Join(config.path, name),
		channels:     make(map[string]*Channel),
		segments:     make(map[uint64]*Segment),
		addChannel:   make(chan *addChannelWork),
		segmentSize:  config.segmentSize,
		messageAdded: make(chan struct{}, 64),
	}
	err := loadStates(t)
	if err != nil {
		return nil, err
	}

	if id := t.state.segmentId; id == 0 {
		if err := t.expand(); err != nil {
			return nil, err
		}
		if err := t.states.syncTopic(); err != nil {
			return nil, err
		}
	} else {
		t.findWritePosition(openSegment(t, id, false))
		t.segments[id] = t.segment
	}
	go t.worker()
	return t, nil
}

func (t *Topic) Write(data []byte) error {
	length := len(data)

	t.dataLock.Lock()
	start := int(t.state.offset)
	dataStart := start + 4
	dataEnd := dataStart + length

	// do we have enough space in the current segment?
	if dataEnd > t.segmentSize {
		if length+int(SEGMENT_HEADER_SIZE) > t.segmentSize {
			return MessageTooLargeErr
		}
		if err := t.expand(); err != nil {
			return err
		}
		start = int(SEGMENT_HEADER_SIZE)
		dataStart = start + 4
		dataEnd = dataStart + length
	}

	// write the message length
	encoder.PutUint32(t.segment.data[start:], uint32(length))
	// write the message
	copy(t.segment.data[dataStart:], data)
	// location of the next write
	t.state.offset = uint32(dataEnd)
	t.dataLock.Unlock()

	// sync the part of the data file we just wrote
	from := start / pageSize * pageSize
	to := dataStart + length - from
	_, _, errno := syscall.Syscall(syscall.SYS_MSYNC, uintptr(unsafe.Pointer(&t.segment.data[from])), uintptr(to), syscall.MS_SYNC)
	if errno != 0 {
		return syscall.Errno(errno)
	}

	// notify channels that a new message is waiting
	t.messageAdded <- blank

	return nil
}

func (t *Topic) Channel(name string, config *ChannelConfiguration) (*Channel, error) {
	if len(name) > MAX_CHANNEL_NAME_SIZE {
		return nil, ChannelNameLenErr
	}
	if config == nil {
		config = ConfigureChannel()
	}
	config.name = name
	res := &addChannelWork{
		c:      make(chan *addChannelWork),
		config: config,
	}
	t.addChannel <- res
	<-res.c
	return res.channel, res.err
}

func (t *Topic) expand() error {
	segment := newSegment(t)
	t.segmentsLock.Lock()
	t.segments[segment.id] = segment
	t.segmentsLock.Unlock()
	if t.segment != nil {
		// create a pointer to the next segment id
		t.segment.nextId = segment.id
		t.segment.size = t.state.offset
		if err := t.segment.syncHeader(); err != nil {
			return err
		}
		t.channelsLock.RLock()
		noChannels := len(t.states.channels) == 0
		t.channelsLock.RUnlock()
		if noChannels {
			t.deleteSegment(t.segment.id)
		}
	}
	t.segment = segment
	t.state.offset = SEGMENT_HEADER_SIZE
	t.state.segmentId = segment.id
	return nil
}

func (t *Topic) worker() {
	for {
		select {
		case work := <-t.addChannel:
			work.channel, work.err = t.createChannel(work.config)
			work.c <- work
		case <-t.messageAdded:
			t.channelsLock.RLock()
			for _, c := range t.channels {
				c.notify()
			}
			t.channelsLock.RUnlock()
		}
	}
}

func (t *Topic) createChannel(config *ChannelConfiguration) (*Channel, error) {
	t.channelsLock.Lock()
	_, exists := t.channels[config.name]
	if exists {
		t.channelsLock.Unlock()
		if config.temp {
			// try again until we've created one
			return t.createChannel(config)
		}
		return nil, ChannelExistsErr
	}
	defer t.channelsLock.Unlock()

	c := newChannel(t, config)
	if config.temp {
		c.state = new(State)
	} else {
		c.state = t.states.getOrCreate(c.name)
		if c.state == nil {
			return nil, ChannelCreateErr
		}
	}

	// if we have a temp channel, or a new channel, set the position to the
	// writer's current position
	if config.temp || c.state.segmentId == 0 {
		t.dataLock.RLock()
		c.state.offset = t.state.offset
		c.state.segmentId = t.state.segmentId
		t.dataLock.RUnlock()
	}
	t.channels[c.name] = c
	return c, nil
}

func (t *Topic) read(channel *Channel) []byte {
	state := channel.state

	// no need to lock the channel here since the state of the channel is only
	// ever changed after this point (either to move to the next segment or message)
	t.dataLock.RLock()
	isCurrentSegment := state.segmentId == t.state.segmentId
	isCurrentOffset := state.offset == t.state.offset
	t.dataLock.RUnlock()

	// are we fully caught up with the topic?
	if isCurrentSegment && isCurrentOffset {
		return nil
	}

	segment := t.loadSegment(state.segmentId)
	// If we aren't on the current segment, we might be at the end of an old
	// segment. We can safely read old segments without locking since the only
	// writer (the topic), is done with it
	if !isCurrentSegment {
		if state.offset >= segment.size {
			previousId := state.segmentId
			segment = t.loadSegment(segment.nextId)
			channel.changeSegment(segment)
			if t.isSegmentUsable(previousId) == false {
				t.deleteSegment(previousId)
			}
		}
	}

	l := encoder.Uint32(segment.data[state.offset:])
	start := state.offset + 4
	end := start + uint32(l)
	return segment.data[start:end]
}

func (t *Topic) loadSegment(id uint64) *Segment {
	t.segmentsLock.RLock()
	segment := t.segments[id]
	t.segmentsLock.RUnlock()
	if segment != nil {
		return segment
	}

	t.segmentsLock.Lock()
	defer t.segmentsLock.Unlock()
	if segment := t.segments[id]; segment != nil {
		return segment
	}
	segment = openSegment(t, id, false)
	t.segments[id] = segment
	return segment
}

// A segment is usable if any channel references it or an earlier one
// How we check this depends on whether or not the channel is active
// (currently consuming message). For an active channel, we ask the channel,
// for an inactive channel, we can query the state direclty knowing that
// no one can be writing to it at the same time
func (t *Topic) isSegmentUsable(id uint64) bool {
	t.channelsLock.RLock()
	defer t.channelsLock.RUnlock()
	for name, state := range t.states.channels {
		if channel, active := t.channels[name]; active {
			if channel.isSegmentUsable(id) == true {
				return true
			}
		} else if state.isSegmentUsable(id) == true {
			return true
		}
	}
	return false
}

// The topic's persisted state could be behind from the actual state. However,
// we can determine the actual state based on the data which is guaranteed to be
// up to date (the next reference id and the data itself)
func (t *Topic) findWritePosition(segment *Segment) {
	for segment.nextId != 0 {
		segment = openSegment(t, segment.nextId, false)
	}
	t.segment = segment
	t.state.segmentId = segment.id

	offset := SEGMENT_HEADER_SIZE
	for offset < uint32(t.segmentSize) {
		l := encoder.Uint32(segment.data[offset:])
		if l == 0 {
			break
		}
		offset += 4 + l
	}
	t.state.offset = offset
}

func (t *Topic) deleteSegment(id uint64) {
	t.segmentsLock.Lock()
	segment := t.segments[id]
	delete(t.segments, id)
	t.segmentsLock.Unlock()
	if segment != nil {
		segment.delete()
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

// afpacket provides a simple example of using afpacket with zero-copy to read
// packet data.
package main

import (
	"flag"
	"log"
	"os"
	"runtime/pprof"

	"github.com/google/gopacket/afpacket"

	_ "github.com/google/gopacket/layers"
)

var (
	iface      = flag.String("i", "eth0", "Interface to read from")
	cpuprofile = flag.String("cpuprofile", "", "If non-empty, write CPU profile here")
	count      = flag.Int64("c", -1, "If >= 0, # of packets to capture before returning")
	verbose    = flag.Int64("log_every", 1, "Write a log every X packets")
	addVLAN    = flag.Bool("add_vlan", false, "If true, add VLAN header")
)

func main() {
	flag.Parse()
	if *cpuprofile != "" {
		log.Printf("Writing CPU profile to %q", *cpuprofile)
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal(err)
		}
		defer pprof.StopCPUProfile()
	}
	log.Printf("Starting on interface %q", *iface)
	source, err := afpacket.NewTPacket(
		afpacket.OptInterface(*iface),
		afpacket.OptBlockSize(1<<20 /*1MB*/),
		afpacket.OptAddVLANHeader(*addVLAN))
	if err != nil {
		log.Fatal(err)
	}
	defer source.Close()
	bytes := uint64(0)
	packets := uint64(0)
	for ; *count != 0; *count-- {
		data, _, err := source.ZeroCopyReadPacketData()
		if err != nil {
			log.Fatal(err)
		}
		bytes += uint64(len(data))
		packets++
		if *count%*verbose == 0 {
			log.Printf("Read in %d bytes in %d packets", bytes, packets)
		}
	}
}
<commit_msg>Extend afpacket example<commit_after>// Copyright 2018 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

// afpacket provides a simple example of using afpacket with zero-copy to read
// packet data.
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"runtime/pprof"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/afpacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"golang.org/x/net/bpf"

	_ "github.com/google/gopacket/layers"
)

var (
	iface      = flag.String("i", "any", "Interface to read from")
	cpuprofile = flag.String("cpuprofile", "", "If non-empty, write CPU profile here")
	snaplen    = flag.Int("s", 0, "Snaplen, if <= 0, use 65535")
	bufferSize = flag.Int("b", 8, "Interface buffersize (MB)")
	filter     = flag.String("f", "port not 22", "BPF filter")
	count      = flag.Int64("c", -1, "If >= 0, # of packets to capture before returning")
	verbose    = flag.Int64("log_every", 1, "Write a log every X packets")
	addVLAN    = flag.Bool("add_vlan", false, "If true, add VLAN header")
)

type afpacketHandle struct {
	TPacket *afpacket.TPacket
}

func newAfpacketHandle(device string, snaplen int, block_size int, num_blocks int,
	timeout time.Duration) (*afpacketHandle, error) {

	h := &afpacketHandle{}
	var err error

	if *iface == "any" {
		h.TPacket, err = afpacket.NewTPacket(
			afpacket.OptFrameSize(snaplen),
			afpacket.OptBlockSize(block_size),
			afpacket.OptNumBlocks(num_blocks),
			afpacket.OptAddVLANHeader(*addVLAN),
			afpacket.OptPollTimeout(timeout),
			afpacket.SocketRaw,
			afpacket.TPacketVersion3)
	} else {
		h.TPacket, err = afpacket.NewTPacket(
			afpacket.OptInterface(*iface),
			afpacket.OptFrameSize(snaplen),
			afpacket.OptBlockSize(block_size),
			afpacket.OptNumBlocks(num_blocks),
			afpacket.OptAddVLANHeader(*addVLAN),
			afpacket.OptPollTimeout(timeout),
			afpacket.SocketRaw,
			afpacket.TPacketVersion3)
	}
	return h, err
}

// ReadPacketData satisfies PacketDataSource interface. Here it will use
// zero copy mode.
func (h *afpacketHandle) ReadPacketData() (data []byte, ci gopacket.CaptureInfo, err error) {
	return h.TPacket.ZeroCopyReadPacketData()
}

// SetBPFFilter translates a BPF filter string into BPF RawInstruction and applies them.
func (h *afpacketHandle) SetBPFFilter(filter string, snaplen int) (err error) {
	pcapBPF, err := pcap.CompileBPFFilter(layers.LinkTypeEthernet, snaplen, filter)
	if err != nil {
		return err
	}
	bpfIns := []bpf.RawInstruction{}
	for _, ins := range pcapBPF {
		bpfIns2 := bpf.RawInstruction{
			Op: ins.Code,
			Jt: ins.Jt,
			Jf: ins.Jf,
			K:  ins.K,
		}
		bpfIns = append(bpfIns, bpfIns2)
	}
	if h.TPacket.SetBPF(bpfIns); err != nil {
		return err
	}
	return h.TPacket.SetBPF(bpfIns)
}

// LinkType returns ethernet link type.
func (h *afpacketHandle) LinkType() layers.LinkType {
	return layers.LinkTypeEthernet
}

// Close will close afpacket source.
func (h *afpacketHandle) Close() {
	h.TPacket.Close()
}

// SocketStats prints received, dropped, queue-freeze packet stats.
func (h *afpacketHandle) SocketStats() (as afpacket.SocketStats, asv afpacket.SocketStatsV3, err error) {
	return h.TPacket.SocketStats()
}

// afpacketComputeSize computes the block_size and the num_blocks in such a way that the
// allocated mmap buffer is close to but smaller than target_size_mb.
// The restriction is that the block_size must be divisible by both the
// frame size and page size.
func afpacketComputeSize(targetSizeMb int, snaplen int, pageSize int) (
	frameSize int, blockSize int, numBlocks int, err error) {

	if snaplen < pageSize {
		frameSize = pageSize / (pageSize / snaplen)
	} else {
		frameSize = (snaplen/pageSize + 1) * pageSize
	}

	// 128 is the default from the gopacket library so just use that
	blockSize = frameSize * 128
	numBlocks = (targetSizeMb * 1024 * 1024) / blockSize

	if numBlocks == 0 {
		return 0, 0, 0, fmt.Errorf("Interface buffersize is too small")
	}

	return frameSize, blockSize, numBlocks, nil
}

func main() {
	flag.Parse()
	if *cpuprofile != "" {
		log.Printf("Writing CPU profile to %q", *cpuprofile)
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal(err)
		}
		defer pprof.StopCPUProfile()
	}
	log.Printf("Starting on interface %q", *iface)
	if *snaplen <= 0 {
		*snaplen = 65535
	}
	szFrame, szBlock, numBlocks, err := afpacketComputeSize(*bufferSize, *snaplen, os.Getpagesize())
	if err != nil {
		log.Fatal(err)
	}
	afpacketHandle, err := newAfpacketHandle(*iface, szFrame, szBlock, numBlocks, pcap.BlockForever)
	if err != nil {
		log.Fatal(err)
	}
	err = afpacketHandle.SetBPFFilter(*filter, *snaplen)
	if err != nil {
		log.Fatal(err)
	}
	source := gopacket.PacketDataSource(afpacketHandle)
	defer afpacketHandle.Close()

	bytes := uint64(0)
	packets := uint64(0)
	for ; *count != 0; *count-- {
		data, _, err := source.ReadPacketData()
		if err != nil {
			log.Fatal(err)
		}
		bytes += uint64(len(data))
		packets++
		if *count%*verbose == 0 {
			_, afpacketStats, err := afpacketHandle.SocketStats()
			if err != nil {
				log.Println(err)
			}
			log.Printf("Read in %d bytes in %d packets", bytes, packets)
			log.Printf("Stats {received dropped queue-freeze}: %d", afpacketStats)
		}
	}
}
<|endoftext|>