<commit_before>package minion

import (
	"os"
	"os/exec"
	"os/signal"
	"log"
	"bytes"
	"time"
	"strings"
	"strconv"
	"path/filepath"
	"encoding/json"

	"code.google.com/p/go-uuid/uuid"
	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
	etcdclient "github.com/coreos/etcd/client"
)

// Minions keyspace in etcd
const EtcdMinionSpace = "/gru/minion"

// Etcd Minion
type etcdMinion struct {
	// Name of this minion
	name string

	// Minion root node in etcd
	rootDir string

	// Minion queue node in etcd
	queueDir string

	// Log directory to keep previously executed tasks
	logDir string

	// Root node for classifiers in etcd
	classifierDir string

	// Minion unique identifier
	id uuid.UUID

	// KeysAPI client to etcd
	kapi etcdclient.KeysAPI
}

// Creates a new etcd minion
func NewEtcdMinion(name string, cfg etcdclient.Config) Minion {
	c, err := etcdclient.New(cfg)
	if err != nil {
		log.Fatal(err)
	}

	kapi := etcdclient.NewKeysAPI(c)
	id := GenerateUUID(name)
	rootDir := filepath.Join(EtcdMinionSpace, id.String())
	queueDir := filepath.Join(rootDir, "queue")
	classifierDir := filepath.Join(rootDir, "classifier")
	logDir := filepath.Join(rootDir, "log")

	m := &etcdMinion{
		name: name,
		rootDir: rootDir,
		queueDir: queueDir,
		classifierDir: classifierDir,
		logDir: logDir,
		id: id,
		kapi: kapi,
	}

	return m
}

// Set the human-readable name of the minion in etcd
func (m *etcdMinion) setName() error {
	nameKey := filepath.Join(m.rootDir, "name")
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), nameKey, m.name, opts)

	return err
}

// Set the time the minion was last seen in seconds since the Epoch
func (m *etcdMinion) setLastseen(s int64) error {
	lastseenKey := filepath.Join(m.rootDir, "lastseen")
	lastseenValue := strconv.FormatInt(s, 10)
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), lastseenKey, lastseenValue, opts)

	return err
}

// Checks for any tasks pending tasks in queue
func (m *etcdMinion) checkQueue(c chan<- *MinionTask) error {
	opts := &etcdclient.GetOptions{
		Recursive: true,
		Sort: true,
	}

	// Get backlog tasks if any
	resp, err := m.kapi.Get(context.Background(), m.queueDir, opts)
	if err != nil {
		return nil
	}

	backlog := resp.Node.Nodes
	if len(backlog) == 0 {
		// No backlog tasks found
		return nil
	}

	log.Printf("Found %d tasks in backlog", len(backlog))
	for _, node := range backlog {
		task, err := EtcdUnmarshalTask(node)
		m.kapi.Delete(context.Background(), node.Key, nil)

		if err != nil {
			continue
		}

		c <- task
	}

	return nil
}

// Runs periodic jobs such as refreshing classifiers and updating lastseen
func (m *etcdMinion) periodicRunner(ticker *time.Ticker) error {
	for {
		// Update classifiers
		for _, classifier := range ClassifierRegistry {
			m.SetClassifier(classifier)
		}

		// Update lastseen time
		now := time.Now().Unix()
		err := m.setLastseen(now)
		if err != nil {
			log.Printf("Failed to update lastseen time: %s\n", err)
		}

		<- ticker.C
	}

	return nil
}

// Processes new tasks
func (m *etcdMinion) processTask(t *MinionTask) error {
	defer m.saveTask(t)

	var buf bytes.Buffer
	cmd := exec.Command(t.Command, t.Args...)
	cmd.Stdout = &buf
	cmd.Stderr = &buf

	log.Printf("Processing task %s\n", t.TaskID)

	cmdError := cmd.Run()
	t.TimeProcessed = time.Now().Unix()
	t.Result = buf.String()

	if cmdError != nil {
		log.Printf("Failed to process task %s\n", t.TaskID)
		t.Error = cmdError.Error()
	} else {
		log.Printf("Finished processing task %s\n", t.TaskID)
	}

	return cmdError
}

// Saves a task in the minion's log
func (m *etcdMinion) saveTask(t *MinionTask) error {
	// Task key in etcd
	taskKey := filepath.Join(m.logDir, t.TaskID.String())

	// Serialize task to JSON
	data, err := json.Marshal(t)
	if err != nil {
		log.Printf("Failed to serialize task %s: %s\n", t.TaskID, err)
		return err
	}

	// Save task result in the minion's space
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}
	_, err = m.kapi.Set(context.Background(), taskKey, string(data), opts)
	if err != nil {
		log.Printf("Failed to save task %s: %s\n", t.TaskID, err)
		return err
	}

	return err
}

// Unmarshals task from etcd
func EtcdUnmarshalTask(node *etcdclient.Node) (*MinionTask, error) {
	task := new(MinionTask)
	err := json.Unmarshal([]byte(node.Value), &task)

	if err != nil {
		log.Printf("Invalid task %s: %s\n", node.Key, err)
	}

	return task, err
}

// Returns the minion unique identifier
func (m *etcdMinion) ID() uuid.UUID {
	return m.id
}

// Returns the assigned name of the minion
func (m *etcdMinion) Name() string {
	return m.name
}

// Classify a minion with a given key and value
func (m *etcdMinion) SetClassifier(c MinionClassifier) error {
	// Classifiers in etcd expire after an hour
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
		TTL: time.Hour,
	}

	// Get classifier values
	key, err := c.GetKey()
	description, err := c.GetDescription()
	value, err := c.GetValue()

	if err != nil {
		return err
	}

	// Create a simple classifier and serialize to JSON
	klassifier := NewSimpleClassifier(key, description, value)
	data, err := json.Marshal(klassifier)

	if err != nil {
		log.Printf("Failed to serialize classifier: %s\n", key)
		return err
	}

	// Set minion classifier in etcd
	klassifierKey := filepath.Join(m.classifierDir, key)
	_, err = m.kapi.Set(context.Background(), klassifierKey, string(data), opts)

	if err != nil {
		log.Printf("Failed to set classifier %s: %s\n", key, err)
	}

	return err
}

// Monitors etcd for new tasks for processing
func (m *etcdMinion) TaskListener(c chan<- *MinionTask) error {
	watcherOpts := &etcdclient.WatcherOptions{
		Recursive: true,
	}
	watcher := m.kapi.Watcher(m.queueDir, watcherOpts)

	for {
		resp, err := watcher.Next(context.Background())
		if err != nil {
			log.Printf("Failed to read task: %s\n", err)
			continue
		}

		// Ignore "delete" events when removing a task from the queue
		action := strings.ToLower(resp.Action)
		if strings.EqualFold(action, "delete") {
			continue
		}

		// Remove task from the queue
		task, err := EtcdUnmarshalTask(resp.Node)
		m.kapi.Delete(context.Background(), resp.Node.Key, nil)

		if err != nil {
			continue
		}

		log.Printf("Received task %s\n", task.TaskID)

		c <- task
	}

	return nil
}

// Processes new tasks
func (m *etcdMinion) TaskRunner(c <-chan *MinionTask) error {
	for {
		task := <-c

		task.TimeReceived = time.Now().Unix()

		if task.IsConcurrent {
			go m.processTask(task)
		} else {
			m.processTask(task)
		}
	}

	return nil
}

// Main entry point of the minion
func (m *etcdMinion) Serve() error {
	// Channel on which we send the quit signal
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, os.Interrupt)

	// Initialize minion
	m.setName()

	log.Printf("Minion %s is ready to serve", m.id)

	// Run periodic tasks every fifteen minutes
	ticker := time.NewTicker(time.Minute * 15)
	go m.periodicRunner(ticker)

	// Check for pending tasks in queue
	tasks := make(chan *MinionTask)
	m.checkQueue(tasks)

	go m.TaskListener(tasks)
	go m.TaskRunner(tasks)

	// Block until a stop signal is received
	s := <-quit
	log.Printf("Received %s signal, shutting down", s)
	close(quit)
	close(tasks)

	return nil
}
<commit_msg>Start the task runner early so that we can start processing tasks<commit_after>package minion

import (
	"os"
	"os/exec"
	"os/signal"
	"log"
	"bytes"
	"time"
	"strings"
	"strconv"
	"path/filepath"
	"encoding/json"

	"code.google.com/p/go-uuid/uuid"
	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
	etcdclient "github.com/coreos/etcd/client"
)

// Minions keyspace in etcd
const EtcdMinionSpace = "/gru/minion"

// Etcd Minion
type etcdMinion struct {
	// Name of this minion
	name string

	// Minion root node in etcd
	rootDir string

	// Minion queue node in etcd
	queueDir string

	// Log directory to keep previously executed tasks
	logDir string

	// Root node for classifiers in etcd
	classifierDir string

	// Minion unique identifier
	id uuid.UUID

	// KeysAPI client to etcd
	kapi etcdclient.KeysAPI
}

// Creates a new etcd minion
func NewEtcdMinion(name string, cfg etcdclient.Config) Minion {
	c, err := etcdclient.New(cfg)
	if err != nil {
		log.Fatal(err)
	}

	kapi := etcdclient.NewKeysAPI(c)
	id := GenerateUUID(name)
	rootDir := filepath.Join(EtcdMinionSpace, id.String())
	queueDir := filepath.Join(rootDir, "queue")
	classifierDir := filepath.Join(rootDir, "classifier")
	logDir := filepath.Join(rootDir, "log")

	m := &etcdMinion{
		name: name,
		rootDir: rootDir,
		queueDir: queueDir,
		classifierDir: classifierDir,
		logDir: logDir,
		id: id,
		kapi: kapi,
	}

	return m
}

// Set the human-readable name of the minion in etcd
func (m *etcdMinion) setName() error {
	nameKey := filepath.Join(m.rootDir, "name")
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), nameKey, m.name, opts)

	return err
}

// Set the time the minion was last seen in seconds since the Epoch
func (m *etcdMinion) setLastseen(s int64) error {
	lastseenKey := filepath.Join(m.rootDir, "lastseen")
	lastseenValue := strconv.FormatInt(s, 10)
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), lastseenKey, lastseenValue, opts)

	return err
}

// Checks for any tasks pending tasks in queue
func (m *etcdMinion) checkQueue(c chan<- *MinionTask) error {
	opts := &etcdclient.GetOptions{
		Recursive: true,
		Sort: true,
	}

	// Get backlog tasks if any
	resp, err := m.kapi.Get(context.Background(), m.queueDir, opts)
	if err != nil {
		return nil
	}

	backlog := resp.Node.Nodes
	if len(backlog) == 0 {
		// No backlog tasks found
		return nil
	}

	log.Printf("Found %d tasks in backlog", len(backlog))
	for _, node := range backlog {
		task, err := EtcdUnmarshalTask(node)
		m.kapi.Delete(context.Background(), node.Key, nil)

		if err != nil {
			continue
		}

		c <- task
	}

	return nil
}

// Runs periodic jobs such as refreshing classifiers and updating lastseen
func (m *etcdMinion) periodicRunner(ticker *time.Ticker) error {
	for {
		// Update classifiers
		for _, classifier := range ClassifierRegistry {
			m.SetClassifier(classifier)
		}

		// Update lastseen time
		now := time.Now().Unix()
		err := m.setLastseen(now)
		if err != nil {
			log.Printf("Failed to update lastseen time: %s\n", err)
		}

		<- ticker.C
	}

	return nil
}

// Processes new tasks
func (m *etcdMinion) processTask(t *MinionTask) error {
	defer m.saveTask(t)

	var buf bytes.Buffer
	cmd := exec.Command(t.Command, t.Args...)
	cmd.Stdout = &buf
	cmd.Stderr = &buf

	log.Printf("Processing task %s\n", t.TaskID)

	cmdError := cmd.Run()
	t.TimeProcessed = time.Now().Unix()
	t.Result = buf.String()

	if cmdError != nil {
		log.Printf("Failed to process task %s\n", t.TaskID)
		t.Error = cmdError.Error()
	} else {
		log.Printf("Finished processing task %s\n", t.TaskID)
	}

	return cmdError
}

// Saves a task in the minion's log
func (m *etcdMinion) saveTask(t *MinionTask) error {
	// Task key in etcd
	taskKey := filepath.Join(m.logDir, t.TaskID.String())

	// Serialize task to JSON
	data, err := json.Marshal(t)
	if err != nil {
		log.Printf("Failed to serialize task %s: %s\n", t.TaskID, err)
		return err
	}

	// Save task result in the minion's space
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}
	_, err = m.kapi.Set(context.Background(), taskKey, string(data), opts)
	if err != nil {
		log.Printf("Failed to save task %s: %s\n", t.TaskID, err)
		return err
	}

	return err
}

// Unmarshals task from etcd
func EtcdUnmarshalTask(node *etcdclient.Node) (*MinionTask, error) {
	task := new(MinionTask)
	err := json.Unmarshal([]byte(node.Value), &task)

	if err != nil {
		log.Printf("Invalid task %s: %s\n", node.Key, err)
	}

	return task, err
}

// Returns the minion unique identifier
func (m *etcdMinion) ID() uuid.UUID {
	return m.id
}

// Returns the assigned name of the minion
func (m *etcdMinion) Name() string {
	return m.name
}

// Classify a minion with a given key and value
func (m *etcdMinion) SetClassifier(c MinionClassifier) error {
	// Classifiers in etcd expire after an hour
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
		TTL: time.Hour,
	}

	// Get classifier values
	key, err := c.GetKey()
	description, err := c.GetDescription()
	value, err := c.GetValue()

	if err != nil {
		return err
	}

	// Create a simple classifier and serialize to JSON
	klassifier := NewSimpleClassifier(key, description, value)
	data, err := json.Marshal(klassifier)

	if err != nil {
		log.Printf("Failed to serialize classifier: %s\n", key)
		return err
	}

	// Set minion classifier in etcd
	klassifierKey := filepath.Join(m.classifierDir, key)
	_, err = m.kapi.Set(context.Background(), klassifierKey, string(data), opts)

	if err != nil {
		log.Printf("Failed to set classifier %s: %s\n", key, err)
	}

	return err
}

// Monitors etcd for new tasks for processing
func (m *etcdMinion) TaskListener(c chan<- *MinionTask) error {
	watcherOpts := &etcdclient.WatcherOptions{
		Recursive: true,
	}
	watcher := m.kapi.Watcher(m.queueDir, watcherOpts)

	for {
		resp, err := watcher.Next(context.Background())
		if err != nil {
			log.Printf("Failed to read task: %s\n", err)
			continue
		}

		// Ignore "delete" events when removing a task from the queue
		action := strings.ToLower(resp.Action)
		if strings.EqualFold(action, "delete") {
			continue
		}

		// Remove task from the queue
		task, err := EtcdUnmarshalTask(resp.Node)
		m.kapi.Delete(context.Background(), resp.Node.Key, nil)

		if err != nil {
			continue
		}

		log.Printf("Received task %s\n", task.TaskID)

		c <- task
	}

	return nil
}

// Processes new tasks
func (m *etcdMinion) TaskRunner(c <-chan *MinionTask) error {
	for {
		task := <-c

		task.TimeReceived = time.Now().Unix()

		if task.IsConcurrent {
			go m.processTask(task)
		} else {
			m.processTask(task)
		}
	}

	return nil
}

// Main entry point of the minion
func (m *etcdMinion) Serve() error {
	// Channel on which we send the quit signal
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, os.Interrupt)

	// Initialize minion
	m.setName()

	log.Printf("Minion %s is ready to serve", m.id)

	// Run periodic tasks every fifteen minutes
	ticker := time.NewTicker(time.Minute * 15)
	go m.periodicRunner(ticker)

	// Check for pending tasks in queue
	tasks := make(chan *MinionTask)
	go m.TaskRunner(tasks)
	m.checkQueue(tasks)

	go m.TaskListener(tasks)

	// Block until a stop signal is received
	s := <-quit
	log.Printf("Received %s signal, shutting down", s)
	close(quit)
	close(tasks)

	return nil
}
<|endoftext|>


<commit_before><commit_msg>EnsureLoadBalancer update instead of recreate existing LBs<commit_after><|endoftext|>


<commit_before>package shadowsocks

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"

	sscore "github.com/shadowsocks/go-shadowsocks2/core"
	sssocks "github.com/shadowsocks/go-shadowsocks2/socks"

	tun2socks "github.com/eycorsican/go-tun2socks"
	"github.com/eycorsican/go-tun2socks/lwip"
)

type tcpHandler struct {
	sync.Mutex

	cipher   sscore.Cipher
	server   string
	conns    map[tun2socks.Connection]net.Conn
	tgtAddrs map[tun2socks.Connection]net.Addr
	tgtSent  map[tun2socks.Connection]bool
}

func (h *tcpHandler) getConn(conn tun2socks.Connection) (net.Conn, bool) {
	h.Lock()
	defer h.Unlock()
	if c, ok := h.conns[conn]; ok {
		return c, true
	}
	return nil, false
}

func (h *tcpHandler) fetchInput(conn tun2socks.Connection, input io.Reader) {
	buf := lwip.NewBytes(lwip.BufSize)

	defer func() {
		h.Close(conn)
		lwip.FreeBytes(buf)
	}()

	for {
		n, err := input.Read(buf)
		if err != nil {
			if err != io.EOF {
				log.Printf("failed to read from Shadowsocks server: %v", err)
				h.Close(conn)
				return
			}
			break
		}
		err = conn.Write(buf[:n])
		if err != nil {
			log.Printf("failed to send data to TUN: %v", err)
			h.Close(conn)
			return
		}
	}
}

func NewTCPHandler(server, cipher, password string) tun2socks.ConnectionHandler {
	ciph, err := sscore.PickCipher(cipher, []byte{}, password)
	if err != nil {
		log.Fatal(err)
	}

	return &tcpHandler{
		cipher:   ciph,
		server:   server,
		conns:    make(map[tun2socks.Connection]net.Conn, 16),
		tgtAddrs: make(map[tun2socks.Connection]net.Addr, 16),
		tgtSent:  make(map[tun2socks.Connection]bool, 16),
	}
}

func (h *tcpHandler) sendTargetAddress(conn tun2socks.Connection) error {
	h.Lock()
	tgtAddr, ok1 := h.tgtAddrs[conn]
	rc, ok2 := h.conns[conn]
	h.Unlock()
	if ok1 && ok2 {
		tgt := sssocks.ParseAddr(tgtAddr.String())
		_, err := rc.Write(tgt)
		if err != nil {
			return errors.New(fmt.Sprintf("sending target address failed: %v", err))
		}
		h.tgtSent[conn] = true
		go h.fetchInput(conn, rc)
	} else {
		return errors.New("target address not found")
	}
	return nil
}

func (h *tcpHandler) Connect(conn tun2socks.Connection, target net.Addr) error {
	rc, err := net.Dial("tcp", h.server)
	if err != nil {
		return errors.New(fmt.Sprintf("dialing remote server failed: %v", err))
	}
	rc = h.cipher.StreamConn(rc)

	h.Lock()
	h.conns[conn] = rc
	h.tgtAddrs[conn] = target
	h.Unlock()
	rc.SetDeadline(time.Time{})
	return nil
}

func (h *tcpHandler) DidReceive(conn tun2socks.Connection, data []byte) error {
	h.Lock()
	rc, ok1 := h.conns[conn]
	sent, ok2 := h.tgtSent[conn]
	h.Unlock()

	if ok1 {
		if !ok2 || !sent {
			h.sendTargetAddress(conn)
		}

		_, err := rc.Write(data)
		if err != nil {
			log.Printf("failed to write data to Shadowsocks server: %v", err)
			h.Close(conn)
			return errors.New(fmt.Sprintf("failed to write data: %v", err))
		}
		return nil
	} else {
		return errors.New(fmt.Sprintf("proxy connection does not exists: %v <-> %v", conn.LocalAddr().String(), conn.RemoteAddr().String()))
	}
}

func (h *tcpHandler) DidSend(conn tun2socks.Connection, len uint16) {
}

func (h *tcpHandler) DidClose(conn tun2socks.Connection) {
	h.Close(conn)
}

func (h *tcpHandler) DidAbort(conn tun2socks.Connection) {
	h.Close(conn)
}

func (h *tcpHandler) DidReset(conn tun2socks.Connection) {
	h.Close(conn)
}

func (h *tcpHandler) LocalDidClose(conn tun2socks.Connection) {
	h.Close(conn)
}

func (h *tcpHandler) Close(conn tun2socks.Connection) {
	if rc, found := h.getConn(conn); found {
		rc.Close()
		h.Lock()
		delete(h.conns, conn)
		h.Unlock()
	}
	delete(h.tgtAddrs, conn)
	delete(h.tgtSent, conn)
}
<commit_msg>fix concurrent accessing map<commit_after>package shadowsocks

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"

	sscore "github.com/shadowsocks/go-shadowsocks2/core"
	sssocks "github.com/shadowsocks/go-shadowsocks2/socks"

	tun2socks "github.com/eycorsican/go-tun2socks"
	"github.com/eycorsican/go-tun2socks/lwip"
)

type tcpHandler struct {
	sync.Mutex

	cipher   sscore.Cipher
	server   string
	conns    map[tun2socks.Connection]net.Conn
	tgtAddrs map[tun2socks.Connection]net.Addr
	tgtSent  map[tun2socks.Connection]bool
}

func (h *tcpHandler) getConn(conn tun2socks.Connection) (net.Conn, bool) {
	h.Lock()
	defer h.Unlock()
	if c, ok := h.conns[conn]; ok {
		return c, true
	}
	return nil, false
}

func (h *tcpHandler) fetchInput(conn tun2socks.Connection, input io.Reader) {
	buf := lwip.NewBytes(lwip.BufSize)

	defer func() {
		h.Close(conn)
		lwip.FreeBytes(buf)
	}()

	for {
		n, err := input.Read(buf)
		if err != nil {
			if err != io.EOF {
				log.Printf("failed to read from Shadowsocks server: %v", err)
				h.Close(conn)
				return
			}
			break
		}
		err = conn.Write(buf[:n])
		if err != nil {
			log.Printf("failed to send data to TUN: %v", err)
			h.Close(conn)
			return
		}
	}
}

func NewTCPHandler(server, cipher, password string) tun2socks.ConnectionHandler {
	ciph, err := sscore.PickCipher(cipher, []byte{}, password)
	if err != nil {
		log.Fatal(err)
	}

	return &tcpHandler{
		cipher:   ciph,
		server:   server,
		conns:    make(map[tun2socks.Connection]net.Conn, 16),
		tgtAddrs: make(map[tun2socks.Connection]net.Addr, 16),
		tgtSent:  make(map[tun2socks.Connection]bool, 16),
	}
}

func (h *tcpHandler) sendTargetAddress(conn tun2socks.Connection) error {
	h.Lock()
	tgtAddr, ok1 := h.tgtAddrs[conn]
	rc, ok2 := h.conns[conn]
	h.Unlock()
	if ok1 && ok2 {
		tgt := sssocks.ParseAddr(tgtAddr.String())
		_, err := rc.Write(tgt)
		if err != nil {
			return errors.New(fmt.Sprintf("sending target address failed: %v", err))
		}
		h.tgtSent[conn] = true
		go h.fetchInput(conn, rc)
	} else {
		return errors.New("target address not found")
	}
	return nil
}

func (h *tcpHandler) Connect(conn tun2socks.Connection, target net.Addr) error {
	rc, err := net.Dial("tcp", h.server)
	if err != nil {
		return errors.New(fmt.Sprintf("dialing remote server failed: %v", err))
	}
	rc = h.cipher.StreamConn(rc)

	h.Lock()
	h.conns[conn] = rc
	h.tgtAddrs[conn] = target
	h.Unlock()
	rc.SetDeadline(time.Time{})
	return nil
}

func (h *tcpHandler) DidReceive(conn tun2socks.Connection, data []byte) error {
	h.Lock()
	rc, ok1 := h.conns[conn]
	sent, ok2 := h.tgtSent[conn]
	h.Unlock()

	if ok1 {
		if !ok2 || !sent {
			h.sendTargetAddress(conn)
		}

		_, err := rc.Write(data)
		if err != nil {
			log.Printf("failed to write data to Shadowsocks server: %v", err)
			h.Close(conn)
			return errors.New(fmt.Sprintf("failed to write data: %v", err))
		}
		return nil
	} else {
		return errors.New(fmt.Sprintf("proxy connection does not exists: %v <-> %v", conn.LocalAddr().String(), conn.RemoteAddr().String()))
	}
}

func (h *tcpHandler) DidSend(conn tun2socks.Connection, len uint16) {
}

func (h *tcpHandler) DidClose(conn tun2socks.Connection) {
	h.Close(conn)
}

func (h *tcpHandler) DidAbort(conn tun2socks.Connection) {
	h.Close(conn)
}

func (h *tcpHandler) DidReset(conn tun2socks.Connection) {
	h.Close(conn)
}

func (h *tcpHandler) LocalDidClose(conn tun2socks.Connection) {
	h.Close(conn)
}

func (h *tcpHandler) Close(conn tun2socks.Connection) {
	if rc, found := h.getConn(conn); found {
		rc.Close()
		h.Lock()
		delete(h.conns, conn)
		h.Unlock()
	}

	h.Lock()
	defer h.Unlock()

	delete(h.tgtAddrs, conn)
	delete(h.tgtSent, conn)
}
<|endoftext|>


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

package main

import (
	"bytes"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"testing"
)

func TestBind(t *testing.T) {
	// mk && rm -rf toto $TMPDIR/gopy-* && gopy bind -output=./toto ./_examples/hi && (echo "=== testing..."; cd toto; cp ../_examples/hi/test.py .; python2 ./test.py && echo "[ok]" || echo "ERR")
	//

	workdir, err := ioutil.TempDir("", "gopy-")
	if err != nil {
		t.Fatalf("could not create workdir: %v\n", err)
	}
	err = os.MkdirAll(workdir, 0644)
	if err != nil {
		t.Fatalf("could not create workdir: %v\n", err)
	}
	defer os.RemoveAll(workdir)

	cmd := exec.Command("gopy", "bind", "-output="+workdir, "./_examples/hi")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err = cmd.Run()
	if err != nil {
		t.Fatalf("error running gopy-bind: %v\n", err)
	}

	cmd = exec.Command(
		"/bin/cp", "./_examples/hi/test.py",
		filepath.Join(workdir, "test.py"),
	)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err = cmd.Run()
	if err != nil {
		t.Fatalf("error copying 'test.py': %v\n", err)
	}

	want := []byte(`hi from go
hello you from go
--- doc(hi)...
package hi exposes a few Go functions to be wrapped and used from Python.

--- doc(hi.Hi)...
Hi()

Hi prints hi from Go

--- hi.Hi()...
--- doc(hi.Hello)...
Hello(str s)

Hello prints a greeting from Go

--- hi.Hello('you')...
--- doc(hi.Add)...
Add(int i, int j) int

Add returns the sum of its arguments.

--- hi.Add(1, 41)...
42
--- hi.Concat('4', '2')...
42
--- doc(hi.Person):
Person is a simple struct

--- p = hi.Person()...
['Age', 'Greet', 'Name', 'String', '__class__', '__delattr__', '__doc__', '__format__', '__getattribute__', '__hash__', '__init__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__']
--- p.Name: None
--- p.Age: None
--- doc(hi.Greet):
Greet() str

Greet sends greetings

--- p.Greet()...
None
--- doc(p):
Person is a simple struct

`)
	buf := new(bytes.Buffer)
	cmd = exec.Command("python2", "./test.py")
	cmd.Dir = workdir
	cmd.Stdin = os.Stdin
	cmd.Stdout = buf
	cmd.Stderr = buf
	err = cmd.Run()
	if err != nil {
		t.Fatalf(
			"error running python module: %v\n%v\n", err,
			string(buf.Bytes()),
		)
	}

	if !reflect.DeepEqual(string(buf.Bytes()), string(want)) {
		t.Fatalf("error running python module:\nwant:\n%s\n\ngot:\n%s\n",
			string(want), string(buf.Bytes()),
		)
	}
}
<commit_msg>test: update<commit_after>// Copyright 2015 The go-python Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"testing"
)

func TestBind(t *testing.T) {
	// mk && rm -rf toto $TMPDIR/gopy-* && gopy bind -output=./toto ./_examples/hi && (echo "=== testing..."; cd toto; cp ../_examples/hi/test.py .; python2 ./test.py && echo "[ok]" || echo "ERR")
	//

	workdir, err := ioutil.TempDir("", "gopy-")
	if err != nil {
		t.Fatalf("could not create workdir: %v\n", err)
	}
	err = os.MkdirAll(workdir, 0644)
	if err != nil {
		t.Fatalf("could not create workdir: %v\n", err)
	}
	defer os.RemoveAll(workdir)

	cmd := exec.Command("gopy", "bind", "-output="+workdir, "./_examples/hi")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err = cmd.Run()
	if err != nil {
		t.Fatalf("error running gopy-bind: %v\n", err)
	}

	cmd = exec.Command(
		"/bin/cp", "./_examples/hi/test.py",
		filepath.Join(workdir, "test.py"),
	)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err = cmd.Run()
	if err != nil {
		t.Fatalf("error copying 'test.py': %v\n", err)
	}

	want := []byte(`hi from go
hello you from go
--- doc(hi)...
package hi exposes a few Go functions to be wrapped and used from Python.

--- doc(hi.Hi)...
Hi()

Hi prints hi from Go

--- hi.Hi()...
--- doc(hi.Hello)...
Hello(str s)

Hello prints a greeting from Go

--- hi.Hello('you')...
--- doc(hi.Add)...
Add(int i, int j) int

Add returns the sum of its arguments.

--- hi.Add(1, 41)...
42
--- hi.Concat('4', '2')...
42
--- doc(hi.Person):
Person is a simple struct

--- p = hi.Person()...
['Age', 'Greet', 'Name', 'String', '__class__', '__delattr__', '__doc__', '__format__', '__getattribute__', '__hash__', '__init__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__']
--- p.Name: None
--- p.Age: None
--- doc(hi.Greet):
Greet() str

Greet sends greetings

--- p.Greet()...
Hello, I am
--- doc(p):
Person is a simple struct

`)
	buf := new(bytes.Buffer)
	cmd = exec.Command("python2", "./test.py")
	cmd.Dir = workdir
	cmd.Stdin = os.Stdin
	cmd.Stdout = buf
	cmd.Stderr = buf
	err = cmd.Run()
	if err != nil {
		t.Fatalf(
			"error running python module: %v\n%v\n", err,
			string(buf.Bytes()),
		)
	}

	if !reflect.DeepEqual(string(buf.Bytes()), string(want)) {
		t.Fatalf("error running python module:\nwant:\n%s\n\ngot:\n%s\n",
			string(want), string(buf.Bytes()),
		)
	}
}
<|endoftext|>


<commit_before>package main_test

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"syscall"
	"time"

	"github.com/cloudfoundry-incubator/runtime-schema/models"
	"github.com/cloudfoundry/yagnats"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/rata"

	"github.com/cloudfoundry-incubator/tps/api"
	"github.com/cloudfoundry-incubator/tps/heartbeat"
)

var _ = Describe("TPS", func() {

	var httpClient *http.Client
	var requestGenerator *rata.RequestGenerator
	var natsClient yagnats.NATSClient

	BeforeEach(func() {
		natsClient = natsRunner.MessageBus
		requestGenerator = rata.NewRequestGenerator(fmt.Sprintf("http://%s", tpsAddr), api.Routes)
		httpClient = &http.Client{
			Transport: &http.Transport{},
		}
	})

	Describe("GET /lrps/:guid", func() {
		Context("when etcd is running", func() {
			BeforeEach(func() {
				bbs.ReportActualLRPAsStarting("some-process-guid", "some-instance-guid-1", "executor-id", 0)

				bbs.ReportActualLRPAsRunning(models.ActualLRP{
					ProcessGuid:  "some-process-guid",
					InstanceGuid: "some-instance-guid-2",

					Index: 1,

					State: models.ActualLRPStateRunning,
				}, "executor-id")

				bbs.ReportActualLRPAsRunning(models.ActualLRP{
					ProcessGuid:  "some-process-guid",
					InstanceGuid: "some-instance-guid-3",

					Index: 2,

					State: models.ActualLRPStateRunning,
				}, "executor-id")

				bbs.ReportActualLRPAsRunning(models.ActualLRP{
					ProcessGuid:  "some-other-process-guid",
					InstanceGuid: "some-instance-guid-3",

					Index: 0,

					State: models.ActualLRPStateRunning,
				}, "executor-id")
			})

			It("reports the state of the given process guid's instances", func() {
				getLRPs, err := requestGenerator.CreateRequest(
					api.LRPStatus,
					rata.Params{"guid": "some-process-guid"},
					nil,
				)
				Ω(err).ShouldNot(HaveOccurred())

				response, err := httpClient.Do(getLRPs)
				Ω(err).ShouldNot(HaveOccurred())

				var lrpInstances []api.LRPInstance
				err = json.NewDecoder(response.Body).Decode(&lrpInstances)
				Ω(err).ShouldNot(HaveOccurred())

				Ω(lrpInstances).Should(HaveLen(3))

				Ω(lrpInstances).Should(ContainElement(api.LRPInstance{
					ProcessGuid:  "some-process-guid",
					InstanceGuid: "some-instance-guid-1",

					Index: 0,
					Since: timeProvider.Time().UnixNano(),

					State: "starting",
				}))

				Ω(lrpInstances).Should(ContainElement(api.LRPInstance{
					ProcessGuid:  "some-process-guid",
					InstanceGuid: "some-instance-guid-2",

					Index: 1,
					Since: timeProvider.Time().UnixNano(),

					State: "running",
				}))

				Ω(lrpInstances).Should(ContainElement(api.LRPInstance{
					ProcessGuid:  "some-process-guid",
					InstanceGuid: "some-instance-guid-3",

					Index: 2,
					Since: timeProvider.Time().UnixNano(),

					State: "running",
				}))
			})
		})

		Context("when etcd is not running", func() {
			BeforeEach(func() {
				etcdRunner.Stop()
			})

			It("returns 500", func() {
				getLRPs, err := requestGenerator.CreateRequest(
					api.LRPStatus,
					rata.Params{"guid": "some-process-guid"},
					nil,
				)
				Ω(err).ShouldNot(HaveOccurred())

				response, err := httpClient.Do(getLRPs)
				Ω(err).ShouldNot(HaveOccurred())

				Ω(response.StatusCode).Should(Equal(http.StatusInternalServerError))
			})
		})
	})

	Context("when the NATS server is running", func() {
		var tpsNatsSubject = "service.announce.tps"
		var announceMsg chan *yagnats.Message

		BeforeEach(func() {
			announceMsg = make(chan *yagnats.Message)
			natsClient.Subscribe(tpsNatsSubject, func(msg *yagnats.Message) {
				announceMsg <- msg
			})
		})

		AfterEach(func() {
			natsClient.UnsubscribeAll(tpsNatsSubject)
		})

		It("heartbeats announcement messages at the predefined interval", func() {
			Eventually(announceMsg, heartbeatInterval+time.Second).Should(Receive())
			Eventually(announceMsg, heartbeatInterval+time.Second).Should(Receive())
		})

		Describe("published HeartbeatMessage", func() {
			var heartbeatMsg heartbeat.HeartbeatMessage

			BeforeEach(func(done Done) {
				heartbeatMsg = heartbeat.HeartbeatMessage{}
				msg := <-announceMsg
				err := json.Unmarshal(msg.Payload, &heartbeatMsg)
				Ω(err).ShouldNot(HaveOccurred())
				close(done)
			})

			It("contains the correct tps address", func() {
				Ω(heartbeatMsg.Addr).Should(Equal(fmt.Sprintf("http://%s", tpsAddr)))
			})

			It("a ttl 3 times longer than the heartbeatInterval, in seconds", func() {
				Ω(heartbeatMsg.TTL).Should(Equal(uint(3)))
			})
		})
	})

	Context("when the NATS server is down while starting up", func() {
		BeforeEach(func() {
			tps.Signal(os.Kill)
			Eventually(tps.Wait()).Should(Receive(nil))

			natsRunner.KillWithFire()

			tps = ifrit.Envoke(runner)
		})

		It("exits imediately", func() {
			Eventually(tps.Wait()).Should(Receive())
		})
	})

	Context("when the NATS server goes down after startup", func() {
		BeforeEach(func() {
			natsRunner.KillWithFire()
			time.Sleep(50 * time.Millisecond)
		})

		It("does not exit", func() {
			Consistently(tps.Wait()).ShouldNot(Receive())
		})

		It("exits when we send a signal", func() {
			tps.Signal(syscall.SIGINT)
			Eventually(tps.Wait()).Should(Receive())
		})
	})
})
<commit_msg>Update nats client to use yagnats apcera/nats wrapper [#76637390]<commit_after>package main_test

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"syscall"
	"time"

	"github.com/cloudfoundry-incubator/runtime-schema/models"
	"github.com/cloudfoundry/yagnats"
	"github.com/apcera/nats"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/rata"

	"github.com/cloudfoundry-incubator/tps/api"
	"github.com/cloudfoundry-incubator/tps/heartbeat"
)

var _ = Describe("TPS", func() {

	var httpClient *http.Client
	var requestGenerator *rata.RequestGenerator
	var natsClient yagnats.ApceraWrapperNATSClient

	BeforeEach(func() {
		natsClient = natsRunner.MessageBus
		requestGenerator = rata.NewRequestGenerator(fmt.Sprintf("http://%s", tpsAddr), api.Routes)
		httpClient = &http.Client{
			Transport: &http.Transport{},
		}
	})

	Describe("GET /lrps/:guid", func() {
		Context("when etcd is running", func() {
			BeforeEach(func() {
				bbs.ReportActualLRPAsStarting("some-process-guid", "some-instance-guid-1", "executor-id", 0)

				bbs.ReportActualLRPAsRunning(models.ActualLRP{
					ProcessGuid:  "some-process-guid",
					InstanceGuid: "some-instance-guid-2",

					Index: 1,

					State: models.ActualLRPStateRunning,
				}, "executor-id")

				bbs.ReportActualLRPAsRunning(models.ActualLRP{
					ProcessGuid:  "some-process-guid",
					InstanceGuid: "some-instance-guid-3",

					Index: 2,

					State: models.ActualLRPStateRunning,
				}, "executor-id")

				bbs.ReportActualLRPAsRunning(models.ActualLRP{
					ProcessGuid:  "some-other-process-guid",
					InstanceGuid: "some-instance-guid-3",

					Index: 0,

					State: models.ActualLRPStateRunning,
				}, "executor-id")
			})

			It("reports the state of the given process guid's instances", func() {
				getLRPs, err := requestGenerator.CreateRequest(
					api.LRPStatus,
					rata.Params{"guid": "some-process-guid"},
					nil,
				)
				Ω(err).ShouldNot(HaveOccurred())

				response, err := httpClient.Do(getLRPs)
				Ω(err).ShouldNot(HaveOccurred())

				var lrpInstances []api.LRPInstance
				err = json.NewDecoder(response.Body).Decode(&lrpInstances)
				Ω(err).ShouldNot(HaveOccurred())

				Ω(lrpInstances).Should(HaveLen(3))

				Ω(lrpInstances).Should(ContainElement(api.LRPInstance{
					ProcessGuid:  "some-process-guid",
					InstanceGuid: "some-instance-guid-1",

					Index: 0,
					Since: timeProvider.Time().UnixNano(),

					State: "starting",
				}))

				Ω(lrpInstances).Should(ContainElement(api.LRPInstance{
					ProcessGuid:  "some-process-guid",
					InstanceGuid: "some-instance-guid-2",

					Index: 1,
					Since: timeProvider.Time().UnixNano(),

					State: "running",
				}))

				Ω(lrpInstances).Should(ContainElement(api.LRPInstance{
					ProcessGuid:  "some-process-guid",
					InstanceGuid: "some-instance-guid-3",

					Index: 2,
					Since: timeProvider.Time().UnixNano(),

					State: "running",
				}))
			})
		})

		Context("when etcd is not running", func() {
			BeforeEach(func() {
				etcdRunner.Stop()
			})

			It("returns 500", func() {
				getLRPs, err := requestGenerator.CreateRequest(
					api.LRPStatus,
					rata.Params{"guid": "some-process-guid"},
					nil,
				)
				Ω(err).ShouldNot(HaveOccurred())

				response, err := httpClient.Do(getLRPs)
				Ω(err).ShouldNot(HaveOccurred())

				Ω(response.StatusCode).Should(Equal(http.StatusInternalServerError))
			})
		})
	})

	Context("when the NATS server is running", func() {
		var tpsNatsSubject = "service.announce.tps"
		var announceMsg chan *nats.Msg
		var subscription *nats.Subscription

		BeforeEach(func() {
			announceMsg = make(chan *nats.Msg)
			var err error
			subscription, err = natsClient.Subscribe(tpsNatsSubject, func(msg *nats.Msg) {
				announceMsg <- msg
			})
			Ω(err).ShouldNot(HaveOccurred())
		})

		AfterEach(func() {
			err := natsClient.Unsubscribe(subscription)
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("heartbeats announcement messages at the predefined interval", func() {
			Eventually(announceMsg, heartbeatInterval+time.Second).Should(Receive())
			Eventually(announceMsg, heartbeatInterval+time.Second).Should(Receive())
		})

		Describe("published HeartbeatMessage", func() {
			var heartbeatMsg heartbeat.HeartbeatMessage

			BeforeEach(func(done Done) {
				heartbeatMsg = heartbeat.HeartbeatMessage{}
				msg := <-announceMsg
				err := json.Unmarshal(msg.Data, &heartbeatMsg)
				Ω(err).ShouldNot(HaveOccurred())
				close(done)
			})

			It("contains the correct tps address", func() {
				Ω(heartbeatMsg.Addr).Should(Equal(fmt.Sprintf("http://%s", tpsAddr)))
			})

			It("a ttl 3 times longer than the heartbeatInterval, in seconds", func() {
				Ω(heartbeatMsg.TTL).Should(Equal(uint(3)))
			})
		})
	})

	Context("when the NATS server is down while starting up", func() {
		BeforeEach(func() {
			tps.Signal(os.Kill)
			Eventually(tps.Wait()).Should(Receive(nil))

			natsRunner.KillWithFire()

			tps = ifrit.Envoke(runner)
		})

		It("exits imediately", func() {
			Eventually(tps.Wait()).Should(Receive())
		})
	})

	Context("when the NATS server goes down after startup", func() {
		BeforeEach(func() {
			natsRunner.KillWithFire()
			time.Sleep(50 * time.Millisecond)
		})

		It("does not exit", func() {
			Consistently(tps.Wait()).ShouldNot(Receive())
		})

		It("exits when we send a signal", func() {
			tps.Signal(syscall.SIGINT)
			Eventually(tps.Wait()).Should(Receive())
		})
	})
})
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
	"time"

	"github.com/kiskolabs/heroku-cloudwatch-drain/logger"
	"github.com/kiskolabs/heroku-cloudwatch-drain/logparser"

	"github.com/stretchr/testify/assert"
)

var l = new(LastMessageLogger)
var parseFunc = func(b []byte) (*logparser.LogEntry, error) {
	return &logparser.LogEntry{Time: time.Now(), Message: ""}, nil
}

var app = &App{
	loggers: map[string]logger.Logger{"app": l},
	parse:   parseFunc,
}
var server = httptest.NewServer(app)

func TestRequestMustNotBeGet(t *testing.T) {
	r, err := http.Get(server.URL + "/app")
	assert.NoError(t, err)
	assert.Equal(t, http.StatusBadRequest, r.StatusCode)
}

func TestRequestPathMustBeAppName(t *testing.T) {
	r, err := http.Post(server.URL+"/", "", nil)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusBadRequest, r.StatusCode)
}

func TestBasicAuth(t *testing.T) {
	app.user = "me"
	app.pass = "SECRET"
	defer func() {
		app.user = ""
		app.pass = ""
	}()

	r, err := http.Post(server.URL+"/app", "", nil)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusForbidden, r.StatusCode)

	uri, _ := url.Parse(server.URL)
	uri.User = url.UserPassword("me", "SECRET")

	r, err = http.Post(uri.String()+"/app", "", nil)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusAccepted, r.StatusCode)
}

func TestNoBasicAuth(t *testing.T) {
	r, err := http.Post(server.URL+"/app", "", nil)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusAccepted, r.StatusCode)
}

func TestSingleLogEntry(t *testing.T) {
	app.parse = logparser.Parse
	defer func() {
		app.parse = parseFunc
	}()

	body := bytes.NewBuffer([]byte(`89 <45>1 2016-10-15T08:59:08.723822+00:00 host heroku web.1 - State changed from up to down`))
	r, err := http.Post(server.URL+"/app", "", body)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusAccepted, r.StatusCode)
	assert.Equal(t, "heroku[web.1]: State changed from up to down", l.m)
}

func TestLogEntryWithEmptyLineAtTheEnd(t *testing.T) {
	app.parse = logparser.Parse
	defer func() {
		app.parse = parseFunc
	}()

	body := bytes.NewBuffer([]byte("89 <45>1 2016-10-15T08:59:08.723822+00:00 host heroku web.1 - State changed from up to down\n"))
	r, err := http.Post(server.URL+"/app", "", body)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusAccepted, r.StatusCode)
	assert.Equal(t, "heroku[web.1]: State changed from up to down", l.m)
}

func TestAnsiCodeStripping(t *testing.T) {
	app.parse = logparser.Parse
	app.stripAnsiCodes = true
	defer func() {
		app.parse = parseFunc
		app.stripAnsiCodes = false
	}()

	body := bytes.NewBuffer([]byte(`89 <45>1 2016-10-15T08:59:08.723822+00:00 host heroku web.1 - (0.1ms) BEGIN`))
	r, err := http.Post(server.URL+"/app", "", body)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusAccepted, r.StatusCode)
	assert.Equal(t, "heroku[web.1]: (0.1ms) BEGIN", l.m)
}

type LastMessageLogger struct {
	m string
}

func (l *LastMessageLogger) Log(t time.Time, s string) {
	l.m = s
}

func (l *LastMessageLogger) Stop() {}
<commit_msg>Try to fix the build<commit_after>package main

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
	"time"

	"github.com/kiskolabs/heroku-cloudwatch-drain/logger"
	"github.com/kiskolabs/heroku-cloudwatch-drain/logparser"

	"github.com/stretchr/testify/assert"
)

var l = new(LastMessageLogger)
var parseFunc = func(b []byte) (*logparser.LogEntry, error) {
	return &logparser.LogEntry{Time: time.Now(), Message: ""}, nil
}

var app = &App{
	loggers: map[string]logger.Logger{"app": l},
	parse:   parseFunc,
}
var server = httptest.NewServer(app)

func TestRequestNotFoundWithGet(t *testing.T) {
	r, err := http.Get(server.URL + "/app")
	assert.NoError(t, err)
	assert.Equal(t, http.StatusNotFound, r.StatusCode)
}

func TestRequestPathMustBeAppName(t *testing.T) {
	r, err := http.Post(server.URL+"/", "", nil)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusBadRequest, r.StatusCode)
}

func TestBasicAuth(t *testing.T) {
	app.user = "me"
	app.pass = "SECRET"
	defer func() {
		app.user = ""
		app.pass = ""
	}()

	r, err := http.Post(server.URL+"/app", "", nil)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusForbidden, r.StatusCode)

	uri, _ := url.Parse(server.URL)
	uri.User = url.UserPassword("me", "SECRET")

	r, err = http.Post(uri.String()+"/app", "", nil)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusAccepted, r.StatusCode)
}

func TestNoBasicAuth(t *testing.T) {
	r, err := http.Post(server.URL+"/app", "", nil)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusAccepted, r.StatusCode)
}

func TestSingleLogEntry(t *testing.T) {
	app.parse = logparser.Parse
	defer func() {
		app.parse = parseFunc
	}()

	body := bytes.NewBuffer([]byte(`89 <45>1 2016-10-15T08:59:08.723822+00:00 host heroku web.1 - State changed from up to down`))
	r, err := http.Post(server.URL+"/app", "", body)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusAccepted, r.StatusCode)
	assert.Equal(t, "heroku[web.1]: State changed from up to down", l.m)
}

func TestLogEntryWithEmptyLineAtTheEnd(t *testing.T) {
	app.parse = logparser.Parse
	defer func() {
		app.parse = parseFunc
	}()

	body := bytes.NewBuffer([]byte("89 <45>1 2016-10-15T08:59:08.723822+00:00 host heroku web.1 - State changed from up to down\n"))
	r, err := http.Post(server.URL+"/app", "", body)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusAccepted, r.StatusCode)
	assert.Equal(t, "heroku[web.1]: State changed from up to down", l.m)
}

func TestAnsiCodeStripping(t *testing.T) {
	app.parse = logparser.Parse
	app.stripAnsiCodes = true
	defer func() {
		app.parse = parseFunc
		app.stripAnsiCodes = false
	}()

	body := bytes.NewBuffer([]byte(`89 <45>1 2016-10-15T08:59:08.723822+00:00 host heroku web.1 - (0.1ms) BEGIN`))
	r, err := http.Post(server.URL+"/app", "", body)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusAccepted, r.StatusCode)
	assert.Equal(t, "heroku[web.1]: (0.1ms) BEGIN", l.m)
}

type LastMessageLogger struct {
	m string
}

func (l *LastMessageLogger) Log(t time.Time, s string) {
	l.m = s
}

func (l *LastMessageLogger) Stop() {}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"

	"github.com/aryann/difflib"
)

func TestMain(t *testing.T) {
	build, _ := filepath.Abs("build")
	filepath.Walk("test", func(path string, info os.FileInfo, err error) error {
		if strings.HasSuffix(path, ".sh") {
			cmd := exec.Command("bash", filepath.Base(path))
			cmd.Dir = filepath.Dir(path)
			cmd.Env = []string{"PATH=" + build}
			stderr := new(bytes.Buffer)
			cmd.Stderr = stderr
			output, err := cmd.Output()
			if err != nil {
				t.Errorf("FAIL: execution failed: " + path + ": " + err.Error())
			} else {
				outfile := strings.TrimSuffix(path, filepath.Ext(path)) + ".txt"
				expected, err := ioutil.ReadFile(outfile)
				if err != nil {
					t.Errorf("FAIL: error on reading output file: " + outfile)
				} else {
					diffs := difflib.Diff(strings.Split(stderr.String()+string(output), "\n"), strings.Split(string(expected), "\n"))
					help := strings.Contains(string(output), "NAME:")
					differs := false
					for _, diff := range diffs {
						differs = differs || (help && diff.Delta == difflib.RightOnly || !help && diff.Delta != difflib.Common)
					}
					if differs {
						buf := bytes.NewBufferString("")
						for _, diff := range diffs {
							if diff.Delta != difflib.Common {
								buf.WriteString(diff.String() + "\n")
							}
						}
						t.Errorf("FAIL: output differs: " + path + "\n" + buf.String())
					} else {
						t.Logf("PASS: " + path + "\n")
					}
				}
			}
		}
		return nil
	})
}
<commit_msg>improve main_test.go: print stderr output<commit_after>package main

import (
	"bytes"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"

	"github.com/aryann/difflib"
)

func TestMain(t *testing.T) {
	build, _ := filepath.Abs("build")
	filepath.Walk("test", func(path string, info os.FileInfo, err error) error {
		if strings.HasSuffix(path, ".sh") {
			cmd := exec.Command("bash", filepath.Base(path))
			cmd.Dir = filepath.Dir(path)
			cmd.Env = []string{"PATH=" + build}
			stderr := new(bytes.Buffer)
			cmd.Stderr = stderr
			output, err := cmd.Output()
			if err != nil {
				t.Errorf("FAIL: execution failed: " + path + ": " + err.Error() + " " + stderr.String())
			} else {
				outfile := strings.TrimSuffix(path, filepath.Ext(path)) + ".txt"
				expected, err := ioutil.ReadFile(outfile)
				if err != nil {
					t.Errorf("FAIL: error on reading output file: " + outfile)
				} else {
					diffs := difflib.Diff(strings.Split(stderr.String()+string(output), "\n"), strings.Split(string(expected), "\n"))
					help := strings.Contains(string(output), "NAME:")
					differs := false
					for _, diff := range diffs {
						differs = differs || (help && diff.Delta == difflib.RightOnly || !help && diff.Delta != difflib.Common)
					}
					if differs {
						buf := bytes.NewBufferString("")
						for _, diff := range diffs {
							if diff.Delta != difflib.Common {
								buf.WriteString(diff.String() + "\n")
							}
						}
						t.Errorf("FAIL: output differs: " + path + "\n" + buf.String())
					} else {
						t.Logf("PASS: " + path + "\n")
					}
				}
			}
		}
		return nil
	})
}
<|endoftext|>


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

package main

import (
	"log"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
	piazza "github.com/venicegeo/pz-gocommon"
	"github.com/venicegeo/pz-gocommon/elasticsearch"
	"github.com/venicegeo/pz-logger/client"
	"github.com/venicegeo/pz-logger/server"
)

const MOCKING = true

type LoggerTester struct {
	suite.Suite

	sys    *piazza.SystemConfig
	logger client.ILoggerService
}

func (suite *LoggerTester) SetupSuite() {
	t := suite.T()
	assert := assert.New(t)

	var required []piazza.ServiceName
	if MOCKING {
		required = []piazza.ServiceName{}
	} else {
		required = []piazza.ServiceName{piazza.PzElasticSearch}
	}
	sys, err := piazza.NewSystemConfig(piazza.PzLogger, required)
	if err != nil {
		log.Fatal(err)
	}

	suite.sys = sys

	esi, err := elasticsearch.NewIndexInterface(suite.sys, "loggertest$", MOCKING)
	assert.NoError(err)

	_ = sys.StartServer(server.CreateHandlers(sys, esi))

	suite.logger, err = client.NewPzLoggerService(sys)
	if err != nil {
		log.Fatal(err)
	}
}

func (suite *LoggerTester) TearDownSuite() {
	//TODO: kill the go routine running the server
}

func TestRunSuite(t *testing.T) {
	s := &LoggerTester{}
	suite.Run(t, s)
}

func checkMessageArrays(t *testing.T, actualMssgs []client.LogMessage, expectedMssgs []client.LogMessage) {
	assert.Equal(t, len(expectedMssgs), len(actualMssgs), "wrong number of log messages")

	for i := 0; i < len(actualMssgs); i++ {
		assert.EqualValues(t, expectedMssgs[i].Address, actualMssgs[i].Address, "message.address %d not equal", i)
		assert.EqualValues(t, expectedMssgs[i].Message, actualMssgs[i].Message, "message.message %d not equal", i)
		assert.EqualValues(t, expectedMssgs[i].Service, actualMssgs[i].Service, "message.service %d not equal", i)
		assert.EqualValues(t, expectedMssgs[i].Severity, actualMssgs[i].Severity, "message.severity %d not equal", i)
		assert.EqualValues(t, expectedMssgs[i].Time.String(), actualMssgs[i].Time.String(), "message.time %d not equal", i)
		assert.EqualValues(t, expectedMssgs[i].String(), actualMssgs[i].String(), "message.string %d not equal", i)
	}
}

func (suite *LoggerTester) TestElasticsearch() {
	t := suite.T()
	assert := assert.New(t)

	esi, err := elasticsearch.NewIndexInterface(suite.sys, "loggertest$", MOCKING)
	assert.NoError(err)

	version := esi.GetVersion()
	assert.Contains("2.2.0", version)
}

// TODO: this test must come first (to preserve counts & ordering)
func (suite *LoggerTester) TestAAAOne() {
	t := suite.T()
	logger := suite.logger
	assert := assert.New(t)

	var err error
	var actualMssgs []client.LogMessage
	var expectedMssgs []client.LogMessage

	////

	data1 := client.LogMessage{
		Service:  "log-tester",
		Address:  "128.1.2.3",
		Time:     time.Now(),
		Severity: "Info",
		Message:  "The quick brown fox",
	}
	err = logger.LogMessage(&data1)
	assert.NoError(err, "PostToMessages")

	actualMssgs, err = logger.GetFromMessages()
	assert.NoError(err, "GetFromMessages")

	expectedMssgs = []client.LogMessage{data1}
	checkMessageArrays(t, actualMssgs, expectedMssgs)

	////

	data2 := client.LogMessage{
		Service:  "log-tester",
		Address:  "128.0.0.0",
		Time:     time.Now(),
		Severity: "Fatal",
		Message:  "The quick brown fox",
	}

	err = logger.LogMessage(&data2)
	assert.NoError(err, "PostToMessages")

	actualMssgs, err = logger.GetFromMessages()
	assert.NoError(err, "GetFromMessages")

	expectedMssgs = []client.LogMessage{data1, data2}
	checkMessageArrays(t, actualMssgs, expectedMssgs)

	stats, err := logger.GetFromAdminStats()
	assert.NoError(err, "GetFromAdminStats")
	assert.Equal(2, stats.NumMessages, "stats check")
	assert.WithinDuration(time.Now(), stats.StartTime, 10*time.Second, "service start time too long ago")
}

func (suite *LoggerTester) TestHelper() {
	t := suite.T()
	logger := suite.logger
	assert := assert.New(t)

	err := logger.Log("mocktest", "0.0.0.0", client.SeverityInfo, time.Now(), "message from logger unit test via piazza.Log()")
	assert.NoError(err, "pzService.Log()")
}

func (suite *LoggerTester) TestClogger() {
	t := suite.T()
	logger := suite.logger
	assert := assert.New(t)

	clogger := client.NewCustomLogger(&logger, "TestingService", "123 Main St.")
	err := clogger.Debug("a DEBUG message")
	assert.NoError(err)
	err = clogger.Info("a INFO message")
	assert.NoError(err)
	err = clogger.Warn("a WARN message")
	assert.NoError(err)
	err = clogger.Error("an ERROR message")
	assert.NoError(err)
	err = clogger.Fatal("a FATAL message")
	assert.NoError(err)
}

func (suite *LoggerTester) TestAdmin() {
	t := suite.T()
	logger := suite.logger
	assert := assert.New(t)

	settings, err := logger.GetFromAdminSettings()
	assert.NoError(err, "GetFromAdminSettings")
	assert.False(settings.Debug, "settings.Debug")

	settings.Debug = true
	err = logger.PostToAdminSettings(settings)
	assert.NoError(err, "PostToAdminSettings")

	settings, err = logger.GetFromAdminSettings()
	assert.NoError(err, "GetFromAdminSettings")
	assert.True(settings.Debug, "settings.Debug")
}
<commit_msg>cleaned up test functions, added setup/teardown, etc<commit_after>// Copyright 2016, RadiantBlue Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
	piazza "github.com/venicegeo/pz-gocommon"
	"github.com/venicegeo/pz-gocommon/elasticsearch"
	"github.com/venicegeo/pz-logger/client"
	"github.com/venicegeo/pz-logger/server"
)

const MOCKING = true

type LoggerTester struct {
	suite.Suite

	esi    elasticsearch.IIndex
	sys    *piazza.SystemConfig
	logger client.ILoggerService
}

func (suite *LoggerTester) setupFixture() {
	t := suite.T()
	assert := assert.New(t)

	var required []piazza.ServiceName
	if MOCKING {
		required = []piazza.ServiceName{}
	} else {
		required = []piazza.ServiceName{piazza.PzElasticSearch}
	}
	sys, err := piazza.NewSystemConfig(piazza.PzLogger, required)
	assert.NoError(err)
	suite.sys = sys

	esi, err := elasticsearch.NewIndexInterface(sys, "loggertest$", MOCKING)
	assert.NoError(err)
	suite.esi = esi

	_ = sys.StartServer(server.CreateHandlers(sys, esi))

	logger, err := client.NewPzLoggerService(sys)
	assert.NoError(err)
	suite.logger = logger
}

func (suite *LoggerTester) teardownFixture() {
	//TODO: kill the go routine running the server

	suite.esi.Close()
	suite.esi.Delete()
}

func TestRunSuite(t *testing.T) {
	s := &LoggerTester{}
	suite.Run(t, s)
}

func checkMessageArrays(t *testing.T, actualMssgs []client.LogMessage, expectedMssgs []client.LogMessage) {
	assert.Equal(t, len(expectedMssgs), len(actualMssgs), "wrong number of log messages")

	for i := 0; i < len(actualMssgs); i++ {
		assert.EqualValues(t, expectedMssgs[i].Address, actualMssgs[i].Address, "message.address %d not equal", i)
		assert.EqualValues(t, expectedMssgs[i].Message, actualMssgs[i].Message, "message.message %d not equal", i)
		assert.EqualValues(t, expectedMssgs[i].Service, actualMssgs[i].Service, "message.service %d not equal", i)
		assert.EqualValues(t, expectedMssgs[i].Severity, actualMssgs[i].Severity, "message.severity %d not equal", i)
		assert.EqualValues(t, expectedMssgs[i].Time.String(), actualMssgs[i].Time.String(), "message.time %d not equal", i)
		assert.EqualValues(t, expectedMssgs[i].String(), actualMssgs[i].String(), "message.string %d not equal", i)
	}
}

func (suite *LoggerTester) Test01Elasticsearch() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	version := suite.esi.GetVersion()
	assert.Contains("2.2.0", version)
}

// TODO: this test must come first (to preserve counts & ordering)
func (suite *LoggerTester) Test02One() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	logger := suite.logger

	var err error

	data1 := client.LogMessage{
		Service:  "log-tester",
		Address:  "128.1.2.3",
		Time:     time.Now(),
		Severity: "Info",
		Message:  "The quick brown fox",
	}

	data2 := client.LogMessage{
		Service:  "log-tester",
		Address:  "128.0.0.0",
		Time:     time.Now(),
		Severity: "Fatal",
		Message:  "The quick brown fox",
	}

	{
		err = logger.LogMessage(&data1)
		assert.NoError(err, "PostToMessages")
	}

	//	time.Sleep(1 * time.Second)

	{
		actualMssgs, err := logger.GetFromMessages()
		assert.NoError(err, "GetFromMessages")
		assert.Len(actualMssgs, 1)
		expectedMssgs := []client.LogMessage{data1}
		checkMessageArrays(t, actualMssgs, expectedMssgs)
	}

	{
		err = logger.LogMessage(&data2)
		assert.NoError(err, "PostToMessages")
	}

	time.Sleep(4 * time.Second)

	{
		actualMssgs, err := logger.GetFromMessages()
		assert.NoError(err, "GetFromMessages")

		expectedMssgs := []client.LogMessage{data1, data2}
		checkMessageArrays(t, actualMssgs, expectedMssgs)
	}

	{
		stats, err := logger.GetFromAdminStats()
		assert.NoError(err, "GetFromAdminStats")
		assert.Equal(2, stats.NumMessages, "stats check")
		assert.WithinDuration(time.Now(), stats.StartTime, 30*time.Second, "service start time too long ago")
	}
}

func (suite *LoggerTester) Test03Help() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	logger := suite.logger

	err := logger.Log("mocktest", "0.0.0.0", client.SeverityInfo, time.Now(), "message from logger unit test via piazza.Log()")
	assert.NoError(err, "pzService.Log()")
}

func (suite *LoggerTester) Test04Clogger() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	logger := suite.logger

	clogger := client.NewCustomLogger(&logger, "TestingService", "123 Main St.")
	err := clogger.Debug("a DEBUG message")
	assert.NoError(err)
	err = clogger.Info("a INFO message")
	assert.NoError(err)
	err = clogger.Warn("a WARN message")
	assert.NoError(err)
	err = clogger.Error("an ERROR message")
	assert.NoError(err)
	err = clogger.Fatal("a FATAL message")
	assert.NoError(err)
}

func (suite *LoggerTester) Test05Admin() {
	t := suite.T()
	assert := assert.New(t)

	suite.setupFixture()
	defer suite.teardownFixture()

	logger := suite.logger

	settings, err := logger.GetFromAdminSettings()
	assert.NoError(err, "GetFromAdminSettings")
	assert.False(settings.Debug, "settings.Debug")

	settings.Debug = true
	err = logger.PostToAdminSettings(settings)
	assert.NoError(err, "PostToAdminSettings")

	settings, err = logger.GetFromAdminSettings()
	assert.NoError(err, "GetFromAdminSettings")
	assert.True(settings.Debug, "settings.Debug")
}
<|endoftext|>


<commit_before>// Copyright 2014 ALTOROS
// Licensed under the AGPLv3, see LICENSE file for details.

package mock

import (
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"sync"

	"github.com/Altoros/gosigma/data"
)

type DriveLibrary struct {
	s sync.Mutex
	m map[string]*data.Drive
	p string
}

var Drives = &DriveLibrary{
	m: make(map[string]*data.Drive),
	p: "/api/2.0/drives",
}
var LibDrives = &DriveLibrary{
	m: make(map[string]*data.Drive),
	p: "/api/2.0/libdrives",
}

func ResetDrives() {
	Drives.Reset()
	LibDrives.Reset()
}

func InitDrive(d *data.Drive) (*data.Drive, error) {
	if d.UUID == "" {
		uuid, err := GenerateUUID()
		if err != nil {
			return nil, err
		}
		d.UUID = uuid
	}
	if d.Status == "" {
		d.Status = "unmounted"
	}

	return d, nil
}

func (d *DriveLibrary) Add(drv *data.Drive) error {
	drv, err := InitDrive(drv)
	if err != nil {
		return err
	}

	d.s.Lock()
	defer d.s.Unlock()

	d.m[drv.UUID] = drv

	return nil
}

func (d *DriveLibrary) AddDrives(dd []data.Drive) []string {
	d.s.Lock()
	defer d.s.Unlock()

	var result []string
	for _, drv := range dd {
		drv, err := InitDrive(&drv)
		if err != nil {
			d.m[drv.UUID] = drv
			result = append(result, drv.UUID)
		}
	}
	return result
}

func (d *DriveLibrary) Remove(uuid string) bool {
	d.s.Lock()
	defer d.s.Unlock()

	_, ok := d.m[uuid]
	delete(d.m, uuid)

	return ok
}

func (d *DriveLibrary) Reset() {
	d.s.Lock()
	defer d.s.Unlock()
	d.m = make(map[string]*data.Drive)
}

func (d *DriveLibrary) SetStatus(uuid, status string) {
	d.s.Lock()
	defer d.s.Unlock()

	drv, ok := d.m[uuid]
	if ok {
		drv.Status = status
	}
}

var ErrNotFound = errors.New("not found")

func (d *DriveLibrary) Clone(uuid string) (string, error) {
	d.s.Lock()
	defer d.s.Unlock()

	drv, ok := d.m[uuid]
	if !ok {
		return "", ErrNotFound
	}

	newUUID, err := GenerateUUID()
	if err != nil {
		return "", err
	}

	var newDrive data.Drive = *drv
	newDrive.Resource = *data.MakeDriveResource(newUUID)
	newDrive.Status = "unmounted"

	if d == LibDrives {
		Drives.Add(&newDrive)
	} else {
		d.m[newUUID] = &newDrive
	}

	return newUUID, nil
}

func (d *DriveLibrary) handleRequest(w http.ResponseWriter, r *http.Request) {
	path := strings.TrimSuffix(r.URL.Path, "/")
	path = strings.TrimPrefix(path, d.p)
	path = strings.TrimPrefix(path, "/")

	switch r.Method {
	case "GET":
		d.handleGet(w, r, path)
	case "POST":
		d.handlePost(w, r, path)
	case "DELETE":
		d.handleDelete(w, r, path)
	}
}

func (d *DriveLibrary) handleGet(w http.ResponseWriter, r *http.Request, path string) {
	switch path {
	case "":
		d.handleDrives(w, r)
	case "detail":
		d.handleDrivesDetail(w, r, 200, nil)
	default:
		d.handleDrive(w, r, 200, path)
	}
}

func (d *DriveLibrary) handlePost(w http.ResponseWriter, r *http.Request, path string) {
	uuid := strings.TrimSuffix(path, "/action")
	d.handleAction(w, r, uuid)
}

func (d *DriveLibrary) handleDelete(w http.ResponseWriter, r *http.Request, uuid string) {
	if ok := d.Remove(uuid); !ok {
		h := w.Header()
		h.Set("Content-Type", "application/json; charset=utf-8")
		w.WriteHeader(404)
		w.Write([]byte(jsonNotFound))
		return
	}
	w.WriteHeader(204)
}

func (d *DriveLibrary) handleDrives(w http.ResponseWriter, r *http.Request) {
	d.s.Lock()
	defer d.s.Unlock()

	var dd data.Drives
	dd.Meta.TotalCount = len(d.m)
	dd.Objects = make([]data.Drive, 0, len(d.m))
	for _, drv := range d.m {
		var drv0 data.Drive
		drv0.Resource = drv.Resource
		drv0.Owner = drv.Owner
		drv0.Status = drv.Status
		dd.Objects = append(dd.Objects, drv0)
	}

	data, err := json.Marshal(&dd)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("500 " + err.Error()))
		return
	}

	h := w.Header()
	h.Set("Content-Type", "application/json; charset=utf-8")
	w.Write(data)
}

func (d *DriveLibrary) handleDrivesDetail(w http.ResponseWriter, r *http.Request, okcode int, filter []string) {
	d.s.Lock()
	defer d.s.Unlock()

	var dd data.Drives

	if len(filter) == 0 {
		dd.Meta.TotalCount = len(d.m)
		dd.Objects = make([]data.Drive, 0, len(d.m))
		for _, drv := range d.m {
			dd.Objects = append(dd.Objects, *drv)
		}
	} else {
		dd.Meta.TotalCount = len(filter)
		dd.Objects = make([]data.Drive, 0, len(filter))
		for _, uuid := range filter {
			if drv, ok := d.m[uuid]; ok {
				dd.Objects = append(dd.Objects, *drv)
			}
		}
	}

	data, err := json.Marshal(&dd)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("500 " + err.Error()))
		return
	}

	h := w.Header()
	h.Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(okcode)
	w.Write(data)
}

func (d *DriveLibrary) handleDrive(w http.ResponseWriter, r *http.Request, okcode int, uuid string) {
	d.s.Lock()
	defer d.s.Unlock()

	h := w.Header()

	drv, ok := d.m[uuid]
	if !ok {
		h.Set("Content-Type", "application/json; charset=utf-8")
		w.WriteHeader(404)
		w.Write([]byte(jsonNotFound))
		return
	}

	data, err := json.Marshal(&drv)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("500 " + err.Error()))
		return
	}

	h.Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(okcode)
	w.Write(data)
}

func (d *DriveLibrary) handleAction(w http.ResponseWriter, r *http.Request, uuid string) {
	vv := r.URL.Query()

	v, ok := vv["do"]
	if !ok || len(v) < 1 {
		w.WriteHeader(400)
		return
	}

	action := v[0]
	switch action {
	case "clone":
		d.handleClone(w, r, uuid)
	default:
		w.WriteHeader(400)
	}
}

func (d *DriveLibrary) handleClone(w http.ResponseWriter, r *http.Request, uuid string) {
	newUUID, err := d.Clone(uuid)
	if err == ErrNotFound {
		h := w.Header()
		h.Set("Content-Type", "application/json; charset=utf-8")
		w.WriteHeader(404)
		w.Write([]byte(jsonNotFound))
		return
	} else if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("500 " + err.Error()))
		return
	}
	Drives.handleDrivesDetail(w, r, 202, []string{newUUID})
}
<commit_msg>added drive resize to mock<commit_after>// Copyright 2014 ALTOROS
// Licensed under the AGPLv3, see LICENSE file for details.

package mock

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"strings"
	"sync"

	"github.com/Altoros/gosigma/data"
)

type DriveLibrary struct {
	s sync.Mutex
	m map[string]*data.Drive
	p string
}

var Drives = &DriveLibrary{
	m: make(map[string]*data.Drive),
	p: "/api/2.0/drives",
}
var LibDrives = &DriveLibrary{
	m: make(map[string]*data.Drive),
	p: "/api/2.0/libdrives",
}

func ResetDrives() {
	Drives.Reset()
	LibDrives.Reset()
}

func InitDrive(d *data.Drive) (*data.Drive, error) {
	if d.UUID == "" {
		uuid, err := GenerateUUID()
		if err != nil {
			return nil, err
		}
		d.UUID = uuid
	}
	if d.Status == "" {
		d.Status = "unmounted"
	}

	return d, nil
}

func (d *DriveLibrary) Add(drv *data.Drive) error {
	drv, err := InitDrive(drv)
	if err != nil {
		return err
	}

	d.s.Lock()
	defer d.s.Unlock()

	d.m[drv.UUID] = drv

	return nil
}

func (d *DriveLibrary) AddDrives(dd []data.Drive) []string {
	d.s.Lock()
	defer d.s.Unlock()

	var result []string
	for _, drv := range dd {
		drv, err := InitDrive(&drv)
		if err != nil {
			d.m[drv.UUID] = drv
			result = append(result, drv.UUID)
		}
	}
	return result
}

func (d *DriveLibrary) Remove(uuid string) bool {
	d.s.Lock()
	defer d.s.Unlock()

	_, ok := d.m[uuid]
	delete(d.m, uuid)

	return ok
}

func (d *DriveLibrary) Reset() {
	d.s.Lock()
	defer d.s.Unlock()
	d.m = make(map[string]*data.Drive)
}

func (d *DriveLibrary) SetStatus(uuid, status string) {
	d.s.Lock()
	defer d.s.Unlock()

	drv, ok := d.m[uuid]
	if ok {
		drv.Status = status
	}
}

var ErrNotFound = errors.New("not found")

func (d *DriveLibrary) Clone(uuid string) (string, error) {
	d.s.Lock()
	defer d.s.Unlock()

	drv, ok := d.m[uuid]
	if !ok {
		return "", ErrNotFound
	}

	newUUID, err := GenerateUUID()
	if err != nil {
		return "", err
	}

	var newDrive data.Drive = *drv
	newDrive.Resource = *data.MakeDriveResource(newUUID)
	newDrive.Status = "unmounted"

	if d == LibDrives {
		Drives.Add(&newDrive)
	} else {
		d.m[newUUID] = &newDrive
	}

	return newUUID, nil
}

func (d *DriveLibrary) Resize(uuid string, size uint64) error {
	d.s.Lock()
	defer d.s.Unlock()

	drv, ok := d.m[uuid]
	if !ok {
		return ErrNotFound
	}

	drv.Size = size

	return nil
}

func (d *DriveLibrary) handleRequest(w http.ResponseWriter, r *http.Request) {
	path := strings.TrimSuffix(r.URL.Path, "/")
	path = strings.TrimPrefix(path, d.p)
	path = strings.TrimPrefix(path, "/")

	switch r.Method {
	case "GET":
		d.handleGet(w, r, path)
	case "POST":
		d.handlePost(w, r, path)
	case "DELETE":
		d.handleDelete(w, r, path)
	}
}

func (d *DriveLibrary) handleGet(w http.ResponseWriter, r *http.Request, path string) {
	switch path {
	case "":
		d.handleDrives(w, r)
	case "detail":
		d.handleDrivesDetail(w, r, 200, nil)
	default:
		d.handleDrive(w, r, 200, path)
	}
}

func (d *DriveLibrary) handlePost(w http.ResponseWriter, r *http.Request, path string) {
	uuid := strings.TrimSuffix(path, "/action")
	d.handleAction(w, r, uuid)
}

func (d *DriveLibrary) handleDelete(w http.ResponseWriter, r *http.Request, uuid string) {
	if ok := d.Remove(uuid); !ok {
		h := w.Header()
		h.Set("Content-Type", "application/json; charset=utf-8")
		w.WriteHeader(404)
		w.Write([]byte(jsonNotFound))
		return
	}
	w.WriteHeader(204)
}

func (d *DriveLibrary) handleDrives(w http.ResponseWriter, r *http.Request) {
	d.s.Lock()
	defer d.s.Unlock()

	var dd data.Drives
	dd.Meta.TotalCount = len(d.m)
	dd.Objects = make([]data.Drive, 0, len(d.m))
	for _, drv := range d.m {
		var drv0 data.Drive
		drv0.Resource = drv.Resource
		drv0.Owner = drv.Owner
		drv0.Status = drv.Status
		dd.Objects = append(dd.Objects, drv0)
	}

	data, err := json.Marshal(&dd)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("500 " + err.Error()))
		return
	}

	h := w.Header()
	h.Set("Content-Type", "application/json; charset=utf-8")
	w.Write(data)
}

func (d *DriveLibrary) handleDrivesDetail(w http.ResponseWriter, r *http.Request, okcode int, filter []string) {
	d.s.Lock()
	defer d.s.Unlock()

	var dd data.Drives

	if len(filter) == 0 {
		dd.Meta.TotalCount = len(d.m)
		dd.Objects = make([]data.Drive, 0, len(d.m))
		for _, drv := range d.m {
			dd.Objects = append(dd.Objects, *drv)
		}
	} else {
		dd.Meta.TotalCount = len(filter)
		dd.Objects = make([]data.Drive, 0, len(filter))
		for _, uuid := range filter {
			if drv, ok := d.m[uuid]; ok {
				dd.Objects = append(dd.Objects, *drv)
			}
		}
	}

	data, err := json.Marshal(&dd)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("500 " + err.Error()))
		return
	}

	h := w.Header()
	h.Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(okcode)
	w.Write(data)
}

func (d *DriveLibrary) handleDrive(w http.ResponseWriter, r *http.Request, okcode int, uuid string) {
	d.s.Lock()
	defer d.s.Unlock()

	h := w.Header()

	drv, ok := d.m[uuid]
	if !ok {
		h.Set("Content-Type", "application/json; charset=utf-8")
		w.WriteHeader(404)
		w.Write([]byte(jsonNotFound))
		return
	}

	data, err := json.Marshal(&drv)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("500 " + err.Error()))
		return
	}

	h.Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(okcode)
	w.Write(data)
}

func (d *DriveLibrary) handleAction(w http.ResponseWriter, r *http.Request, uuid string) {
	vv := r.URL.Query()

	v, ok := vv["do"]
	if !ok || len(v) < 1 {
		w.WriteHeader(400)
		return
	}

	action := v[0]
	switch action {
	case "clone":
		d.handleClone(w, r, uuid)
	case "resize":
		d.handleResize(w, r, uuid)
	default:
		w.WriteHeader(400)
	}
}

func (d *DriveLibrary) handleClone(w http.ResponseWriter, r *http.Request, uuid string) {
	newUUID, err := d.Clone(uuid)
	if err == ErrNotFound {
		h := w.Header()
		h.Set("Content-Type", "application/json; charset=utf-8")
		w.WriteHeader(404)
		w.Write([]byte(jsonNotFound))
		return
	} else if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("500 " + err.Error()))
		return
	}
	Drives.handleDrivesDetail(w, r, 202, []string{newUUID})
}

func (d *DriveLibrary) handleResize(w http.ResponseWriter, r *http.Request, uuid string) {
	bb, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("500 " + err.Error()))
		return
	}

	drv, err := data.ReadDrive(bytes.NewReader(bb))
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("500 " + err.Error()))
		return
	}

	err = d.Resize(uuid, drv.Size)
	if err == ErrNotFound {
		h := w.Header()
		h.Set("Content-Type", "application/json; charset=utf-8")
		w.WriteHeader(404)
		w.Write([]byte(jsonNotFound))
		return
	} else if err != nil {
		w.WriteHeader(500)
		w.Write([]byte("500 " + err.Error()))
		return
	}

	d.handleDrivesDetail(w, r, 202, []string{uuid})
}
<|endoftext|>


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

package dicom

import (
	"bytes"
	"io"
	"os"
	"path"
	"reflect"
	"testing"

	"encoding/binary"

)

var (
	nestedDataSetElement = createDataElement(0x00081155, UIVR,
		[]string{"1.2.840.10008.5.1.4.1.1.4"}, 26)
	nestedSeq         = createSingletonSequence(nestedDataSetElement)
	seq               = createSingletonSequence(createDataElement(0x00081150, SQVR, &nestedSeq, 42))
	bufferedPixelData = createDataElement(PixelDataTag, OWVR, [][]byte{{0x11, 0x11, 0x22, 0x22}}, 4)
	expectedElements  = []*DataElement{
		createDataElement(0x00020000, ULVR, []uint32{198}, 4),
		createDataElement(0x00020001, OBVR, [][]byte{{0, 1}}, 2),
		createDataElement(0x00020002, UIVR, []string{"1.2.840.10008.5.1.4.1.1.4"}, 26),
		createDataElement(0x00020003, UIVR,
			[]string{"1.2.840.113619.2.176.3596.3364818.7271.1259708501.876"}, 54),
		createDataElement(0x00020010, UIVR, []string{"1.2.840.10008.1.2.1"}, 20),
		createDataElement(0x00020012, UIVR, []string{"1.2.276.0.7230010.3.0.3.5.4"}, 28),
		createDataElement(0x00020013, SHVR, []string{"OFFIS_DCMTK_354"}, 16),
		createDataElement(0x00081110, SQVR, &seq, 62),
		bufferedPixelData,
	}
)

func TestIterator_NextElement(t *testing.T) {
	tests := []struct {
		name   string
		file   string
		syntax transferSyntax
		want   *DataSet
	}{
		{
			"Explicit Lengths, Explicit VR, Little Endian",
			"ExplicitVRLittleEndian.dcm",
			explicitVRLittleEndian,
			createExpectedDataSet(bufferedPixelData, ExplicitVRLittleEndianUID),
		},
		{
			"Undefined Sequence & Item lengths, Explicit VR, Little Endian",
			"ExplicitVRLittleEndianUndefLen.dcm",
			explicitVRLittleEndian,
			createExpectedDataSet(bufferedPixelData, ExplicitVRLittleEndianUID),
		},
		{
			"Explicit Length, Explicit VR, Big Endian",
			"ExplicitVRBigEndian.dcm",
			explicitVRBigEndian,
			createExpectedDataSet(bufferedPixelData, ExplicitVRBigEndianUID),
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			iter, err := createIteratorFromFile(tc.file)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			for elem, err := iter.NextElement(); err != io.EOF; elem, err = iter.NextElement() {
				if err != nil {
					t.Fatalf("NextElement() => %v", err)
				}
				compareDataElements(elem, tc.want.Elements[uint32(elem.Tag)], tc.syntax.ByteOrder, t)
			}
		})
	}
}

func TestIterator_Close(t *testing.T) {
	iter, err := createIteratorFromFile("ExplicitVRLittleEndian.dcm")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if err := iter.Close(); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if _, err := iter.NextElement(); err != io.EOF {
		t.Fatalf("got %v, want %v", err, io.EOF)
	}
}

func TestIterator_atEndOfInput(t *testing.T) {
	iter, err := newDataElementIterator(dcmReaderFromBytes(nil), defaultEncoding)
	if err != nil {
		t.Fatalf("unexepcted error: %v", err)
	}
	_, err = iter.NextElement()
	if err != io.EOF {
		t.Fatalf("expected iterator to return EOF if at end of input: got %v, want %v", err, io.EOF)
	}
}

func TestEmptyIterator(t *testing.T) {
	iter := &dataElementIterator{empty: true, metaHeader: emptyElementIterator{}}
	_, err := iter.NextElement()
	if err != io.EOF {
		t.Fatalf("expected empty iterator to return io.EOF: got %v, want %v", err, io.EOF)
	}
}

func compareDataElements(e1 *DataElement, e2 *DataElement, order binary.ByteOrder, t *testing.T) {
	if e1 == nil || e2 == nil {
		if e1 != e2 {
			t.Fatalf("expected both elements to be nil: got %v, want %v", e1, e2)
		}
		return
	}
	if e1.VR != e2.VR {
		t.Fatalf("expected VRs to be equal: got %v, want %v", e1.VR, e2.VR)
	}
	if e1.Tag != e2.Tag {
		t.Fatalf("expected tags to be equal: got %v, want %v", e1.Tag, e2.Tag)
	}

	var err error
	e1, err = processElement(e1, order)
	if err != nil {
		t.Fatalf("unexpected error unstreaming data element: %v", err)
	}
	e2, err = processElement(e2, order)
	if err != nil {
		t.Fatalf("unexpected error unstreaming data element: %v", err)
	}

	if e1.VR != SQVR {
		if !reflect.DeepEqual(e1.ValueField, e2.ValueField) {
			t.Fatalf("expected ValueFields to be equal: got %v, want %v",
				e1.ValueField, e2.ValueField)
		}
	} else {
		compareSequences(e1.ValueField.(*Sequence), e2.ValueField.(*Sequence), order, t)
	}
}

func compareSequences(s1 *Sequence, s2 *Sequence, order binary.ByteOrder, t *testing.T) {
	if len(s1.Items) != len(s2.Items) {
		t.Fatalf("expected sequences to have same length: got %v, want %v",
			len(s1.Items), len(s2.Items))
	}

	for i := range s1.Items {
		compareDataSets(s1.Items[i], s2.Items[i], order, t)
	}
}

func compareDataSets(d1 *DataSet, d2 *DataSet, order binary.ByteOrder, t *testing.T) {
	k1, k2 := getKeys(d1.Elements), getKeys(d2.Elements)

	if !reflect.DeepEqual(k1, k2) {
		t.Fatalf("expected datasets to have same keys: got %v, want %v", k1, k2)
	}

	for k := range k1 {
		compareDataElements(d1.Elements[k], d2.Elements[k], order, t)
	}
}

func getKeys(m map[uint32]*DataElement) map[uint32]bool {
	ret := make(map[uint32]bool)
	for k := range m {
		ret[k] = true
	}
	return ret
}

func createIteratorFromFile(file string) (DataElementIterator, error) {
	r, err := openFile(file)
	if err != nil {
		return nil, err
	}

	return NewDataElementIterator(r)
}

func openFile(name string) (io.Reader, error) {
	var p = path.Join("../",
		"testdata/+name)

	return os.Open(p)
}

var sampleBytes = []byte{1, 2, 3, 4}

func countReaderFromBytes(data []byte) *countReader {
	return &countReader{
		bytes.NewBuffer(data),
		0,
	}
}

func dcmReaderFromBytes(data []byte) *dcmReader {
	return newDcmReader(bytes.NewBuffer(data))
}

func createDataElement(tag uint32, vr *VR, value interface{}, length uint32) *DataElement {
	return &DataElement{DataElementTag(tag), vr, value, length}
}

func createSingletonSequence(elements ...*DataElement) Sequence {
	ds := DataSet{map[uint32]*DataElement{}}
	for _, elem := range elements {
		ds.Elements[uint32(elem.Tag)] = elem
	}
	return Sequence{[]*DataSet{&ds}}
}
<commit_msg>Update copybara to be iterative after first commit. Also fix destination to .git<commit_after>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package dicom

import (
	"bytes"
	"io"
	"os"
	"path"
	"reflect"
	"testing"

	"encoding/binary"

)

var (
	nestedDataSetElement = createDataElement(0x00081155, UIVR,
		[]string{"1.2.840.10008.5.1.4.1.1.4"}, 26)
	nestedSeq         = createSingletonSequence(nestedDataSetElement)
	seq               = createSingletonSequence(createDataElement(0x00081150, SQVR, &nestedSeq, 42))
	bufferedPixelData = createDataElement(PixelDataTag, OWVR, [][]byte{{0x11, 0x11, 0x22, 0x22}}, 4)
	expectedElements  = []*DataElement{
		createDataElement(0x00020000, ULVR, []uint32{198}, 4),
		createDataElement(0x00020001, OBVR, [][]byte{{0, 1}}, 2),
		createDataElement(0x00020002, UIVR, []string{"1.2.840.10008.5.1.4.1.1.4"}, 26),
		createDataElement(0x00020003, UIVR,
			[]string{"1.2.840.113619.2.176.3596.3364818.7271.1259708501.876"}, 54),
		createDataElement(0x00020010, UIVR, []string{"1.2.840.10008.1.2.1"}, 20),
		createDataElement(0x00020012, UIVR, []string{"1.2.276.0.7230010.3.0.3.5.4"}, 28),
		createDataElement(0x00020013, SHVR, []string{"OFFIS_DCMTK_354"}, 16),
		createDataElement(0x00081110, SQVR, &seq, 62),
		bufferedPixelData,
	}
)

func TestIterator_NextElement(t *testing.T) {
	tests := []struct {
		name   string
		file   string
		syntax transferSyntax
		want   *DataSet
	}{
		{
			"Explicit Lengths, Explicit VR, Little Endian",
			"ExplicitVRLittleEndian.dcm",
			explicitVRLittleEndian,
			createExpectedDataSet(bufferedPixelData, ExplicitVRLittleEndianUID),
		},
		{
			"Undefined Sequence & Item lengths, Explicit VR, Little Endian",
			"ExplicitVRLittleEndianUndefLen.dcm",
			explicitVRLittleEndian,
			createExpectedDataSet(bufferedPixelData, ExplicitVRLittleEndianUID),
		},
		{
			"Explicit Length, Explicit VR, Big Endian",
			"ExplicitVRBigEndian.dcm",
			explicitVRBigEndian,
			createExpectedDataSet(bufferedPixelData, ExplicitVRBigEndianUID),
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			iter, err := createIteratorFromFile(tc.file)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			for elem, err := iter.NextElement(); err != io.EOF; elem, err = iter.NextElement() {
				if err != nil {
					t.Fatalf("NextElement() => %v", err)
				}
				compareDataElements(elem, tc.want.Elements[uint32(elem.Tag)], tc.syntax.ByteOrder, t)
			}
		})
	}
}

func TestIterator_Close(t *testing.T) {
	iter, err := createIteratorFromFile("ExplicitVRLittleEndian.dcm")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if err := iter.Close(); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if _, err := iter.NextElement(); err != io.EOF {
		t.Fatalf("got %v, want %v", err, io.EOF)
	}
}

func TestIterator_atEndOfInput(t *testing.T) {
	iter, err := newDataElementIterator(dcmReaderFromBytes(nil), defaultEncoding)
	if err != nil {
		t.Fatalf("unexepcted error: %v", err)
	}
	_, err = iter.NextElement()
	if err != io.EOF {
		t.Fatalf("expected iterator to return EOF if at end of input: got %v, want %v", err, io.EOF)
	}
}

func TestEmptyIterator(t *testing.T) {
	iter := &dataElementIterator{empty: true, metaHeader: emptyElementIterator{}}
	_, err := iter.NextElement()
	if err != io.EOF {
		t.Fatalf("expected empty iterator to return io.EOF: got %v, want %v", err, io.EOF)
	}
}

func compareDataElements(e1 *DataElement, e2 *DataElement, order binary.ByteOrder, t *testing.T) {
	if e1 == nil || e2 == nil {
		if e1 != e2 {
			t.Fatalf("expected both elements to be nil: got %v, want %v", e1, e2)
		}
		return
	}
	if e1.VR != e2.VR {
		t.Fatalf("expected VRs to be equal: got %v, want %v", e1.VR, e2.VR)
	}
	if e1.Tag != e2.Tag {
		t.Fatalf("expected tags to be equal: got %v, want %v", e1.Tag, e2.Tag)
	}

	var err error
	e1, err = processElement(e1, order)
	if err != nil {
		t.Fatalf("unexpected error unstreaming data element: %v", err)
	}
	e2, err = processElement(e2, order)
	if err != nil {
		t.Fatalf("unexpected error unstreaming data element: %v", err)
	}

	if e1.VR != SQVR {
		if !reflect.DeepEqual(e1.ValueField, e2.ValueField) {
			t.Fatalf("expected ValueFields to be equal: got %v, want %v",
				e1.ValueField, e2.ValueField)
		}
	} else {
		compareSequences(e1.ValueField.(*Sequence), e2.ValueField.(*Sequence), order, t)
	}
}

func compareSequences(s1 *Sequence, s2 *Sequence, order binary.ByteOrder, t *testing.T) {
	if len(s1.Items) != len(s2.Items) {
		t.Fatalf("expected sequences to have same length: got %v, want %v",
			len(s1.Items), len(s2.Items))
	}

	for i := range s1.Items {
		compareDataSets(s1.Items[i], s2.Items[i], order, t)
	}
}

func compareDataSets(d1 *DataSet, d2 *DataSet, order binary.ByteOrder, t *testing.T) {
	k1, k2 := getKeys(d1.Elements), getKeys(d2.Elements)

	if !reflect.DeepEqual(k1, k2) {
		t.Fatalf("expected datasets to have same keys: got %v, want %v", k1, k2)
	}

	for k := range k1 {
		compareDataElements(d1.Elements[k], d2.Elements[k], order, t)
	}
}

func getKeys(m map[uint32]*DataElement) map[uint32]bool {
	ret := make(map[uint32]bool)
	for k := range m {
		ret[k] = true
	}
	return ret
}

func createIteratorFromFile(file string) (DataElementIterator, error) {
	r, err := openFile(file)
	if err != nil {
		return nil, err
	}

	return NewDataElementIterator(r)
}

func openFile(name string) (io.Reader, error) {
	var p = path.Join("../",
		"testdata/"+name)

	return os.Open(p)
}

var sampleBytes = []byte{1, 2, 3, 4}

func countReaderFromBytes(data []byte) *countReader {
	return &countReader{
		bytes.NewBuffer(data),
		0,
	}
}

func dcmReaderFromBytes(data []byte) *dcmReader {
	return newDcmReader(bytes.NewBuffer(data))
}

func createDataElement(tag uint32, vr *VR, value interface{}, length uint32) *DataElement {
	return &DataElement{DataElementTag(tag), vr, value, length}
}

func createSingletonSequence(elements ...*DataElement) Sequence {
	ds := DataSet{map[uint32]*DataElement{}}
	for _, elem := range elements {
		ds.Elements[uint32(elem.Tag)] = elem
	}
	return Sequence{[]*DataSet{&ds}}
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package walking

import (
	"context"
	"encoding/base64"
	"fmt"
	"io"
	"math/rand"
	"time"

	"github.com/containerd/containerd/archive"
	"github.com/containerd/containerd/archive/compression"
	"github.com/containerd/containerd/content"
	"github.com/containerd/containerd/diff"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/mount"
	digest "github.com/opencontainers/go-digest"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
)

type walkingDiff struct {
	store content.Store
}

var emptyDesc = ocispec.Descriptor{}
var uncompressed = "containerd.io/uncompressed"

// NewWalkingDiff is a generic implementation of diff.Comparer.  The diff is
// calculated by mounting both the upper and lower mount sets and walking the
// mounted directories concurrently. Changes are calculated by comparing files
// against each other or by comparing file existence between directories.
// NewWalkingDiff uses no special characteristics of the mount sets and is
// expected to work with any filesystem.
func NewWalkingDiff(store content.Store) diff.Comparer {
	return &walkingDiff{
		store: store,
	}
}

// Compare creates a diff between the given mounts and uploads the result
// to the content store.
func (s *walkingDiff) Compare(ctx context.Context, lower, upper []mount.Mount, opts ...diff.Opt) (d ocispec.Descriptor, err error) {
	var config diff.Config
	for _, opt := range opts {
		if err := opt(&config); err != nil {
			return emptyDesc, err
		}
	}

	if config.MediaType == "" {
		config.MediaType = ocispec.MediaTypeImageLayerGzip
	}

	var isCompressed bool
	switch config.MediaType {
	case ocispec.MediaTypeImageLayer:
	case ocispec.MediaTypeImageLayerGzip:
		isCompressed = true
	default:
		return emptyDesc, errors.Wrapf(errdefs.ErrNotImplemented, "unsupported diff media type: %v", config.MediaType)
	}

	var ocidesc ocispec.Descriptor
	if err := mount.WithTempMount(ctx, lower, func(lowerRoot string) error {
		return mount.WithTempMount(ctx, upper, func(upperRoot string) error {
			var newReference bool
			if config.Reference == "" {
				newReference = true
				config.Reference = uniqueRef()
			}

			cw, err := s.store.Writer(ctx,
				content.WithRef(config.Reference),
				content.WithDescriptor(ocispec.Descriptor{
					MediaType: config.MediaType, // most contentstore implementations just ignore this
				}))
			if err != nil {
				return errors.Wrap(err, "failed to open writer")
			}
			defer func() {
				if err != nil {
					cw.Close()
					if newReference {
						if abortErr := s.store.Abort(ctx, config.Reference); abortErr != nil {
							log.G(ctx).WithError(abortErr).WithField("ref", config.Reference).Warnf("failed to delete diff upload")
						}
					}
				}
			}()
			if !newReference {
				if err = cw.Truncate(0); err != nil {
					return err
				}
			}

			if isCompressed {
				dgstr := digest.SHA256.Digester()
				var compressed io.WriteCloser
				compressed, err = compression.CompressStream(cw, compression.Gzip)
				if err != nil {
					return errors.Wrap(err, "failed to get compressed stream")
				}
				err = archive.WriteDiff(ctx, io.MultiWriter(compressed, dgstr.Hash()), lowerRoot, upperRoot)
				compressed.Close()
				if err != nil {
					return errors.Wrap(err, "failed to write compressed diff")
				}

				if config.Labels == nil {
					config.Labels = map[string]string{}
				}
				config.Labels[uncompressed] = dgstr.Digest().String()
			} else {
				if err = archive.WriteDiff(ctx, cw, lowerRoot, upperRoot); err != nil {
					return errors.Wrap(err, "failed to write diff")
				}
			}

			var commitopts []content.Opt
			if config.Labels != nil {
				commitopts = append(commitopts, content.WithLabels(config.Labels))
			}

			dgst := cw.Digest()
			if err := cw.Commit(ctx, 0, dgst, commitopts...); err != nil {
				if !errdefs.IsAlreadyExists(err) {
					return errors.Wrap(err, "failed to commit")
				}
			}

			info, err := s.store.Info(ctx, dgst)
			if err != nil {
				return errors.Wrap(err, "failed to get info from content store")
			}
			if info.Labels == nil {
				info.Labels = make(map[string]string)
			}
			// Set uncompressed label if digest already existed without label
			if _, ok := info.Labels[uncompressed]; !ok {
				info.Labels[uncompressed] = config.Labels[uncompressed]
				if _, err := s.store.Update(ctx, info, "labels."+uncompressed); err != nil {
					return errors.Wrap(err, "error setting uncompressed label")
				}
			}

			ocidesc = ocispec.Descriptor{
				MediaType: config.MediaType,
				Size:      info.Size,
				Digest:    info.Digest,
			}
			return nil
		})
	}); err != nil {
		return emptyDesc, err
	}

	return ocidesc, nil
}

func uniqueRef() string {
	t := time.Now()
	var b [3]byte
	// Ignore read failures, just decreases uniqueness
	rand.Read(b[:])
	return fmt.Sprintf("%d-%s", t.UnixNano(), base64.URLEncoding.EncodeToString(b[:]))
}
<commit_msg>diff/walking: fix defer cleanup<commit_after>/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package walking

import (
	"context"
	"encoding/base64"
	"fmt"
	"io"
	"math/rand"
	"time"

	"github.com/containerd/containerd/archive"
	"github.com/containerd/containerd/archive/compression"
	"github.com/containerd/containerd/content"
	"github.com/containerd/containerd/diff"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/mount"
	digest "github.com/opencontainers/go-digest"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
)

type walkingDiff struct {
	store content.Store
}

var emptyDesc = ocispec.Descriptor{}
var uncompressed = "containerd.io/uncompressed"

// NewWalkingDiff is a generic implementation of diff.Comparer.  The diff is
// calculated by mounting both the upper and lower mount sets and walking the
// mounted directories concurrently. Changes are calculated by comparing files
// against each other or by comparing file existence between directories.
// NewWalkingDiff uses no special characteristics of the mount sets and is
// expected to work with any filesystem.
func NewWalkingDiff(store content.Store) diff.Comparer {
	return &walkingDiff{
		store: store,
	}
}

// Compare creates a diff between the given mounts and uploads the result
// to the content store.
func (s *walkingDiff) Compare(ctx context.Context, lower, upper []mount.Mount, opts ...diff.Opt) (d ocispec.Descriptor, err error) {
	var config diff.Config
	for _, opt := range opts {
		if err := opt(&config); err != nil {
			return emptyDesc, err
		}
	}

	if config.MediaType == "" {
		config.MediaType = ocispec.MediaTypeImageLayerGzip
	}

	var isCompressed bool
	switch config.MediaType {
	case ocispec.MediaTypeImageLayer:
	case ocispec.MediaTypeImageLayerGzip:
		isCompressed = true
	default:
		return emptyDesc, errors.Wrapf(errdefs.ErrNotImplemented, "unsupported diff media type: %v", config.MediaType)
	}

	var ocidesc ocispec.Descriptor
	if err := mount.WithTempMount(ctx, lower, func(lowerRoot string) error {
		return mount.WithTempMount(ctx, upper, func(upperRoot string) (retErr error) {
			var newReference bool
			if config.Reference == "" {
				newReference = true
				config.Reference = uniqueRef()
			}

			cw, err := s.store.Writer(ctx,
				content.WithRef(config.Reference),
				content.WithDescriptor(ocispec.Descriptor{
					MediaType: config.MediaType, // most contentstore implementations just ignore this
				}))
			if err != nil {
				return errors.Wrap(err, "failed to open writer")
			}
			defer func() {
				if retErr != nil {
					cw.Close()
					if newReference {
						if abortErr := s.store.Abort(ctx, config.Reference); abortErr != nil {
							log.G(ctx).WithError(abortErr).WithField("ref", config.Reference).Warnf("failed to delete diff upload")
						}
					}
				}
			}()
			if !newReference {
				if err = cw.Truncate(0); err != nil {
					return err
				}
			}

			if isCompressed {
				dgstr := digest.SHA256.Digester()
				var compressed io.WriteCloser
				compressed, err = compression.CompressStream(cw, compression.Gzip)
				if err != nil {
					return errors.Wrap(err, "failed to get compressed stream")
				}
				err = archive.WriteDiff(ctx, io.MultiWriter(compressed, dgstr.Hash()), lowerRoot, upperRoot)
				compressed.Close()
				if err != nil {
					return errors.Wrap(err, "failed to write compressed diff")
				}

				if config.Labels == nil {
					config.Labels = map[string]string{}
				}
				config.Labels[uncompressed] = dgstr.Digest().String()
			} else {
				if err = archive.WriteDiff(ctx, cw, lowerRoot, upperRoot); err != nil {
					return errors.Wrap(err, "failed to write diff")
				}
			}

			var commitopts []content.Opt
			if config.Labels != nil {
				commitopts = append(commitopts, content.WithLabels(config.Labels))
			}

			dgst := cw.Digest()
			if err := cw.Commit(ctx, 0, dgst, commitopts...); err != nil {
				if !errdefs.IsAlreadyExists(err) {
					return errors.Wrap(err, "failed to commit")
				}
			}

			info, err := s.store.Info(ctx, dgst)
			if err != nil {
				return errors.Wrap(err, "failed to get info from content store")
			}
			if info.Labels == nil {
				info.Labels = make(map[string]string)
			}
			// Set uncompressed label if digest already existed without label
			if _, ok := info.Labels[uncompressed]; !ok {
				info.Labels[uncompressed] = config.Labels[uncompressed]
				if _, err := s.store.Update(ctx, info, "labels."+uncompressed); err != nil {
					return errors.Wrap(err, "error setting uncompressed label")
				}
			}

			ocidesc = ocispec.Descriptor{
				MediaType: config.MediaType,
				Size:      info.Size,
				Digest:    info.Digest,
			}
			return nil
		})
	}); err != nil {
		return emptyDesc, err
	}

	return ocidesc, nil
}

func uniqueRef() string {
	t := time.Now()
	var b [3]byte
	// Ignore read failures, just decreases uniqueness
	rand.Read(b[:])
	return fmt.Sprintf("%d-%s", t.UnixNano(), base64.URLEncoding.EncodeToString(b[:]))
}
<|endoftext|>


<commit_before>package triggers

type Donators struct{}

func (d *Donators) Triggers() []string {
	return []string{
		"donators",
		"donations",
		"donate",
		"supporters",
		"support",
		"patreon",
		"patreons",
		"credits",
	}
}

func (d *Donators) Response(trigger string, content string) string {
	return "<:robyulblush:327206930437373952> **These awesome people support me:**\nKakkela 💕\nSunny 💓\nsomicidal minaiac 💞\nOokami 🖤\nKeldra 💗\nTN 💝\nseulguille 💘\nSlenn 💜\nFugu ❣\nWoori 💞\nhikari 💙\nAshton 💖\nKay 💝\njamie 💓\nHomeboywill 💘\nRimbol 💕\nGenisphere 💖\nekgus 💗\nCPark 💞\njungoo 💕\nShawn 💗\nSaltiestPeach 💘\nBae Nja Min 💖\nhaerts 💓\ncasker 💞\nJean 💖\npooth 💘\nThank you so much!\n_You want to be in this list? <https://www.patreon.com/sekl>!_"
}
<commit_msg>[donators] adds shb117!<commit_after>package triggers

type Donators struct{}

func (d *Donators) Triggers() []string {
	return []string{
		"donators",
		"donations",
		"donate",
		"supporters",
		"support",
		"patreon",
		"patreons",
		"credits",
	}
}

func (d *Donators) Response(trigger string, content string) string {
	return "<:robyulblush:327206930437373952> **These awesome people support me:**\nKakkela 💕\nSunny 💓\nsomicidal minaiac 💞\nOokami 🖤\nKeldra 💗\nTN 💝\nseulguille 💘\nSlenn 💜\nFugu ❣\nWoori 💞\nhikari 💙\nAshton 💖\nKay 💝\njamie 💓\nHomeboywill 💘\nRimbol 💕\nGenisphere 💖\nekgus 💗\nCPark 💞\njungoo 💕\nShawn 💗\nSaltiestPeach 💘\nBae Nja Min 💖\nhaerts 💓\ncasker 💞\nJean 💖\npooth 💘\nshb117 💕\nThank you so much!\n_You want to be in this list? <https://www.patreon.com/sekl>!_"
}
<|endoftext|>